"""Post-execution diagnostic enrichment.

After a script runs, this module:
1. AST-parses the script to find selector strings
2. Re-runs each selector against the actual page to get match counts + candidates
3. Classifies the failure category
4. Generates Reflexion-style lesson bullets
"""

from __future__ import annotations

import ast
import re
from typing import Any

from agentic.models import (
    AttemptDiagnostics,
    ExecutionResult,
    FieldComparison,
    SelectorProbe,
)


def extract_selectors_from_script(script: str) -> list[dict]:
    """AST-parse a script to find all selector strings passed to .find() or .find_all().

    Returns list of {selector: str, line_number: int}.
    """
    selectors = []
    try:
        tree = ast.parse(script)
    except SyntaxError:
        return selectors

    for node in ast.walk(tree):
        if not isinstance(node, ast.Call):
            continue
        # Look for .find('...') or .find_all('...')
        func = node.func
        if not isinstance(func, ast.Attribute):
            continue
        if func.attr not in ("find", "find_all"):
            continue
        # Get the first positional argument (the selector string)
        if node.args and isinstance(node.args[0], ast.Constant) and isinstance(node.args[0].value, str):
            selectors.append(
                {
                    "selector": node.args[0].value,
                    "line_number": node.lineno,
                }
            )

    return selectors


def probe_selectors(selectors: list[dict], page: Any) -> list[SelectorProbe]:
    """Re-run each selector against the page and record match counts + candidates."""
    probes = []
    for sel_info in selectors:
        selector = sel_info["selector"]
        try:
            matches = page.find_all(selector)
            candidates = []
            for el in list(matches)[:5]:  # cap at 5
                candidate = {
                    "text": (el.text or "")[:80] if hasattr(el, "text") else "",
                    "bbox": el.bbox if hasattr(el, "bbox") else None,
                }
                if hasattr(el, "size"):
                    candidate["size"] = el.size
                if hasattr(el, "bold"):
                    candidate["bold"] = el.bold
                if hasattr(el, "fontname"):
                    candidate["fontname"] = el.fontname
                candidates.append(candidate)

            remaining = len(matches) - 5
            probes.append(
                SelectorProbe(
                    selector=selector,
                    match_count=len(matches),
                    candidates=candidates,
                    line_number=sel_info["line_number"],
                )
            )
        except Exception as e:
            # Selector is invalid — record the error, not "0 matches"
            probes.append(
                SelectorProbe(
                    selector=selector,
                    match_count=-1,  # -1 signals error, not "no matches"
                    candidates=[{"error": f"{type(e).__name__}: {e}"}],
                    line_number=sel_info["line_number"],
                )
            )

    return probes


def classify_failure(
    execution: ExecutionResult,
    comparisons: list[FieldComparison],
    probes: list[SelectorProbe],
) -> str | None:
    """Categorize the primary failure mode."""
    if not execution.success:
        if execution.error_type == "timeout":
            return "TIMEOUT"
        if execution.error_type == "lint_block":
            return "LINT_BLOCK"
        return "CRASH"

    wrong_fields = [c for c in comparisons if not c.match]
    if not wrong_fields:
        return None  # all correct

    no_match = [p for p in probes if p.match_count == 0]
    multi_match = [p for p in probes if p.match_count > 1]

    if no_match:
        return "NO_MATCH"
    if multi_match:
        return "MULTI_MATCH"
    return "WRONG_VALUE"


def detect_approach(script: str) -> str:
    """Heuristic detection of which extraction approach the script uses."""
    indicators = {
        "table": ["extract_table", "to_df", "guides()", "from_whitespace", "from_content"],
        "layout": ["analyze_layout", "region[type="],
        "regex": ["re.search", "re.findall", "re.match"],
        "form_cells": ["detect_form_cells", "form_cell"],
    }

    found = set()
    script_lower = script.lower()
    for approach, keywords in indicators.items():
        for kw in keywords:
            if kw.lower() in script_lower:
                found.add(approach)
                break

    has_spatial = any(
        op in script for op in [".right(", ".below(", ".above(", ".left("]
    )
    has_find = ".find(" in script

    if has_find and has_spatial:
        found.add("anchor_spatial")

    if len(found) > 1:
        return "mixed"
    if found:
        return found.pop()
    if has_find:
        return "anchor_spatial"
    return "unknown"


def generate_lessons(
    comparisons: list[FieldComparison],
    probes: list[SelectorProbe],
    category: str | None,
) -> list[str]:
    """Auto-generate Reflexion-style lesson bullets from diagnostics."""
    lessons = []

    # Selector-level lessons
    for probe in probes:
        if probe.match_count == 0:
            lessons.append(
                f"Selector '{probe.selector}' (line {probe.line_number}) matched 0 elements "
                f"— selector is wrong or too specific"
            )
        elif probe.match_count > 1:
            texts = [c.get("text", "")[:40] for c in probe.candidates[:3]]
            lessons.append(
                f"Selector '{probe.selector}' (line {probe.line_number}) matched "
                f"{probe.match_count} elements: {texts} — need disambiguation"
            )

    # Field-level lessons
    for comp in comparisons:
        if comp.match_type == "wrong":
            lessons.append(
                f"Field '{comp.field_name}': extracted '{comp.raw_actual}' "
                f"but expected '{comp.raw_expected}' — wrong element selected"
            )
        elif comp.match_type == "missing":
            lessons.append(
                f"Field '{comp.field_name}': not extracted — may need different approach"
            )

    return lessons


def enrich_diagnostics(
    script: str,
    execution: ExecutionResult,
    comparisons: list[FieldComparison],
    page: Any,  # Natural PDF Page object
) -> AttemptDiagnostics:
    """Full diagnostic enrichment: probe selectors, classify failure, generate lessons."""
    # Extract and probe selectors
    selectors = extract_selectors_from_script(script)
    probes = probe_selectors(selectors, page) if page is not None else []

    # Classify failure
    category = classify_failure(execution, comparisons, probes)

    # Generate lessons
    lessons = generate_lessons(comparisons, probes, category)

    # Detect approach
    approach = detect_approach(script)

    return AttemptDiagnostics(
        failure_category=category,
        selector_probes=probes,
        lessons=lessons,
        approach_used=approach,
    )
