src/serving/api.py

"""FastAPI application for scholarship recommendation serving.

Endpoints:
    POST /recommend          — Recommend top-K scholarships for a student profile
    POST /recommend-single   — Match student to a single scholarship
    POST /parse-cv           — Parse a CV/resume (PDF or image) and extract student data
    POST /refresh            — Rebuild scholarship embedding cache from CSV (admin)
    POST /retrain            — Retrain model with new data (async, admin)
    GET  /health             — Health check (includes retraining status)

After successful retraining, data + model artifacts are pushed to HuggingFace.
"""
from __future__ import annotations

import io
import threading
from http import HTTPStatus
from typing import Optional

import pandas as pd
from fastapi import Body, Depends, FastAPI, File, HTTPException, Query, UploadFile
from fastapi.responses import HTMLResponse
import os
from fastapi.middleware.cors import CORSMiddleware
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
from pydantic import BaseModel, Field

from scripts.hf_sync import push_data_artifacts, push_model_artifacts
from src.serving.cv_parser import parse_cv
from src.serving.inference_engine import InferenceEngine
from src.serving.llm_client import LLMClient

# Bearer security scheme for auth-protected endpoints
security = HTTPBearer(auto_error=False)


def _get_auth_token(credentials: Optional[HTTPAuthorizationCredentials] = Depends(security)) -> Optional[str]:
    """Dependency to extract Bearer token from request headers."""
    return credentials.credentials if credentials else None


# ── Pydantic schemas ──────────────────────────────────────────────────────

class StudentProfile(BaseModel):
    """Student profile for scholarship recommendation.

    Required fields: nationality, age, high_school_track
    All other fields are optional with sensible defaults.
    """

    # ── Required fields ──────────────────────────────────────────────────
    nationality: str = Field(
        ...,
        description="Student's nationality (e.g. 'Indonesia', 'Malaysia')",
        examples=["Indonesia"],
    )
    age: int = Field(
        ...,
        description="Student's age in years",
        examples=[17],
        ge=15,
        le=20,
    )
    high_school_track: str = Field(
        ...,
        description="High school track (e.g. 'science', 'social')",
        examples=["science"],
    )

    # ── Academic performance (optional) ──────────────────────────────────
    overall_report_card_average: float = Field(
        default=0.0,
        description="Overall report card average score",
        examples=[85.0],
    )
    math_score: float = Field(
        default=0.0,
        description="Mathematics exam score",
        examples=[80.0],
    )
    english_score: float = Field(
        default=0.0,
        description="English exam score",
        examples=[75.0],
    )
    major_subject_average: float = Field(
        default=0.0,
        description="Average score in major subjects",
        examples=[82.0],
    )

    # ── Language & competitions (optional) ───────────────────────────────
    toefl_score: float = Field(
        default=0.0,
        description="TOEFL score (0-120)",
        examples=[65.0],
    )
    ielts_score: float = Field(
        default=0.0,
        description="IELTS score (0-9)",
        examples=[7.0],
    )
    language_proficiency: Optional[str] = Field(
        default=None,
        description="Language proficiency level or certification (deprecated, use toefl_score/ielts_score instead)",
        examples=["TOEFL 65"],
    )
    olympiad_level: str = Field(
        default="",
        description="Highest olympiad level achieved",
        examples=["national"],
    )
    olympiad_subjects: Optional[str] = Field(
        default=None,
        description="Olympiad subjects participated in",
        examples=["physics,mathematics"],
    )

    # ── Experience counts (optional) ─────────────────────────────────────
    leadership_experience_count: int = Field(
        default=0,
        description="Number of leadership positions held",
        examples=[3],
    )
    volunteer_experience_count: int = Field(
        default=0,
        description="Number of volunteer activities",
        examples=[5],
    )
    competition_wins_count: int = Field(
        default=0,
        description="Number of competition wins",
        examples=[2],
    )

    # ── Background (optional) ────────────────────────────────────────────
    school_tier: str = Field(
        default="",
        description="School tier classification",
        examples=["accredited_a"],
    )
    family_income_category: str = Field(
        default="",
        description="Family income category",
        examples=["upper_middle"],
    )
    from_underrepresented_region: bool = Field(
        default=False,
        description="Whether student is from an underrepresented region",
    )

    # ── Preferences (optional) ───────────────────────────────────────────
    intended_career_track: str = Field(
        default="",
        description="Intended career or field of study",
        examples=["computer_science"],
    )
    willing_to_return_home: bool = Field(
        default=False,
        description="Willingness to return home after studying abroad",
    )
    target_countries: Optional[str] = Field(
        default=None,
        description="Preferred destination countries",
        examples=["Japan,Germany"],
    )
    needs_full_funding: bool = Field(
        default=False,
        description="Whether full funding is required",
    )
    can_self_fund_living: bool = Field(
        default=False,
        description="Ability to self-fund living expenses",
    )

    # ── Text narratives (optional) ───────────────────────────────────────
    personal_statement: str = Field(
        default="",
        description="Personal statement or motivation letter",
        examples=["Passionate about STEM innovation and technology"],
    )
    achievements_narrative: str = Field(
        default="",
        description="Description of key achievements",
        examples=["National science olympiad participant"],
    )
    future_goals: str = Field(
        default="",
        description="Future career goals and aspirations",
        examples=["Technology entrepreneur building AI solutions"],
    )


