"""Gold standard generation via VLM/LLM.

Supports:
- Manual gold (from_dict) for testing
- LLM-based gold generation from PDF text (via Pydantic AI structured output)
- Preflight check to flag possible hallucinations
- Caching to disk (agentic/output/gold/) for reuse
"""

from __future__ import annotations

import json
import re
from pathlib import Path
from typing import Any, Literal

from pydantic import BaseModel, Field

from agentic.config import AgentConfig
from agentic.models import FieldSchema, GoldStandard
from agentic.scorer import normalize_text, parse_number


# ── Pydantic models for structured gold output ─────────────────────────────

class ExtractedField(BaseModel):
    """A single extracted field from the document."""
    name: str = Field(description="snake_case field name")
    value: Any = Field(description="Extracted value — use exact text from the document. For tables, use a list of row dicts.")
    field_type: Literal["string", "number", "currency", "date", "boolean", "address", "list", "table"] = Field(
        default="string", description="Type of the field value. Use 'table' for tabular data (value should be list of row objects)."
    )


class GoldExtractionResult(BaseModel):
    """Structured extraction of all meaningful fields from a document."""
    fields: list[ExtractedField] = Field(
        description="All meaningful fields with actual values (not blank/empty)"
    )


def generate_gold_manual(fields: dict[str, Any]) -> GoldStandard:
    """Create gold from a manually provided dict. For testing and bootstrapping."""
    return GoldStandard.from_dict(fields)


