"""The main extraction session — the feedback loop orchestrator.

ExtractionRunner.run() is the entry point:
  1. Open PDF, get .to_llm()
  2. Generate or accept gold standard
  3. Loop: generate script → lint → execute → score → diagnose → iterate
  4. Return ExtractionSession with full history
"""

from __future__ import annotations

import re
import time
from dataclasses import asdict
from pathlib import Path
from typing import Any

from dotenv import load_dotenv

# Load .env from agentic directory
load_dotenv(Path(__file__).parent / ".env")

from agentic.config import AgentConfig
from agentic.diagnostics import enrich_diagnostics
from agentic.executor import execute_script
from agentic.gold import generate_gold_manual, generate_gold_via_llm, preflight_check_gold
from agentic.linter import lint_script
from agentic.models import (
    ExtractionAttempt,
    ExtractionSession,
    ExecutionResult,
    FieldSchema,
    GoldStandard,
)
from agentic.prompt import build_system_prompt, build_user_message
from agentic.scorer import overall_score, score_extraction


def call_agent_llm(
    system_prompt: str,
    user_message: str,
    config: AgentConfig,
) -> tuple[str, int, int]:
    """Call the agent LLM to generate extraction code.

    Returns (script_text, generation_time_ms, tokens_used).

    Uses the anthropic SDK by default. Override this function for other providers.
    """
    start = time.monotonic()

    try:
        import anthropic

        client = anthropic.Anthropic()
        response = client.messages.create(
            model=config.model,
            max_tokens=config.max_tokens,
            temperature=config.temperature,
            system=system_prompt,
            messages=[{"role": "user", "content": user_message}],
        )
        gen_time = int((time.monotonic() - start) * 1000)

        script = response.content[0].text
        tokens = response.usage.input_tokens + response.usage.output_tokens

        # Strip markdown fences if the model wrapped the code
        script = strip_markdown_fences(script)

        # Trim anything before `def extract` (models sometimes add commentary)
        lines = script.split("\n")
        for i, line in enumerate(lines):
            if line.strip().startswith("def extract"):
                script = "\n".join(lines[i:])
                break

        return script, gen_time, tokens

    except ImportError:
        raise ImportError(
            "anthropic package not installed. Install with: pip install anthropic"
        )


def strip_markdown_fences(text: str) -> str:
    """Extract code from markdown fences if present, otherwise return as-is."""
    # Try to extract content between ```python (or ```) fences
    match = re.search(r'```(?:python)?\s*\n(.*?)```', text, re.DOTALL)
    if match:
        return match.group(1).strip()
    return text.strip()


def syntax_retry(
    script: str,
    system_prompt: str,
    user_message: str,
    config: AgentConfig,
) -> tuple[str, int, int]:
    """Cheap retry loop for syntax errors. Returns (script, total_gen_time, total_tokens).

    Sends only the broken script + error message (~500 tokens) instead of the
    full user message (potentially 16K+). The LLM doesn't need PDF context to
    fix a missing colon.
    """
    import ast

    total_time = 0
    total_tokens = 0

    for retry in range(config.syntax_retries):
        try:
            ast.parse(script)
            return script, total_time, total_tokens  # valid — done
        except SyntaxError as e:
            # Minimal retry — just the broken code + error, no PDF context
            retry_msg = (
                f"You previously generated a Python script that failed with a SyntaxError.\n"
                f"Error on line {e.lineno}: {e.msg}\n\n"
                f"Here is the script you generated:\n```python\n{script}\n```\n\n"
                f"Return ONLY the corrected Python code for the extract function. "
                f"Keep the logic identical. No markdown fences."
            )
            script, gen_time, gen_tokens = call_agent_llm(
                system_prompt, retry_msg, config
            )
            total_time += gen_time
            total_tokens += gen_tokens

    return script, total_time, total_tokens


def should_early_stop(attempts: list[ExtractionAttempt], config: AgentConfig) -> bool:
    """Check if we should stop early due to repeated failure category."""
    if len(attempts) < 2:
        return False

    # Check last N attempts for same failure category with no improvement
    recent = attempts[-config.approach_shift_after :]
    if len(recent) < config.approach_shift_after:
        return False

    categories = [a.diagnostics.failure_category for a in recent]
    scores = [a.overall_score for a in recent]

    # Same category and no score improvement
    if len(set(categories)) == 1 and categories[0] is not None:
        if max(scores) - min(scores) < 0.01:
            return True

    return False


def determine_stop_reason(
    attempts: list[ExtractionAttempt],
    config: AgentConfig,
    converged: bool,
) -> str:
    """Determine why the loop stopped."""
    if converged:
        return "converged"
    if len(attempts) >= config.max_iterations:
        return "max_iterations"
    if should_early_stop(attempts, config):
        return "repeated_failure"
    return "unknown"


