"""
FastAPI Server for AES-Feedback.
Serves the IndoBERT + IndoSBERT + Feedback Engine pipeline.
Optimized for HuggingFace Docker Space.

Usage (localhost):
    python -m api.app

Usage (HF Space):
    uvicorn api.app:app --host 0.0.0.0 --port 7860
"""
import os
import sys
import tempfile
import torch
from contextlib import asynccontextmanager
from fastapi import FastAPI, HTTPException, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from huggingface_hub import snapshot_download

from .schemas import (
    PredictRequest, PredictResponse, HealthResponse,
    FilePredictResponse, FilePredictResult,
)
from model.predict import AESPredictor
from model.config import API_HOST, API_PORT, SAVED_MODELS_DIR_PAIR, SAVED_MODELS_DIR_PAIR_FOCAL
from .file_parser import StructureDetector


predictor = None


@asynccontextmanager
async def lifespan(app: FastAPI):
    global predictor

    print("\nStarting AES-Feedback API Server...")

    model_dir = os.path.join(SAVED_MODELS_DIR_PAIR_FOCAL, "best_model")

    if not os.path.exists(model_dir):
        print("Downloading model from arkhangelos/aes-indobert-pair-focal...")
        os.makedirs(SAVED_MODELS_DIR_PAIR_FOCAL, exist_ok=True)
        try:
            snapshot_download(
                "arkhangelos/aes-indobert-pair-focal",
                local_dir=model_dir,
            )
            print("[OK] Model downloaded.")
        except Exception as e:
            print(f"[WARN] Download failed: {e}")

    legacy_dir = os.path.join(SAVED_MODELS_DIR_PAIR, "best_model")
    if not os.path.exists(model_dir) and os.path.exists(legacy_dir):
        model_dir = legacy_dir

    if os.path.exists(model_dir):
        try:
            predictor = AESPredictor(model_dir=model_dir)
            print(f"[OK] Model loaded from {model_dir}\n")
        except Exception as e:
            print(f"[WARN] Gagal load model: {e}\n")
            predictor = None
    else:
        print(f"[WARN] Model tidak ditemukan di {model_dir}\n")
        predictor = None

    yield

    if predictor is not None:
        del predictor
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    print("AES-Feedback API Server stopped.")


app = FastAPI(
    title="BantuGuru API — AES-Feedback",
    description=(
        "Automated Essay Scoring dengan Formative Feedback.\n\n"
        "Menggunakan IndoBERT untuk prediksi skor dan IndoSBERT untuk analisis koherensi.\n"
        "Menghasilkan feedback formatif untuk 3 aspek: "
        "Argument Structure, Reasoning, dan Evidence Use."
    ),
    version="1.0.0",
    lifespan=lifespan,
)

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


@app.get("/", tags=["General"])
async def root():
    return {
        "name": "BantuGuru API — AES-Feedback",
        "version": "1.0.0",
        "docs": "/docs",
        "health": "/health",
    }


@app.get("/health", response_model=HealthResponse, tags=["General"])
async def health():
    status = "ok" if predictor is not None else "no_model_loaded"
    return HealthResponse(
        status=status,
        model_loaded=predictor is not None,
        device=str(torch.device("cuda" if torch.cuda.is_available() else "cpu")),
    )


@app.post("/api/predict", response_model=PredictResponse, tags=["Prediction"])
async def predict(request: PredictRequest):
    if predictor is None:
        raise HTTPException(
            status_code=503,
            detail="Model belum dimuat. Tunggu startup selesai atau hubungi admin.",
        )

    try:
        result = predictor.predict(
            request.essay,
            id_soal=request.id_soal,
            soal=request.soal,
            kunci_jawaban=request.kunci_jawaban,
        )

        return PredictResponse(
            overall_score=result["overall_score"],
            normalized_score=result["normalized_score"],
            max_score=result["max_score"],
            id_soal=result["id_soal"],
            coherence=result["coherence"],
            feedback=result["feedback"],
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Prediction error: {str(e)}")


@app.post("/api/predict/file", response_model=FilePredictResponse, tags=["Prediction"])
async def predict_file(file: UploadFile):
    ext = file.filename.split(".")[-1].lower() if file.filename else ""
    if ext not in ("docx", "pdf"):
        raise HTTPException(
            status_code=400,
            detail="Format file tidak didukung. Gunakan .docx atau .pdf."
        )

    if predictor is None:
        raise HTTPException(
            status_code=503,
            detail="Model belum dimuat. Tunggu startup selesai."
        )

    tmp = tempfile.NamedTemporaryFile(delete=False, suffix=f".{ext}")
    try:
        content = await file.read()
        tmp.write(content)
        tmp.close()

        result = StructureDetector.parse(tmp.name)

        file_results = []
        for qa in result.pairs:
            pred = predictor.predict(qa.jawaban, soal=qa.soal)
            matched = pred.get("matched", False)
            id_soal = pred.get("id_soal")
            file_results.append(FilePredictResult(
                soal=qa.soal[:200] if qa.soal else "",
                jawaban=qa.jawaban[:200] if qa.jawaban else "",
                id_soal=id_soal,
                matched=matched,
                prediction=PredictResponse(
                    overall_score=pred["overall_score"],
                    normalized_score=pred["normalized_score"],
                    max_score=pred["max_score"],
                    id_soal=pred["id_soal"],
                    coherence=pred["coherence"],
                    feedback=pred["feedback"],
                ) if pred else None,
            ))

        return FilePredictResponse(
            filename=file.filename or "unknown",
            parsed_as="multi_qa" if result.confidence >= 0.7 else "single",
            strategy=result.strategy,
            total_pairs=len(result.pairs),
            results=file_results,
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"File processing error: {str(e)}")
    finally:
        if os.path.exists(tmp.name):
            os.unlink(tmp.name)


if __name__ == "__main__":
    import uvicorn
    port = int(os.environ.get("PORT", API_PORT))
    print(f"Starting server on http://{API_HOST}:{port}")
    print(f"Docs: http://localhost:{port}/docs\n")
    uvicorn.run(
        "api.app:app",
        host=API_HOST,
        port=port,
        reload=False,
    )