"""ARDEM LLM inference adapter for the existing Pennsylvania DB pipeline.

The public functions intentionally match ``inference.py`` so ``main_db.py``
can keep its S3, validation, database, dashboard, audit, and email workflow.
"""

from __future__ import annotations

import hashlib
import json
import os
import shutil
import tempfile
import uuid
from pathlib import Path
from typing import Any, Dict, Iterable, List

import boto3
from PIL import Image

from new_inference_main import infer_from_api
from prompts.pa_prompt import PA_MODEL_PROMPT


STATE = "PA"
ARDEM_S3_BUCKET = os.getenv("ARDEM_S3_BUCKET", "ardemllm")
ARDEM_INPUT_PREFIX = os.getenv("ARDEM_INPUT_PREFIX", "VITU/ardem-input").strip("/")
ARDEM_OUTPUT_PREFIX = os.getenv("ARDEM_OUTPUT_PREFIX", "VITU/ardem-output").strip("/")

CRITICAL_FIELDS = {
    "Taxable Selling Price / Purchase Price",
    "Process Date",
    "Sales Tax",
    "Plate Fee",
    "Motor Vehicle Fees",
    "Local Use / County Fee",
}


def _page_number(path: str) -> int:
    stem = Path(path).stem.lower()
    marker = "_page_"
    if marker in stem:
        try:
            return int(stem.rsplit(marker, 1)[1])
        except ValueError:
            pass
    return 10**9


def _image_to_pdf(image_path: str, destination: Path) -> None:
    with Image.open(image_path) as image:
        image.convert("RGB").save(destination, "PDF", resolution=300.0)


def _s3_client():
    return boto3.client(
        "s3",
        aws_access_key_id=os.getenv("AWS_ACCESS_KEY"),
        aws_secret_access_key=os.getenv("AWS_SECRET_KEY"),
        region_name=os.getenv("AWS_REGION"),
    )


def _as_extraction(result: Any) -> Dict[str, Any]:
    """Handle both parsed ARDEM responses and legacy response envelopes."""
    if not isinstance(result, dict):
        return {}
    nested = result.get("extraction_result")
    if isinstance(nested, dict):
        return nested
    if isinstance(nested, str):
        try:
            parsed = json.loads(nested)
            return parsed if isinstance(parsed, dict) else {}
        except json.JSONDecodeError:
            return {}
    return result


def _clean_fields(fields: Dict[str, Any]) -> Dict[str, str]:
    return {
        key: str(value).strip()
        for key, value in fields.items()
        if not key.startswith("_") and value is not None and str(value).strip()
    }


def _transaction_md5(value: Any) -> str:
    if not value:
        return ""
    normalized = " ".join(str(value).split()).upper()
    return hashlib.md5(normalized.encode("utf-8")).hexdigest()


def _infer_one_page(image_path: str) -> Dict[str, str]:
    """Upload one temporary split-page PDF and return strict ARDEM JSON fields."""
    tag = uuid.uuid4().hex
    stem = Path(image_path).stem
    input_key = f"{ARDEM_INPUT_PREFIX}/{STATE}/{stem}_{tag}.pdf"
    output_key = f"{ARDEM_OUTPUT_PREFIX}/{STATE}/{stem}_{tag}/"

    with tempfile.TemporaryDirectory(prefix="vitu_pa_ardem_") as temp_dir:
        page_pdf = Path(temp_dir) / f"{stem}.pdf"
        _image_to_pdf(image_path, page_pdf)
        _s3_client().upload_file(str(page_pdf), ARDEM_S3_BUCKET, input_key)
        response = infer_from_api(
            s3_input_path=f"s3://{ARDEM_S3_BUCKET}/{input_key}",
            s3_output_path=f"s3://{ARDEM_S3_BUCKET}/{output_key}",
            prompt=PA_MODEL_PROMPT,
            state=STATE,
        )
    return _clean_fields(_as_extraction(response))


