from fastapi import FastAPI, UploadFile, File, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from contextlib import asynccontextmanager
import torch
import uvicorn

from app.predictor import ProteinPredictor
from app.schemas import PredictionResponse, HealthResponse

# ── Global predictor instance ──────────────────────────────────────────────
predictor: ProteinPredictor = None


@asynccontextmanager
async def lifespan(app: FastAPI):
    """Load model once at startup, release on shutdown."""
    global predictor
    print("🔄 Loading model...")
    predictor = ProteinPredictor()
    predictor.load()
    print("✅ Model ready!")
    yield
    print("🛑 Shutting down...")


# ── App ─────────────────────────────────────────────────────────────────────
app = FastAPI(
    title="Protein Function Prediction API",
    description="Predicts GO terms for a protein given its FASTA sequence using GraphSAGE + ESM2.",
    version="1.0.0",
    lifespan=lifespan,
)

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


# ── Routes ───────────────────────────────────────────────────────────────────
@app.get("/health", response_model=HealthResponse, tags=["Health"])
def health():
    """Check that the API and model are up."""
    return HealthResponse(
        status="ok",
        model_loaded=predictor is not None and predictor.is_ready,
        device=str(predictor.device) if predictor else "unknown",
    )


@app.post("/predict", response_model=PredictionResponse, tags=["Prediction"])
async def predict(file: UploadFile = File(..., description="FASTA file (.fasta or .fa)")):
    """
    Upload a FASTA file and get predicted GO terms.

    - Accepts single or multi-sequence FASTA files.
    - Returns GO term IDs + their confidence scores.
    - Threshold: 0.5 (configurable via config.py).
    """
    if not file.filename.endswith((".fasta", ".fa", ".txt")):
        raise HTTPException(status_code=400, detail="File must be a .fasta or .fa file.")

    content = await file.read()
    try:
        fasta_text = content.decode("utf-8")
    except UnicodeDecodeError:
        raise HTTPException(status_code=400, detail="Could not decode file. Make sure it's a valid FASTA text file.")

    if not fasta_text.strip():
        raise HTTPException(status_code=400, detail="File is empty.")

    try:
        result = predictor.predict(fasta_text)
    except ValueError as e:
        raise HTTPException(status_code=422, detail=str(e))
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Prediction failed: {str(e)}")

    return result


if __name__ == "__main__":
    uvicorn.run("app.main:app", host="0.0.0.0", port=8000, reload=False)