# app/pipelines/kpi_extraction.py
from __future__ import annotations

import json
from typing import Dict, List, Any

import pandas as pd
from google.cloud import storage
import vertexai
from vertexai.generative_models import GenerativeModel
from io import BytesIO

from config import (
    GCP_PROJECT_ID,
    GCS_BUCKET_NAME,
    KPI_RESULTS_PREFIX,
    VERTEX_MODEL_NAME,
    GCP_REGION,
)

storage_client = storage.Client(project=GCP_PROJECT_ID) if GCP_PROJECT_ID else storage.Client()
_KPI_FINAL_FILE = "applicant_extracted_data.csv"


def _kpi_result_blob_name(applicant_id: str) -> str:
    # Aggregated KPI file holding all applicants
    return f"{KPI_RESULTS_PREFIX}{_KPI_FINAL_FILE}"


def get_existing_kpi_result_uri(applicant_id: str) -> str | None:
    bucket = storage_client.bucket(GCS_BUCKET_NAME)
    blob = bucket.blob(_kpi_result_blob_name(applicant_id))
    if blob.exists():
        return f"gs://{GCS_BUCKET_NAME}/{blob.name}"
    return None


def _load_existing_kpis() -> pd.DataFrame | None:
    """Load aggregated KPI CSV if it exists; return None if missing."""
    bucket = storage_client.bucket(GCS_BUCKET_NAME)
    blob = bucket.blob(_kpi_result_blob_name("ignored"))
    if not blob.exists():
        return None
    data = blob.download_as_bytes()
    return pd.read_csv(BytesIO(data))


def _load_ocr_json(gs_uri: str) -> Dict[str, Any]:
    """Download and parse the OCR JSON for one applicant (output of ocr_pipeline)."""
    if not gs_uri.startswith("gs://"):
        raise ValueError(f"Expected gs:// URI, got {gs_uri}")
    bucket_name, blob_name = gs_uri[5:].split("/", 1)

    bucket = storage_client.bucket(bucket_name)
    blob = bucket.blob(blob_name)
    content = blob.download_as_text()
    return json.loads(content)


def _init_vertex():
    vertexai.init(project=GCP_PROJECT_ID, location=GCP_REGION)
    return GenerativeModel(VERTEX_MODEL_NAME)


def _prep_kpi_list(config_df: pd.DataFrame) -> List[str]:
    """Mirror notebook: first col = document_type, second = kpi_description."""
    if config_df.shape[1] < 2:
        raise ValueError("KPI config must have at least 2 columns (document_type, kpi_description).")

    doc_col, kpi_col = config_df.columns[:2]
    df = config_df.rename(columns={doc_col: "document_type", kpi_col: "kpi_description"})
    df["document_type"] = df["document_type"].astype(str).str.strip()
    df["kpi_description"] = df["kpi_description"].astype(str).str.strip()

    return (
        df["kpi_description"]
        .dropna()
        .astype(str)
        .str.strip()
        .drop_duplicates()
        .tolist()
    )


