from fastapi import FastAPI
from fastapi.responses import JSONResponse
from pydantic import BaseModel  # ✅ FIX: Use BaseModel not model
from app.model_loader import load_model
import torch

# Initialize app and model
app = FastAPI()
model, tokenizer = load_model()

# ✅ Define request body schema using Pydantic
class InputText(BaseModel):
    input: str

# 🚀 Inference endpoint
@app.post("/predict")
async def predict(input_text: InputText):
    # Create prompt
    prompt = (
        "You are a neuroscience research assistant.\n"
        "Determine if the following abstract was modified by an AI model.\n"
        "Respond with 'yes' if it was modified, or 'no' if it is original.\n\n"
        f"Abstract:\n{input_text.input}\n\nAnswer:"
    )

    # Tokenize
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

    # Generate
    with torch.no_grad():
        output = model.generate(
            **inputs,
            max_new_tokens=50,
            do_sample=False,
            temperature=0.3,
            pad_token_id=tokenizer.eos_token_id,
        )

    # Decode and extract answer
    decoded_output = tokenizer.decode(output[0], skip_special_tokens=True)
    response_text = decoded_output[len(prompt):].strip().lower()

    if "yes" in response_text:
        answer = "yes"
    elif "no" in response_text:
        answer = "no"
    else:
        answer = "unknown"

    return {"output": response_text, "answer": answer}