class ScholarshipResult(BaseModel):
    """Single scholarship recommendation result."""
    scholarship_id: str
    score: float
    rank: int = 1
    recommendation: str = Field(
        default="",
        description="Personalized recommendation explaining why this scholarship is a match",
    )
    metadata: dict
    fit_scores: Optional[dict] = Field(
        default=None,
        description="Fit scores for academic, leadership, and language alignment (0-1 scale)",
    )


class RecommendationResponse(BaseModel):
    """Response for /recommend endpoint."""
    recommendations: list[ScholarshipResult]
    k: int


class SingleScholarshipResult(ScholarshipResult):
    """Single scholarship match result for /recommend-single endpoint."""
    pass


class RecommendSingleResponse(BaseModel):
    """Response for /recommend-single endpoint."""
    result: SingleScholarshipResult


class RefreshResponse(BaseModel):
    """Response for /refresh endpoint."""
    status: str = Field(
        description="Operation status",
        examples=["refreshed"],
    )
    total_scholarships: int = Field(
        description="Total number of scholarships in cache after refresh",
        examples=[150],
    )


class HealthResponse(BaseModel):
    """Response for /health endpoint."""
    status: str = Field(description="Service health status")
    student_tower_loaded: bool = Field(description="Whether student tower model is loaded")
    scholarship_tower_loaded: bool = Field(description="Whether scholarship tower model is loaded")
    cached_scholarships: int = Field(description="Number of scholarships in cache")
    llm: str = Field(
        description="LLM availability status ('available' or 'unavailable')",
    )
    retraining: RetrainingInfo = Field(
        description="Current retraining job status and metadata",
    )


class RetrainStartResponse(BaseModel):
    """Response for /retrain endpoint (immediate, async)."""
    status: str
    message: str


class RetrainingInfo(BaseModel):
    """Retraining status info."""
    status: str  # idle | training | done | error
    started_at: Optional[str] = None
    finished_at: Optional[str] = None
    error: Optional[str] = None


# ── CV Parser schemas ─────────────────────────────────────────────────────

class LanguageCertificate(BaseModel):
    """A language test result (e.g. IELTS, TOEFL)."""
    test_type: str = Field(description="Test type", examples=["IELTS", "TOEFL"])
    score: Optional[float] = Field(default=None, description="Test score")
    valid_until: Optional[str] = Field(
        default=None,
        description="Expiry date (YYYY-MM-DD)",
        examples=["2026-12-31"],
    )


class CVPersonalInfo(BaseModel):
    """Personal information extracted from CV."""
    full_name: Optional[str] = Field(default=None)
    gender: Optional[str] = None
    date_of_birth: Optional[str] = None
    province: Optional[str] = None
    economic_background: Optional[str] = None
    from_underrepresented_region: Optional[bool] = None


