import json
import torch
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from transformers import AutoTokenizer, AutoModelForSequenceClassification

MODEL_PATH = "./final_model"   # ← points to folder inside Docker
TOP_K = 3
MAX_LENGTH = 512

print("Loading model...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_PATH)
model.eval()

with open(f"{MODEL_PATH}/label_mapping.json", "r", encoding="utf-8") as f:
    mapping = json.load(f)
id2label = mapping["id2label"]
print("Model loaded. API ready!")

app = FastAPI(title="Field Classification API")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

class PredictRequest(BaseModel):
    text: str

class Prediction(BaseModel):
    rank: int
    label: str
    confidence_percent: float

class PredictResponse(BaseModel):
    predictions: list[Prediction]
    confidence_sum_percent: float
    input_text_preview: str

@app.get("/")
def root():
    return {"status": "ok", "message": "Field Classification API is running!"}

@app.post("/predict", response_model=PredictResponse)
def predict(request: PredictRequest):
    text = request.text.strip()
    if not text:
        raise HTTPException(status_code=400, detail="text must not be empty.")

    inputs = tokenizer(
        text,
        return_tensors="pt",
        truncation=True,
        padding=True,
        max_length=MAX_LENGTH,
    )

    with torch.no_grad():
        outputs = model(**inputs)

    logits = outputs.logits[0]
    top_values, top_indices = torch.topk(logits, TOP_K)
    percentages = torch.softmax(top_values, dim=0) * 100

    predictions = [
        Prediction(
            rank=i + 1,
            label=id2label[str(idx.item())],
            confidence_percent=round(pct.item(), 2),
        )
        for i, (idx, pct) in enumerate(zip(top_indices, percentages))
    ]

    return PredictResponse(
        predictions=predictions,
        confidence_sum_percent=round(percentages.sum().item(), 2),
        input_text_preview=text[:120],
    )