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

import json
from pathlib import Path
from typing import Optional

from google.cloud import storage, vision

from config import (
    GCP_PROJECT_ID,
    GCS_BUCKET_NAME,
    PDF_INPUT_PREFIX,
    VISION_OUTPUT_PREFIX,
    OCR_RESULTS_PREFIX,
)

storage_client = storage.Client(project=GCP_PROJECT_ID) if GCP_PROJECT_ID else storage.Client()
vision_client = vision.ImageAnnotatorClient()


def _normalize_prefix(prefix: str) -> str:
    """Ensure GCS folder-style prefixes end with a slash."""
    return prefix if prefix.endswith("/") else f"{prefix}/"


def _applicant_prefix(applicant_id: str) -> str:
    # PDFs live under applicant-documents-pdf/applicant_1/, etc.
    return f"{_normalize_prefix(PDF_INPUT_PREFIX)}applicant_{applicant_id}/"


def _ocr_result_blob_name(applicant_id: str) -> str:
    # You can change this if your notebook writes something else
    return f"{_normalize_prefix(OCR_RESULTS_PREFIX)}applicant_{applicant_id}.json"


def get_existing_ocr_result_uri(applicant_id: str) -> Optional[str]:
    """Check if combined OCR JSON already exists in GCS."""
    bucket = storage_client.bucket(GCS_BUCKET_NAME)
    blob_name = _ocr_result_blob_name(applicant_id)
    blob = bucket.blob(blob_name)
    if blob.exists():
        return f"gs://{GCS_BUCKET_NAME}/{blob_name}"
    return None


def _list_applicant_pdf_blobs(bucket: storage.Bucket, applicant_id: str):
    """Return PDF blobs for an applicant, mirroring the notebook's scan logic."""
    prefix = _applicant_prefix(applicant_id)
    pdf_exts = {".pdf"}
    blobs = []

    for blob in bucket.list_blobs(prefix=prefix):
        if blob.name.endswith("/"):
            continue
        if Path(blob.name).suffix.lower() not in pdf_exts:
            continue
        blobs.append(blob)

    return blobs


def _ocr_pdf_gcs(gcs_source_uri: str, gcs_destination_uri: str, timeout: int = 600) -> str:
    """Run async Vision OCR on a PDF in GCS and return combined text for all pages."""
    mime_type = "application/pdf"
    feature = vision.Feature(type_=vision.Feature.Type.DOCUMENT_TEXT_DETECTION)

    gcs_source = vision.GcsSource(uri=gcs_source_uri)
    input_config = vision.InputConfig(gcs_source=gcs_source, mime_type=mime_type)

    gcs_destination = vision.GcsDestination(uri=gcs_destination_uri)
    output_config = vision.OutputConfig(gcs_destination=gcs_destination, batch_size=10)

    async_request = vision.AsyncAnnotateFileRequest(
        features=[feature],
        input_config=input_config,
        output_config=output_config,
    )

    operation = vision_client.async_batch_annotate_files(requests=[async_request])
    operation.result(timeout=timeout)

    dest_bucket, dest_prefix = gcs_destination_uri.replace("gs://", "").split("/", 1)
    blobs = list(storage_client.list_blobs(dest_bucket, prefix=dest_prefix))
    if not blobs:
        raise RuntimeError(f"No OCR output files found at {gcs_destination_uri}")

    full_text: list[str] = []
    for blob in blobs:
        json_str = blob.download_as_text(encoding="utf-8")
        result = json.loads(json_str)
        for resp in result.get("responses", []):
            annotation = resp.get("fullTextAnnotation")
            if annotation and "text" in annotation:
                full_text.append(annotation["text"])

    return "\n".join(full_text)


def run_ocr_for_applicant(applicant_id: str) -> str:
    """
    Runs OCR for all PDFs for a given applicant_id, writes a combined JSON
    to GCS, and returns its gs:// URI.

    This function is where you paste/adapt the core logic from
    gcs_pdf_vision_ocr_capstone.ipynb:
      - list PDFs under applicant prefix
      - call Vision API for each PDF
      - aggregate page text into a dict/list
      - save a JSON file to OCR_RESULTS_PREFIX
    """
    bucket = storage_client.bucket(GCS_BUCKET_NAME)
    input_prefix = _applicant_prefix(applicant_id)

    pdf_blobs = _list_applicant_pdf_blobs(bucket, applicant_id)
    if not pdf_blobs:
        raise ValueError(f"No PDFs found under prefix {input_prefix}")

    combined_result: dict = {"applicant_id": applicant_id, "documents": {}}

    for blob in pdf_blobs:
        doc_basename = Path(blob.name).stem
        gcs_source_uri = f"gs://{GCS_BUCKET_NAME}/{blob.name}"

        dest_prefix = f"{_normalize_prefix(VISION_OUTPUT_PREFIX)}applicant_{applicant_id}/{doc_basename}/"
        gcs_destination_uri = f"gs://{GCS_BUCKET_NAME}/{dest_prefix}"

        try:
            text = _ocr_pdf_gcs(gcs_source_uri, gcs_destination_uri)
        except Exception as exc:  # pragma: no cover - defensive log
            # Mirror notebook behavior: continue even if one document fails
            print(f"Error OCR'ing {gcs_source_uri}: {exc}")
            text = ""

        combined_result["documents"][doc_basename] = text

    blob_name = _ocr_result_blob_name(applicant_id)
    blob = bucket.blob(blob_name)
    blob.upload_from_string(json.dumps(combined_result, ensure_ascii=False, indent=2), content_type="application/json")

    return f"gs://{GCS_BUCKET_NAME}/{blob_name}"