class CVAcademicInfo(BaseModel):
    """Academic information extracted from CV."""
    school_level: Optional[str] = None
    major_program: Optional[str] = None
    grade_class: Optional[str] = None
    school_name: Optional[str] = None
    school_tier_accreditation: Optional[str] = None
    expected_graduation_year: Optional[int] = None
    average_grade: Optional[float] = None
    math_score: Optional[float] = None
    english_score: Optional[float] = None
    major_subject_average: Optional[float] = None
    extracurricular_achievements: Optional[str] = None
    olympiad_level: Optional[str] = None
    intended_career_track: Optional[str] = None
    willing_to_return_home: Optional[bool] = None
    needs_full_funding: Optional[bool] = None


class CVSkillsInfo(BaseModel):
    """Skills information extracted from CV."""
    hard_skills: list[str] = Field(default_factory=list)
    soft_skills: list[str] = Field(default_factory=list)
    languages: list[str] = Field(default_factory=list)
    language_certificates: list[LanguageCertificate] = Field(default_factory=list)
    target_countries: list[str] = Field(default_factory=list)


class ParsedCVResponse(BaseModel):
    """Response for the /parse-cv endpoint."""
    personal: CVPersonalInfo
    academic: CVAcademicInfo
    skills: CVSkillsInfo
    parsing_note: Optional[str] = Field(
        default=None,
        description="Note about the parsing process (e.g., 'PDF text-only', 'image-based')",
    )


# ── Application factory ───────────────────────────────────────────────────

