# app/main.py
from __future__ import annotations

from typing import Optional
from pathlib import Path

from fastapi import FastAPI, HTTPException, Query
from pydantic import BaseModel
import pandas as pd

from ocr_pipeline import get_existing_ocr_result_uri, run_ocr_for_applicant
from kpi_extraction import extract_kpis_from_ocr, _load_existing_kpis
from applicant_validation import validate_applicant_fields
from risk_model import predict_risk_for_applicant

# Load KPI config from Excel once (first two columns: document_type, kpi_description)
_CONFIG_PATH = Path(__file__).resolve().parent.parent.parent / "Pre Disbursement Checklist.xlsx"
if not _CONFIG_PATH.exists():
    raise FileNotFoundError(f"KPI config Excel not found at {_CONFIG_PATH}")
config_df = pd.read_excel(_CONFIG_PATH, sheet_name="Pre-Disbursement")


class ModelOutput(BaseModel):
    default_probability: float


class ApplicantResponse(BaseModel):
    applicant_id: str
    ocr_uri: str
    kpi_rows: int
    validated_rows: int
    model_output: ModelOutput


app = FastAPI(title="Loan Applicant Validation & Risk API")


@app.get("/health")
def health():
    return {"status": "ok"}


@app.get("/applicants/{applicant_id}", response_model=ApplicantResponse)
def process_applicant(
    applicant_id: str,
    force_rerun_ocr: bool = Query(False),
    force_rerun_kpi: bool = Query(False),
):
    print(f"Processing applicant {applicant_id}")

    # 1. OCR
    if force_rerun_ocr:
        print("Step 1: OCR - forcing rerun")
    else:
        print("Step 1: OCR - checking for existing result")
    ocr_uri = None if force_rerun_ocr else get_existing_ocr_result_uri(applicant_id)

    if not ocr_uri:
        try:
            print("Step 1: OCR - running pipeline")
            ocr_uri = run_ocr_for_applicant(applicant_id)
        except Exception as e:
            raise HTTPException(status_code=500, detail=f"OCR failed: {e}")
    else:
        print(f"Step 1: OCR - reusing existing result at {ocr_uri}")

    # 2. KPI extraction
    try:
        if not force_rerun_kpi:
            print("Step 2: KPI extraction - loading existing KPI rows")
            existing_kpi = _load_existing_kpis()
            if existing_kpi is not None:
                kpi_df = existing_kpi[existing_kpi["applicant_id"].astype(str) == str(applicant_id)].copy()
                if kpi_df.empty:
                    raise ValueError("No KPI rows found for applicant in aggregated file; rerun KPI extraction.")
            else:
                raise ValueError("Aggregated KPI file not found; rerun KPI extraction.")
        else:
            print("Step 2: KPI extraction - running from OCR output")
            kpi_df = extract_kpis_from_ocr(
                applicant_id=applicant_id,
                ocr_gs_uri=ocr_uri,
                config_df=config_df,
                save_to_gcs=True,
                reuse_existing=False,
            )
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"KPI extraction failed: {e}")

    # 3. Validation
    try:
        print("Step 3: Validation - validating KPI rows")
        validated_df = validate_applicant_fields(applicant_id=applicant_id, kpi_df=kpi_df, save_to_gcs=True)
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Validation failed: {e}")

    # 4. Model prediction
    try:
        print("Step 4: Risk model - running prediction")
        model_out = predict_risk_for_applicant(applicant_id=applicant_id, validated_df=validated_df)
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Model prediction failed: {e}")

    print("Completed request")
    return ApplicantResponse(
        applicant_id=applicant_id,
        ocr_uri=ocr_uri,
        kpi_rows=len(kpi_df),
        validated_rows=len(validated_df),
        model_output=ModelOutput(**model_out),
    )
