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

from typing import Dict, Any, List, Optional
from io import BytesIO

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

import json

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

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


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


VALIDATION_INSTRUCTION = """You are a data validation assistant for a loan underwriting team.

You receive JSON data for a SINGLE applicant. Each item in `records` is a field extracted
from one document and contains:

- applicant_id
- document_type
- kpi_description  (a human-readable description of what the field means)
- field_name
- value

IMPORTANT:
- The same real-world field (e.g., applicant name, SSN, home address, employer)
  may appear with slightly different `field_name` values in different documents.
- Use BOTH `field_name` AND `kpi_description` to determine when two records represent
  the SAME underlying concept.

Examples of equivalent concepts:
- "Borrower Name" ~ "Applicant Full Name"
- "SSN" ~ "Social Security Number"
- "Home Addr" ~ "Residential Address"
- "Employer Name" ~ "Current Company"

YOUR TASK:
For this SINGLE applicant, determine whether each record's `value` is consistent
with other records that represent the same underlying concept.

You must return a JSON ARRAY where each element corresponds 1:1 to an input record,
with EXACTLY these fields:

- "applicant_id"  (string or number)
- "document_type" (string)
- "field_name"    (string)
- "value"         (string)
- "validated"     (boolean)

Rules for `validated`:
- true  -> This value is consistent with other semantically equivalent fields
           for the same applicant (either identical or clearly equivalent).
- false -> There is a clear conflict, a major discrepancy, or not enough evidence
           to be confident it matches.

CRITICAL:
- Do NOT add extra fields.
- Do NOT drop any records.
- The number of objects in the output array MUST equal the number of input records.
- The ENTIRE response MUST be valid JSON. Do NOT include code, comments, or Markdown.
""".strip()


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


def _validation_result_blob_name(applicant_id: str) -> str:
    return f"{VALIDATION_RESULTS_PREFIX}applicant_{applicant_id}_validated.csv"


def _load_applicant_kpis_from_gcs(applicant_id: str) -> pd.DataFrame:
    """Fetch aggregated KPI CSV and filter to applicant."""
    blob_name = _kpi_csv_blob_name(applicant_id)
    bucket = storage_client.bucket(GCS_BUCKET_NAME)
    blob = bucket.blob(blob_name)
    data = blob.download_as_bytes()
    df = pd.read_csv(BytesIO(data))
    return df[df["applicant_id"].astype(str) == str(applicant_id)].copy()


def _build_applicant_payload(df_applicant: pd.DataFrame) -> Dict[str, Any]:
    """Shape applicant KPI rows into the JSON payload expected by Gemini."""
    df_applicant = df_applicant.sort_values(
        ["field_name", "document_type", "kpi_description"]
    ).reset_index(drop=True)
    records = df_applicant.to_dict(orient="records")
    applicant_id = str(df_applicant["applicant_id"].iloc[0])
    return {"applicant_id": applicant_id, "records": records}


def _validate_applicant_records(model: GenerativeModel, payload: Dict[str, Any]) -> List[Dict[str, Any]]:
    """Send a single applicant payload to Gemini and return normalized results."""
    prompt = f"""{VALIDATION_INSTRUCTION}

Here is the applicant data as JSON:

{json.dumps(payload, ensure_ascii=False, default=str)}
"""
    response = model.generate_content(
        [prompt],
        generation_config={"response_mime_type": "application/json"},
    )

    raw_text = response.text or ""
    parsed = json.loads(raw_text)

    if isinstance(parsed, dict):
        parsed_list = [parsed]
    elif isinstance(parsed, list):
        parsed_list = parsed
    else:
        raise ValueError(f"Model returned JSON that is not an object or array: {type(parsed)}")

    required_keys = ["applicant_id", "document_type", "field_name", "value", "validated"]
    normalized: List[Dict[str, Any]] = []

    for item in parsed_list:
        if not isinstance(item, dict):
            continue

        norm = {k: item.get(k) for k in required_keys}

        val = norm.get("validated")
        if isinstance(val, str):
            lowered = val.lower()
            if lowered in {"true", "yes", "y", "1"}:
                norm["validated"] = True
            elif lowered in {"false", "no", "n", "0"}:
                norm["validated"] = False

        normalized.append(norm)

    return normalized


def validate_applicant_fields(
    applicant_id: str,
    kpi_df: Optional[pd.DataFrame] = None,
    save_to_gcs: bool = True,
) -> pd.DataFrame:
    """
    Uses Gemini to validate that extracted KPIs / fields are consistent across
    documents for the same applicant, mirroring applicant_info_validation_with_gemini.ipynb.

    If kpi_df is not provided, the function will load the per-applicant KPI CSV
    from GCS (output of kpi_extraction.py).
    """
    if kpi_df is None:
        kpi_df = _load_applicant_kpis_from_gcs(applicant_id)

    expected_cols = ["applicant_id", "document_type", "kpi_description", "field_name", "value"]
    missing = [c for c in expected_cols if c not in kpi_df.columns]
    if missing:
        raise ValueError(f"Missing required columns: {missing}")

    # filter to applicant_id to mimic notebook per-applicant run
    df_applicant = kpi_df[kpi_df["applicant_id"].astype(str) == str(applicant_id)].copy()
    if df_applicant.empty:
        raise ValueError(f"No KPI records found for applicant_id={applicant_id}")

    model = _init_vertex()
    payload = _build_applicant_payload(df_applicant)

    try:
        validated_records = _validate_applicant_records(model, payload)
    except Exception as exc:
        # fallback: mark all as not validated for this applicant
        validated_records = []
        for _, row in df_applicant.iterrows():
            validated_records.append(
                {
                    "applicant_id": row["applicant_id"],
                    "document_type": row["document_type"],
                    "field_name": row["field_name"],
                    "value": row["value"],
                    "validated": False,
                }
            )
        print(f"[ERROR] Validation failed for applicant_id={applicant_id}: {exc}")

    validated_df = pd.DataFrame(validated_records)

    final_cols = ["applicant_id", "document_type", "field_name", "value", "validated"]
    validated_df = validated_df[final_cols]

    if save_to_gcs:
        bucket = storage_client.bucket(GCS_BUCKET_NAME)
        blob_name = _validation_result_blob_name(applicant_id)
        blob = bucket.blob(blob_name)
        blob.upload_from_string(validated_df.to_csv(index=False), content_type="text/csv")

    return validated_df