def _call_gemini_for_document(
    model: GenerativeModel,
    document_label: str,
    ocr_text: str,
    kpi_descriptions: List[str],
    max_ocr_chars: int = 12000,
    max_kpis: int = 200,
) -> List[Dict[str, Any]]:
    """Port of call_gemini_for_document from the notebook."""
    ocr_text_input = (ocr_text or "")[:max_ocr_chars]
    kpis_input = kpi_descriptions[:max_kpis]

    payload = {
        "document_label": document_label,
        "kpi_descriptions": kpis_input,
        "ocr_text": ocr_text_input,
    }

    system_msg = """
You are an assistant for a loan document validation system.

You are given:
- A document label (e.g., "Disbursement memo", "Loan Application Form", "APPLICATION FORM_46").
- OCR text for that document.
- A list of KPI descriptions that describe checks or pieces of information
  that might appear in *some* document types.

Your task for THIS SINGLE DOCUMENT is:
1. Decide which KPI descriptions are actually relevant to this document,
   based on the document label and the OCR text.
2. For each relevant KPI description:
   - Return the KPI description verbatim.
   - Infer a concise field name that represents the main value to extract.
   - Extract the best-guess value from the OCR text.
3. If a KPI clearly does NOT apply to this document, do not include it
   in the output.
4. If a KPI seems relevant but the value isn't clearly present, include it
   with value = null.
5. If a field like Marital status has a checkmark for any of the values (i.e. Single, Married, Others)
6. If Age is a field that needs to be extracted, it's usually followed by the "Yrs" suffix, but if it does not appear. Use the date of birth and substract it to today's date .
7. If the field is Years in Current Employment ,  look for the value of 'YEARS IN PRESENT EMPLOYMENT'
8. If Annual Income is the field, look for synonyms like Salary, etc. Make sure the returned value is Annual (meaning if Income is monthly, convert it to annual)
9. If the  document in ocr_text is APLICANT FORM_# make sure to extract the values from document_label = 'Loan Application Form'
10. For all documents, try to Extract the following fields from ANY part of the text:

    Annual Income

    Applicant Age

    Marital Status

    House Ownership (Owned / Rented / Mortgage / Living with Parents)

    Occupation

    Residence City

    Residence State

    Years in Current Employment

    Important Extraction Rules

    Look for synonyms and variations.

    Examples:

    Income, Monthly Salary, Annual Earnings, Gross Pay, etc.

    Age, DOB  convert DOB to age if the year is present.

    House ownership, Living status, Residence type, etc.

    If multiple values appear, choose the most recent or most explicit.

    Normalize outputs into clean text (e.g., Married, Single), not raw OCR strings.

    If a field is impossible to determine, output null.


Return STRICT JSON with this exact structure:
{
  "document_label": "string",
  "kpis": [
    {
      "kpi_description": "string",
      "field_name": "string",
      "value": "string or null"
    },
    ...
  ]
}

Do not add any other top-level keys. Do not include explanations.
""".strip()

    prompt = (
        system_msg
        + "\n\n--- INPUT JSON ---\n"
        + json.dumps(payload, ensure_ascii=False)
        + "\n------------------\n"
        "Respond ONLY with the JSON object in the required format."
    )

    response = model.generate_content(prompt)
    raw = (response.text or "").strip()

    try:
        parsed = json.loads(raw)
    except Exception:
        start = raw.find("{")
        end = raw.rfind("}") + 1
        parsed = json.loads(raw[start:end])

    kpi_entries = parsed.get("kpis", []) or []

    records: List[Dict[str, Any]] = []
    for item in kpi_entries:
        kpi_desc = str(item.get("kpi_description", "")).strip()
        field_name = str(item.get("field_name", "")).strip() or kpi_desc
        value = item.get("value", None)
        value_str = "NA" if value is None else (str(value).strip() or "NA")

        if not kpi_desc:
            continue

        records.append(
            {
                "kpi_description": kpi_desc,
                "field_name": field_name,
                "value": value_str,
            }
        )

    return records


def extract_kpis_from_ocr(
    applicant_id: str,
    ocr_gs_uri: str,
    config_df: pd.DataFrame,
    save_to_gcs: bool = True,
    reuse_existing: bool = True,
) -> pd.DataFrame:
    """
    Core logic from gemini_llm_kpi_anymatch_extraction.ipynb.

    - Load OCR JSON for this applicant.
    - For each document & each KPI description in config_df:
      prompt Gemini to extract the KPI if present.
    - Return a DataFrame with columns at least:
        ['applicant_id', 'document_type', 'kpi_description', 'field_name', 'value']
    """
    if reuse_existing:
        existing = _load_existing_kpis()
        if existing is not None:
            filtered = existing[existing["applicant_id"].astype(str) == str(applicant_id)].copy()
            if not filtered.empty:
                return filtered
    ocr_data = _load_ocr_json(ocr_gs_uri)

    model = _init_vertex()
    all_kpis = _prep_kpi_list(config_df)

    all_records: List[Dict[str, Any]] = []

    applicant_from_json = ocr_data.get("applicant_id")
    documents = ocr_data.get("documents") or {}
    if not isinstance(documents, dict):
        raise ValueError(f"Unexpected documents structure in OCR JSON: {type(documents)}")

    for doc_type, text in documents.items():
        document_label = str(doc_type)
        extracted = _call_gemini_for_document(
            model=model,
            document_label=document_label,
            ocr_text=text or "",
            kpi_descriptions=all_kpis,
        )

        for entry in extracted:
            all_records.append(
                {
                    "applicant_id": applicant_from_json or applicant_id,
                    "document_type": document_label,
                    "kpi_description": entry["kpi_description"],
                    "field_name": entry["field_name"],
                    "value": entry["value"],
                }
            )

    results_df = pd.DataFrame(all_records)

    # Ensure standard columns
    if "applicant_id" not in results_df.columns:
        results_df["applicant_id"] = applicant_id

    if save_to_gcs:
        # Merge into aggregated CSV, replacing any existing rows for this applicant
        existing = _load_existing_kpis()
        if existing is not None:
            existing = existing[existing["applicant_id"].astype(str) != str(applicant_id)]
            merged = pd.concat([existing, results_df], ignore_index=True)
        else:
            merged = results_df

        bucket = storage_client.bucket(GCS_BUCKET_NAME)
        blob_name = _kpi_result_blob_name(applicant_id)
        blob = bucket.blob(blob_name)
        csv_bytes = merged.to_csv(index=False).encode("utf-8")
        blob.upload_from_string(csv_bytes, content_type="text/csv")

    return results_df