def infer_multiple_images(image_paths: List[str]) -> Dict[str, Any]:
    """Implement PA's tested priority/fallback page strategy using ARDEM LLM.

    It skips the cover page when further pages exist, processes page 3 and the
    final page first, and processes page 2 only if critical values are absent.
    This follows ``vitu_test_PA.py`` while retaining the existing pipeline
    function contract.
    """
    pages = sorted((str(path) for path in image_paths), key=_page_number)
    if not pages:
        return {"_ardem_error": "No pages available for PA inference."}

    non_cover_pages = [path for path in pages if _page_number(path) > 1] or pages
    by_page = {_page_number(path): path for path in non_cover_pages}
    priority_numbers = [number for number in (3, _page_number(non_cover_pages[-1])) if number in by_page]
    # A single page may be both page 3 and the final page.
    priority_numbers = list(dict.fromkeys(priority_numbers))
    if not priority_numbers:
        priority_numbers = [_page_number(non_cover_pages[-1])]

    page_results: Dict[int, Dict[str, str]] = {}
    errors: List[str] = []
    for number in priority_numbers:
        try:
            print(f"[ARDEM] PA inference for page {number}")
            page_results[number] = _infer_one_page(by_page[number])
        except Exception as exc:
            errors.append(f"page {number}: {exc}")
            page_results[number] = {}

    found = {
        field
        for result in page_results.values()
        for field in CRITICAL_FIELDS
        if result.get(field, "").strip()
    }
    if CRITICAL_FIELDS - found and 2 in by_page and 2 not in page_results:
        try:
            print("[ARDEM] PA critical fields missing; running page 2 fallback")
            page_results[2] = _infer_one_page(by_page[2])
        except Exception as exc:
            errors.append(f"page 2: {exc}")
            page_results[2] = {}

    if not any(page_results.values()):
        return {"_ardem_error": "; ".join(errors) or "ARDEM returned no PA fields."}

    # Preserve the PA merge policy from vitu_test_PA.py: price comes from the
    # first available processed page; remaining values come from the last page.
    final: Dict[str, Any] = {}
    price_key = "Taxable Selling Price / Purchase Price"
    for page_number in sorted(page_results):
        value = page_results[page_number].get(price_key, "").strip()
        if value:
            final[price_key] = value
            break

    last_result = page_results[max(page_results)]
    for field, value in last_result.items():
        if field not in {price_key, "Total State Taxes"}:
            final[field] = value

    # Account for prompt variants while retaining the legacy PA normalizer's
    # canonical field key.
    if price_key not in final and "Taxable Selling Price/ Purchase Price" in last_result:
        final[price_key] = last_result["Taxable Selling Price/ Purchase Price"]

    final["_txn_md5"] = _transaction_md5(final.get("Transaction ID"))
    if errors:
        final["_ardem_warnings"] = "; ".join(errors)
    return final


def save_json_and_excel(data: Any, json_out: str, excel_out: str) -> None:
    """Export batch results using the required Backend_db PA template."""
    from openpyxl import Workbook, load_workbook

    Path(json_out).parent.mkdir(parents=True, exist_ok=True)
    with open(json_out, "w", encoding="utf-8") as handle:
        json.dump(data, handle, indent=2, ensure_ascii=False)

    records = data if isinstance(data, list) else [data]
    records = [record for record in records if isinstance(record, dict)]
    output_path = Path(excel_out)
    template = Path(__file__).parent / "Backend_db" / "static" / "templates" / "ARDEM_VITU_TEMPLATE.xlsx"
    if not template.exists():
        raise FileNotFoundError(f"Required PA Excel template is missing: {template}")
    if not output_path.exists():
        shutil.copy2(template, output_path)

    workbook = load_workbook(output_path) if output_path.exists() else Workbook()
    sheet = workbook.active
    header_row = 3 if sheet.max_row >= 3 and any(cell.value for cell in sheet[3]) else 1
    headers = [cell.value for cell in sheet[header_row]]
    if not any(headers):
        headers = ["Filename"]
        sheet.cell(row=header_row, column=1, value="Filename")

    template_to_field = {
        "Taxable Selling Price / Purchase Price": "Taxable Selling Price / Purchase Price",
        "Sales Tax Credit Reciprocity ": "Sales Tax Credit Reciprocity",
        "Local Use / County Fee": "Local Use / County Fee",
        "Additional Charge ": "Additional Charge",
    }
    for record in records:
        row = []
        for header in headers:
            if not header:
                row.append("")
                continue
            source = template_to_field.get(header, header)
            row.append(record.get(source, record.get(header, "")))
        sheet.append(row)
    workbook.save(output_path)