class ExtractionRunner:
    """Orchestrates the extraction feedback loop."""

    def __init__(self, config: AgentConfig | None = None):
        self.config = config or AgentConfig()

    def run(
        self,
        pdf_path: str,
        page_number: int = 0,
        schema: list[FieldSchema] | list[dict] | None = None,
        gold: dict | GoldStandard | None = None,
    ) -> ExtractionSession:
        """Run the extraction loop.

        Args:
            pdf_path: Path to PDF file.
            page_number: Page to extract from (0-indexed).
            schema: List of FieldSchema or dicts with {name, description, field_type}.
            gold: Pre-computed gold standard (dict or GoldStandard). If None, generated via LLM.

        Returns:
            ExtractionSession with full history and best script.
        """
        session_start = time.monotonic()

        # Normalize schema
        if schema is None:
            schema = []
        schema = [
            FieldSchema(**f) if isinstance(f, dict) else f for f in schema
        ]

        # Open PDF and get .to_llm()
        from natural_pdf import PDF

        pdf = PDF(pdf_path)

        # Build .to_llm() for all pages
        to_llm_parts = []
        page_texts = []
        for page in pdf.pages:
            to_llm_parts.append(page.to_llm(detail="standard", include_hints="api"))
            page_texts.append(page.extract_text(layout=True))
        to_llm_output = "\n\n".join(to_llm_parts)
        all_page_text = "\n\n".join(page_texts)

        # Gold standard
        if isinstance(gold, dict):
            gold_obj = generate_gold_manual(gold)
        elif isinstance(gold, GoldStandard):
            gold_obj = gold
        else:
            # Auto-generate gold and schema via LLM
            generated_schema, gold_obj = generate_gold_via_llm(pdf_path, self.config)
            if not schema:
                schema = generated_schema

        # Preflight check — use actual page text
        gold_obj = preflight_check_gold(gold_obj, all_page_text)

        # Build system prompt (stable across iterations)
        system_prompt = build_system_prompt(schema, self.config)

        # Iteration loop
        attempts: list[ExtractionAttempt] = []
        converged = False

        for iteration in range(self.config.max_iterations):
            iter_start = time.monotonic()

            # Build user message
            user_message = build_user_message(
                to_llm_output=to_llm_output,
                schema=schema,
                gold=gold_obj,
                attempts=attempts,
                iteration=iteration,
            )

            # Call agent LLM
            script, gen_time, gen_tokens = call_agent_llm(
                system_prompt, user_message, self.config
            )

            # Syntax micro-loop — cheap retries on parse errors
            if self.config.syntax_retries > 0:
                script, retry_time, retry_tokens = syntax_retry(
                    script, system_prompt, user_message, self.config
                )
                gen_time += retry_time
                gen_tokens += retry_tokens

            # Lint
            violations = lint_script(script, schema, self.config.lint_rules)

            # Execute (skip if lint errors)
            lint_errors = [v for v in violations if v.severity == "error"]
            if self.config.lint_block_on_error and lint_errors:
                execution = ExecutionResult(
                    success=False,
                    error_type="lint_block",
                    error_message=f"{len(lint_errors)} lint error(s): "
                    + "; ".join(v.message for v in lint_errors[:3]),
                )
            else:
                execution = execute_script(script, pdf_path, self.config)

            # Score
            comparisons = score_extraction(
                execution.output, gold_obj, schema, self.config
            )
            score = overall_score(comparisons)

            # Diagnostics — probe against first page (most selectors target page 0)
            diagnostics = enrich_diagnostics(script, execution, comparisons, pdf.pages[0])

            # Record attempt
            iter_time = int((time.monotonic() - iter_start) * 1000)
            attempt = ExtractionAttempt(
                iteration=iteration,
                script=script,
                lint_violations=violations,
                execution=execution,
                field_comparisons=comparisons,
                diagnostics=diagnostics,
                overall_score=score,
                generation_time_ms=gen_time,
                generation_tokens=gen_tokens,
                total_time_ms=iter_time,
            )
            attempts.append(attempt)

            # Track convergence but keep iterating for refinement
            if score >= self.config.target_score:
                converged = True

            # Only stop early on repeated failure, not on success
            if not converged and should_early_stop(attempts, self.config):
                break

        # Close PDF
        pdf.close()

        # Find best attempt
        best = max(attempts, key=lambda a: a.overall_score) if attempts else None
        total_time = int((time.monotonic() - session_start) * 1000)

        return ExtractionSession(
            pdf_path=pdf_path,
            page_number=page_number,
            schema=schema,
            to_llm_output=to_llm_output,
            config=asdict(self.config),
            gold=gold_obj,
            attempts=attempts,
            best_script=best.script if best else None,
            best_score=best.overall_score if best else 0.0,
            converged=converged,
            stop_reason=determine_stop_reason(attempts, self.config, converged),
            total_tokens=sum(a.generation_tokens for a in attempts),
            total_time_ms=total_time,
        )