def create_app(engine: InferenceEngine) -> FastAPI:
    """Create and configure the FastAPI application.

    Args:
        engine: A warmed-up InferenceEngine instance.

    Returns:
        Configured FastAPI application.
    """
    # Initialize LLM client for recommendation text generation
    llm = LLMClient(engine.cfg if engine.cfg else engine._load_config())

    app = FastAPI(
        title="ScholarshipID Recommendation API",
        description="Two-tower retrieval model for matching students to scholarships.",
        version="0.1.0",
    )

    # ── CORS Middleware ──────────────────────────────────────────────────
    app.add_middleware(
        CORSMiddleware,
        allow_origins=[engine.server_config.cors_origins]
            if engine.server_config.cors_origins != "*"
            else ["*"],
        allow_credentials=True,
        allow_methods=["*"],
        allow_headers=["*"],
    )

    # ── Endpoints ────────────────────────────────────────────────────────

    @app.get("/", response_class=HTMLResponse)
    async def root():
        """Landing page for ScholarshipID API — served from static/index.html."""
        html_path = os.path.join(os.path.dirname(__file__), "static", "index.html")
        with open(html_path, "r", encoding="utf-8") as f:
            return HTMLResponse(content=f.read(), status_code=200)

    @app.post("/recommend", response_model=RecommendationResponse)
    async def recommend(
        auth_token: Optional[str] = Depends(_get_auth_token),
        student: StudentProfile = Body(..., description="Student profile for recommendation"),
        k: int = Query(5, ge=1, le=50),
    ):
        """Return top-K scholarships for the given student profile.

        The student body is encoded through the student tower, then matched
        against cached scholarship embeddings via dot-product (cosine similarity).
        Note: This service uses a hosted Local AI as the LLM provider. Contact yusrmuttaqien to enable the LLM service.
        """
        # Auth check
        if engine.server_config.auth_required:
            if auth_token is None or auth_token != engine.server_config.auth_token:
                raise HTTPException(
                    status_code=HTTPStatus.UNAUTHORIZED,
                    detail="Invalid or missing authorization token",
                )
        try:
            results = engine.recommend(student.model_dump(), k=k)
        except Exception as exc:
            raise HTTPException(status_code=500, detail=str(exc))

        # Enrich results with LLM-generated fit scores and recommendations.
        # Batch scholarship metadata and call LLM once per batch (BATCH_SIZE=10)
        # to minimize round-trips to the single-slot llama.cpp server.
        student_data = student.model_dump()
        BATCH_SIZE = 10

        if llm.is_available:
            all_metadata = [r.get("metadata", {}) for r in results]
            for batch_start in range(0, len(all_metadata), BATCH_SIZE):
                batch_meta = all_metadata[batch_start:batch_start + BATCH_SIZE]
                print(f"[LLM] Processing batch {batch_start // BATCH_SIZE + 1} "
                      f"({len(batch_meta)} scholarships)...", flush=True)
                fit_scores_list, rec_list = llm.generate_batch_recommendation(
                    student_data, batch_meta
                )
                for i, idx in enumerate(range(batch_start, min(batch_start + len(batch_meta), len(results)))):
                    results[idx]["fit_scores"] = fit_scores_list[i]
                    results[idx]["recommendation"] = rec_list[i]
                    if not rec_list[i]:
                        print(f"[LLM] Empty recommendation for {results[idx].get('scholarship_id')}", flush=True)

        return RecommendationResponse(
            recommendations=[ScholarshipResult(**r) for r in results],
            k=k,
        )

    @app.post("/recommend-single", response_model=RecommendSingleResponse)
    async def recommend_single(
        auth_token: Optional[str] = Depends(_get_auth_token),
        student: StudentProfile = Body(..., description="Student profile for single scholarship match"),
        scholarship_id: str = Query(
            ...,
            description="The ID of the scholarship to evaluate",
            examples=["LPDP-2024-001"],
        ),
    ):
        """Return match score for a single scholarship given a student profile.

        The student is encoded through the student tower, then matched against
        the specific scholarship embedding via dot-product (cosine similarity).
        Note: This service uses a hosted Local AI as the LLM provider. Contact yusrmuttaqien to enable the LLM service.
        """
        # Auth check
        if engine.server_config.auth_required:
            if auth_token is None or auth_token != engine.server_config.auth_token:
                raise HTTPException(
                    status_code=HTTPStatus.UNAUTHORIZED,
                    detail="Invalid or missing authorization token",
                )
        try:
            score_data = engine.get_scholarship_score(
                scholarship_id, student.model_dump()
            )
        except Exception as exc:
            raise HTTPException(status_code=500, detail=str(exc))

        if score_data is None:
            raise HTTPException(
                status_code=404,
                detail=f"Scholarship with id '{scholarship_id}' not found",
            )

        # Enrich with LLM-generated recommendation and fit scores.
        # Use the combined method (fit_scores + recommendation in one LLM call)
        # to halve the number of round-trips to the single-slot llama.cpp server.
        student_data = student.model_dump()
        metadata = score_data.get("metadata", {})
        if llm.is_available:
            try:
                print(f"[LLM] Processing {scholarship_id}", flush=True)
                fit_scores, rec = llm.generate_recommendation_with_fit_scores(student_data, metadata)
                score_data["fit_scores"] = fit_scores
                score_data["recommendation"] = rec
                if not rec:
                    print(f"[LLM] Empty recommendation for {scholarship_id}", flush=True)
            except Exception as e:
                print(
                    f"[LLM] Enrichment failed for {scholarship_id}: {e}", flush=True
                )

        # Ensure fit_scores is present (even without LLM)
        if "fit_scores" not in score_data:
            score_data["fit_scores"] = None

        return RecommendSingleResponse(
            result=SingleScholarshipResult(**score_data),
        )

    @app.post("/refresh", response_model=RefreshResponse)
    async def refresh(
        auth_token: Optional[str] = Depends(_get_auth_token),
        scholarships: UploadFile = File(..., description="Scholarship data CSV file"),
    ):
        """Rebuild the scholarship embedding cache from an uploaded CSV file.

        Call this endpoint after adding new scholarships to the data source.
        The model will use updated embeddings on the next /recommend call.

        New scholarships are merged into existing ones by scholarship_id (new replaces old).

        Requires admin authentication when auth_required is enabled.
        """
        # Validate file extension
        if not scholarships.filename or not scholarships.filename.lower().endswith(".csv"):
            raise HTTPException(
                status_code=422,
                detail="Only .csv files are allowed",
            )

        # Auth check
        if engine.server_config.auth_required:
            if auth_token is None or auth_token != engine.server_config.auth_token:
                raise HTTPException(
                    status_code=HTTPStatus.UNAUTHORIZED,
                    detail="Invalid or missing authorization token",
                )

        try:
            content = await scholarships.read()
            csv_text = content.decode("utf-8")

            # Incrementally merge new scholarships with existing ones (new replaces old by ID)
            cfg = engine.cfg if engine.cfg else engine._load_config()
            raw_path = cfg["data"]["raw_path"]
            existing_df = pd.read_csv(f"{raw_path}/scholarships.csv")
            new_df = pd.read_csv(io.StringIO(csv_text))
            merged_df = pd.concat([existing_df, new_df]).drop_duplicates(
                subset=["scholarship_id"], keep="last"
            )
            merged_df.to_csv(f"{raw_path}/scholarships.csv", index=False)

            # Rebuild in-memory embedding cache from the merged CSV on disk
            engine.refresh_from_csv(merged_df.to_csv(index=False))
        except Exception as exc:
            raise HTTPException(status_code=500, detail=str(exc))

        # Push updated data artifacts to HuggingFace on success (data only)
        try:
            push_data_artifacts(
                config_path=engine.config_path,
                message="Auto-push after /refresh",
            )
        except Exception as e:
            print(f"Warning: HF push failed after refresh: {e}", flush=True)

        return RefreshResponse(
            status="refreshed",
            total_scholarships=len(engine._sch_ids),
        )

    @app.post("/retrain", response_model=RetrainStartResponse)
    async def retrain(
        auth_token: Optional[str] = Depends(_get_auth_token),
        students: Optional[UploadFile] = File(None),
        scholarships: Optional[UploadFile] = File(None),
        feedbacks: Optional[UploadFile] = File(None),
    ):
        """Initiate model retraining with new data.

        Accepts 3 optional CSV files in the request body (multipart/form-data):
        - students: New student records (by student_id)
        - scholarships: New scholarship records (by scholarship_id)
        - feedbacks: New feedback records (by timestamp or feedback_id)

        At least one file is required. Training runs asynchronously in a background thread.
        Requires admin authentication when auth_required is enabled.

        The retraining process:
        1. Merges new data with existing disk data by ID
        2. Saves merged data back to disk
        3. Precomputes text embeddings on full merged datasets
        4. Trains model (finetuning from existing weights) on all feedback data
        5. Exports updated scholarship embeddings

        After completion, call /recommend without refresh — the engine auto-loads
        the latest embeddings from disk.
        """
        # Validate at least one file provided
        if not any([students, scholarships, feedbacks]):
            raise HTTPException(
                status_code=422,
                detail="At least one file is required (students, scholarships, or feedbacks)",
            )

        # Auth check
        if engine.server_config.auth_required:
            if auth_token is None or auth_token != engine.server_config.auth_token:
                raise HTTPException(
                    status_code=HTTPStatus.UNAUTHORIZED,
                    detail="Invalid or missing authorization token",
                )

        # Check if retraining is already in progress
        retrain_info = engine.get_retraining_status()
        if retrain_info["status"] == "training":
            raise HTTPException(
                status_code=409,
                detail="Retraining already in progress. Check /health for status.",
            )

        def _do_retrain():
            csv_parts = {}
            if students and students.filename:
                csv_parts["students"] = students.file.read().decode("utf-8")
            if scholarships and scholarships.filename:
                csv_parts["scholarships"] = scholarships.file.read().decode("utf-8")
            if feedbacks and feedbacks.filename:
                csv_parts["feedbacks"] = feedbacks.file.read().decode("utf-8")

            result = engine.retrain_from_csvs(
                students_csv_text=csv_parts.get("students", None),
                scholarships_csv_text=csv_parts.get("scholarships", None),
                feedbacks_csv_text=csv_parts.get("feedbacks", None),
            )

            # Push updated data + model artifacts to HuggingFace on success
            if result.get("status") == "done":
                try:
                    push_data_artifacts(
                        config_path=engine.config_path,
                        message="Auto-push after API retraining",
                    )
                    push_model_artifacts(
                        config_path=engine.config_path,
                        message="Auto-push after API retraining",
                    )
                except Exception as e:
                    print(f"Warning: HF push failed: {e}", flush=True)

        # Start training in background thread
        threading.Thread(target=_do_retrain, daemon=True).start()

        return RetrainStartResponse(
            status="training_started",
            message="Retraining initiated. Check /health for status.",
        )

    # ── CV Parser endpoint ────────────────────────────────────────────────

    _ALLOWED_CV_EXTENSIONS = {".pdf", ".png", ".jpg", ".jpeg", ".webp"}
    _MAX_FILE_SIZE = 10 * 1024 * 1024  # 10 MB

    @app.post("/parse-cv", response_model=ParsedCVResponse)
    async def parse_cv_endpoint(
        auth_token: Optional[str] = Depends(_get_auth_token),
        file: UploadFile = File(..., description="CV or resume file (PDF, PNG, JPG, WEBP)"),
    ):
        """Parse a CV/resume and extract student profile data using multimodal LLM.

        Accepts PDF files (text-based or scanned) and image files (PNG, JPG, WEBP).
        Returns structured student data matching the frontend form schema.
        Note: This service uses a hosted Local AI as the LLM provider. Contact yusrmuttaqien to enable the LLM service.
        """
        # Auth check
        if engine.server_config.auth_required:
            if auth_token is None or auth_token != engine.server_config.auth_token:
                raise HTTPException(
                    status_code=HTTPStatus.UNAUTHORIZED,
                    detail="Invalid or missing authorization token",
                )

        # Validate file extension
        if not file.filename:
            raise HTTPException(status_code=422, detail="No filename provided")

        ext = file.filename.lower().rsplit(".", 1)[-1] if "." in file.filename else ""
        if f".{ext}" not in _ALLOWED_CV_EXTENSIONS:
            allowed = ", ".join(_ALLOWED_CV_EXTENSIONS)
            raise HTTPException(
                status_code=422,
                detail=f"Unsupported file type. Allowed: {allowed}",
            )

        # Read and validate size
        file_bytes = await file.read()
        if len(file_bytes) == 0:
            raise HTTPException(status_code=422, detail="Empty file")

        if len(file_bytes) > _MAX_FILE_SIZE:
            raise HTTPException(
                status_code=422,
                detail=f"File too large. Maximum size is {_MAX_FILE_SIZE // (1024*1024)} MB",
            )

        # Parse CV using multimodal LLM
        try:
            parsed = parse_cv(llm, file_bytes, filename=file.filename)
        except Exception as exc:
            raise HTTPException(status_code=500, detail=f"CV parsing failed: {exc}")

        if not parsed:
            raise HTTPException(
                status_code=502,
                detail="CV parsing requires the Local AI LLM service, which is not currently configured. Contact yusrmuttaqien to enable.",
            )

        # Normalize the parsed dict to match our Pydantic model
        personal = parsed.get("personal", {}) or {}
        academic = parsed.get("academic", {}) or {}
        skills = parsed.get("skills", {}) or {}

        # Ensure language_certificates is a list of dicts
        lang_certs_raw = skills.get("language_certificates") or []
        if isinstance(lang_certs_raw, list):
            lang_certs = [
                LanguageCertificate(**c) if isinstance(c, dict) else LanguageCertificate()
                for c in lang_certs_raw
            ]
        else:
            lang_certs = []

        return ParsedCVResponse(
            personal=CVPersonalInfo(**personal),
            academic=CVAcademicInfo(**academic),
            skills=CVSkillsInfo(
                hard_skills=skills.get("hard_skills") or [],
                soft_skills=skills.get("soft_skills") or [],
                languages=skills.get("languages") or [],
                language_certificates=lang_certs,
                target_countries=skills.get("target_countries") or [],
            ),
            parsing_note=parsed.get("_parsing_note"),
        )

    @app.get("/health", response_model=HealthResponse)
    async def health():
        """Health check — returns model loading status, LLM availability, and retraining info."""
        retrain_info = engine.get_retraining_status()

        llm_status = "available" if llm.is_reachable() else "unavailable"

        return HealthResponse(
            status="healthy",
            student_tower_loaded=engine.student_tower is not None,
            scholarship_tower_loaded=engine.scholarship_tower is not None,
            cached_scholarships=len(engine._sch_ids),
            llm=llm_status,
            retraining=retrain_info,
        )

    return app