from contextlib import asynccontextmanager
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from main import load_models, gradient_attention_matrix, annotate, explain_token, SKIP_TOKENS

@asynccontextmanager
async def lifespan(app: FastAPI):
    print("loading models")
    app.state.tokenizer, app.state.model, app.state.nlp_en, app.state.nlp_ja = load_models()
    print("models loaded")

    yield

app = FastAPI(lifespan=lifespan)

origins = [
    "http://localhost:5173",
    "https://xai-translation-decoder.vercel.app",
]

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

class ExplainRequest(BaseModel):
    text: str

class TokenExplanation(BaseModel):
    tgt_token: str
    tgt_surface: str
    rationale: str
    top_sources: list[dict]
    alternatives: list[str]
    ja_annotation: dict | None

class ExplainResponse(BaseModel):
    translation: str
    src_tokens: list[str]
    tgt_tokens: list[str]
    matrix: list[list[float]]
    tokens: list[TokenExplanation]

@app.post("/explain", response_model=ExplainResponse)
def explain(req: ExplainRequest):
    #{"translation": translation, "src_tokens":  src_tokens, "tgt_tokens": tgt_tokens, "matrix": matrix, "logits": logits.detach()}
    grad_res = gradient_attention_matrix(req.text, app.state.tokenizer, app.state.model)

    #{"en_tokens": en_tokens, "ja_tokens": ja_tokens}
    annots = annotate(req.text, grad_res["translation"], app.state.nlp_en, app.state.nlp_ja)
    
    tokens = []
    for i in range(len(grad_res["tgt_tokens"])):
        if grad_res["tgt_tokens"][i] in SKIP_TOKENS:
            continue
        exp = explain_token(i, grad_res["tgt_tokens"], grad_res["src_tokens"], grad_res["matrix"], annots, grad_res["logits"], app.state.tokenizer)
        tokens.append(exp)

    return {
        "translation": grad_res["translation"],
        "src_tokens": grad_res["src_tokens"],
        "tgt_tokens": grad_res["tgt_tokens"],
        "matrix": grad_res["matrix"].tolist(),
        "tokens": tokens
    }