def _extract_pdf_text(pdf_path: str, budget: int = 15000) -> str:
    """Extract text from all pages with per-page budget."""
    from natural_pdf import PDF

    pdf = PDF(pdf_path)
    per_page = max(budget // max(len(pdf.pages), 1), 500)
    page_chunks = []
    for i, page in enumerate(pdf.pages):
        text = page.extract_text(layout=True)
        if len(text) > per_page:
            text = text[:per_page] + "\n... (page truncated)"
        page_chunks.append(f"--- PAGE {i + 1} ---\n{text}")
    pdf.close()
    return "\n\n".join(page_chunks)


class FormFilledCheck(BaseModel):
    """Quick check of whether a document has actual filled-in data."""
    is_filled: bool = Field(description="True if the document contains actual filled-in data values, not just blank form labels/templates")
    reason: str = Field(description="One sentence explaining why")


def preflight_check_filled(pdf_path: str) -> tuple[bool, str]:
    """Quick Haiku check: is this PDF a filled form or an empty template?

    Returns (is_filled, reason). Use before expensive gold generation.
    """
    from dotenv import load_dotenv
    from pydantic_ai import Agent
    from pydantic_ai.settings import ModelSettings

    load_dotenv(Path(__file__).parent / ".env")

    text = _extract_pdf_text(pdf_path, budget=3000)
    agent = Agent(
        "anthropic:claude-haiku-4-5-20251001",
        output_type=FormFilledCheck,
        instructions=(
            "Look at this document text. Determine if it contains actual filled-in data "
            "(names, dates, values, responses) or if it's just an empty form template "
            "with only labels, instructions, and blank fields."
        ),
        retries=3,
    )
    result = agent.run_sync(
        f"Document text:\n{text}",
        model_settings=ModelSettings(temperature=0),
    )
    return result.output.is_filled, result.output.reason


def generate_gold_via_llm(
    pdf_path: str,
    config: AgentConfig,
    cache_dir: str | None = None,
) -> tuple[list[FieldSchema], GoldStandard]:
    """Generate gold standard by sending PDF text to an LLM.

    Uses Pydantic AI for structured output — no markdown fence stripping needed.

    Returns (schema, gold). Caches results under cache_dir for reuse.

    Args:
        pdf_path: Path to the PDF file.
        config: AgentConfig with gold_model and gold_temperature.
        cache_dir: Directory to cache gold files. Defaults to agentic/output/gold/.
    """
    from dotenv import load_dotenv
    from pydantic_ai import Agent
    from pydantic_ai.settings import ModelSettings

    load_dotenv(Path(__file__).parent / ".env")

    # Check cache first
    if cache_dir is None:
        cache_dir = str(Path(__file__).parent / "output" / "gold")
    cache_path = _gold_cache_path(pdf_path, cache_dir)
    if cache_path.exists():
        return _load_cached_gold(cache_path)

    # Quick Haiku check: is this a filled form or empty template?
    is_filled, reason = preflight_check_filled(pdf_path)
    if not is_filled:
        raise ValueError(f"PDF appears to be an empty/unfilled form: {reason}")

    # Extract text from all pages with per-page budget
    all_text = _extract_pdf_text(pdf_path)

    # Map config model name to pydantic-ai model string
    model_name = config.gold_model
    if not model_name.startswith(("openai:", "anthropic:", "google-gla:")):
        model_name = f"anthropic:{model_name}"

    # Use Pydantic AI for structured output
    agent = Agent(
        model_name,
        output_type=GoldExtractionResult,
        instructions=(
            "Extract all meaningful fields from the document. "
            "Only extract fields that have actual values (not blank/empty). "
            "Use the exact text from the document for values — don't paraphrase. "
            "Ignore page numbers, headers, footers, and boilerplate legal text. "
            "Focus on form fields, key-value pairs, and tabular data."
        ),
        retries=3,
    )

    result = agent.run_sync(
        f"Document text:\n{all_text}",
        model_settings=ModelSettings(temperature=config.gold_temperature),
    )
    extraction = result.output

    # Convert to our internal types
    schema = []
    gold_fields = {}
    for f in extraction.fields:
        schema.append(FieldSchema(name=f.name, field_type=f.field_type))
        gold_fields[f.name] = f.value

    gold = GoldStandard(
        fields=gold_fields,
        model_used=config.gold_model,
        mode="text",
        tokens_used=result.usage().input_tokens + result.usage().output_tokens,
    )

    # Cache for reuse
    _save_gold_cache(cache_path, schema, gold_fields)

    return schema, gold


def _gold_cache_path(pdf_path: str, cache_dir: str) -> Path:
    """Derive a cache file path from the PDF path.

    Uses nested layout: cache_dir/category/doc_id.json
    where category = parent dir name and doc_id = file stem.
    """
    p = Path(pdf_path)
    return Path(cache_dir) / p.parent.name / f"{p.stem}.json"


def _save_gold_cache(cache_path: Path, schema: list[FieldSchema], gold_fields: dict) -> None:
    """Save gold to disk for reuse."""
    cache_path.parent.mkdir(parents=True, exist_ok=True)
    data = {
        "schema": [{"name": f.name, "field_type": f.field_type, "description": f.description} for f in schema],
        "gold": gold_fields,
    }
    cache_path.write_text(json.dumps(data, indent=2))


def _load_cached_gold(cache_path: Path) -> tuple[list[FieldSchema], GoldStandard]:
    """Load gold from a cached file."""
    data = json.loads(cache_path.read_text())
    schema = [FieldSchema(**f) for f in data["schema"]]
    gold = GoldStandard(fields=data["gold"], model_used="cached", mode="cached")
    return schema, gold


def preflight_check_gold(gold: GoldStandard, page_text: str) -> GoldStandard:
    """Flag gold values that don't appear in the page text.

    Uses the same normalization as the scorer to avoid false flags
    from formatting differences ($1,234.56 vs 1234.56).
    """
    text_lower = normalize_text(page_text)

    for field_name, value in gold.fields.items():
        if value is None:
            continue

        value_str = str(value)
        found = False

        # Check normalized text match
        if normalize_text(value_str) in text_lower:
            found = True

        # Check numeric match (strip currency, compare as number)
        if not found:
            num = parse_number(value_str)
            if num is not None:
                page_numbers = re.findall(r"[\d,]+\.?\d*", page_text)
                for pn in page_numbers:
                    pn_val = parse_number(pn)
                    if pn_val is not None and abs(pn_val - num) < 0.01:
                        found = True
                        break

        if not found:
            gold.warnings.append(
                f"Gold value '{value_str}' for field '{field_name}' "
                f"not found in page text — may be VLM hallucination"
            )
            gold.suspect_fields.add(field_name)

    return gold
