"""Code review step: send extraction script to Claude Code for improvement.

Takes a working extraction script and sends it to a Claude Code instance that
has full access to the Natural PDF codebase. The reviewer can read source code,
run the script, poke at intermediate results, and suggest improvements.

Output: an improved script + explanation of changes. These pairs (naive, improved)
are training data for fine-tuning.
"""

from __future__ import annotations

import asyncio
import json
import os
import signal
import tempfile
import time
from dataclasses import asdict, dataclass, field
from pathlib import Path
from typing import Any


@dataclass
class ReviewResult:
    """Result of a code review pass."""

    original_script: str
    improved_script: str | None = None
    explanation: str = ""
    changes_made: list[str] = field(default_factory=list)
    duration_ms: int = 0
    success: bool = False
    error: str | None = None
    raw_output: str = ""

    def to_dict(self) -> dict:
        return asdict(self)


def build_review_prompt(
    script: str,
    pdf_path: str,
    page_number: int,
    schema: list[dict],
    gold: dict[str, Any],
    to_llm_output: str,
) -> str:
    """Build the prompt for the Claude Code review agent."""
    schema_desc = "\n".join(
        f"  {f['name']}: {f.get('description', '')} ({f.get('field_type', 'string')})"
        for f in schema
    )

    return f"""\
You are reviewing and improving a PDF extraction script that uses the Natural PDF library.

## Context
- PDF: {pdf_path} (page {page_number})
- The script below correctly extracts all required fields. Your job is to make it better.

## Target Fields
{schema_desc}

## What the page looks like (.to_llm() output)
{to_llm_output}

## The Script to Review
```python
{script}
```

## Your Task
1. Read the Natural PDF library source code to understand what APIs are available.
   Key locations:
   - natural_pdf/core/page.py — Page methods
   - natural_pdf/elements/ — Element, Region, ElementCollection
   - natural_pdf/describe/to_llm.py — how .to_llm() works
   - docs/tutorials/ — usage patterns
   - docs/cookbooks/ — real-world examples

2. Run the script to verify it works:
   ```
   .venv/bin/python -c "
   from natural_pdf import PDF
   import json
   pdf = PDF('{pdf_path}')
   page = pdf.pages[{page_number}]

   # paste extract function here, then:
   result = extract(page)
   print(json.dumps(result, indent=2, default=str))
   pdf.close()
   "
   ```

3. Improve the script. Look for:
   - Simpler Natural PDF idioms (e.g., chaining, vectorized pandas instead of loops)
   - More robust selectors (will this work on similar documents with slight variations?)
   - Unnecessary complexity that can be removed
   - Better use of the library's features (guides, form cells, regions, etc.)
   - Cleaner code structure

4. Write the improved script to: {Path(pdf_path).parent.parent / 'agentic' / 'output' / 'reviewed_script.py'}

5. After writing the file, print a summary of what you changed and why.

## Rules
- The improved script MUST still extract all fields correctly.
- Keep the same function signature: def extract(page) -> dict
- Run the improved version to verify it still works before finishing.
- Don't add try/except or defensive guards — keep the "let it crash" philosophy.
"""


async def run_review(
    script: str,
    pdf_path: str,
    page_number: int = 0,
    schema: list[dict] | None = None,
    gold: dict[str, Any] | None = None,
    to_llm_output: str = "",
    model: str = "sonnet",
    max_turns: int = 15,
    timeout_seconds: int = 600,
    cwd: str | None = None,
) -> ReviewResult:
    """Send a script to Claude Code for review and improvement.

    Args:
        script: The extraction script to review.
        pdf_path: Path to the PDF file.
        page_number: Page number (0-indexed).
        schema: List of field dicts [{name, description, field_type}].
        gold: Gold standard values dict.
        to_llm_output: The .to_llm() representation of the page.
        model: Claude model to use.
        max_turns: Max conversation turns for the review agent.
        timeout_seconds: Max time for the review.
        cwd: Working directory for Claude Code.

    Returns:
        ReviewResult with original and improved scripts.
    """
    prompt = build_review_prompt(
        script=script,
        pdf_path=pdf_path,
        page_number=page_number,
        schema=schema or [],
        gold=gold or {},
        to_llm_output=to_llm_output,
    )

    # Output path for the reviewed script
    output_path = Path(pdf_path).parent.parent / "agentic" / "output" / "reviewed_script.py"

    cmd = [
        "claude", "-p", "-",
        "--max-turns", str(max_turns),
        "--model", model,
        "--allowedTools", "Read,Write,Edit,Bash,Glob,Grep",
    ]

    start = time.monotonic()

    try:
        proc = await asyncio.create_subprocess_exec(
            *cmd,
            stdin=asyncio.subprocess.PIPE,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
            cwd=cwd or str(Path(pdf_path).parent.parent),
            start_new_session=True,
        )

        try:
            stdout, stderr = await asyncio.wait_for(
                proc.communicate(input=prompt.encode()),
                timeout=timeout_seconds,
            )
        except (asyncio.TimeoutError, TimeoutError):
            try:
                os.killpg(proc.pid, signal.SIGKILL)
            except OSError:
                proc.kill()
            await proc.wait()
            return ReviewResult(
                original_script=script,
                duration_ms=int((time.monotonic() - start) * 1000),
                error=f"Review timed out after {timeout_seconds}s",
            )

        duration_ms = int((time.monotonic() - start) * 1000)
        raw_output = stdout.decode("utf-8", errors="replace")

        # Check if the reviewed script was written
        improved_script = None
        if output_path.exists():
            improved_script = output_path.read_text()

        return ReviewResult(
            original_script=script,
            improved_script=improved_script,
            explanation=raw_output[-3000:] if raw_output else "",
            duration_ms=duration_ms,
            success=improved_script is not None,
            raw_output=raw_output[-5000:] if raw_output else "",
        )

    except FileNotFoundError:
        return ReviewResult(
            original_script=script,
            duration_ms=int((time.monotonic() - start) * 1000),
            error="'claude' CLI not found. Install Claude Code: https://claude.ai/code",
        )
    except Exception as e:
        return ReviewResult(
            original_script=script,
            duration_ms=int((time.monotonic() - start) * 1000),
            error=str(e),
        )


def review_sync(
    script: str,
    pdf_path: str,
    page_number: int = 0,
    schema: list[dict] | None = None,
    gold: dict[str, Any] | None = None,
    to_llm_output: str = "",
    **kwargs,
) -> ReviewResult:
    """Synchronous wrapper for run_review."""
    return asyncio.run(run_review(
        script=script,
        pdf_path=pdf_path,
        page_number=page_number,
        schema=schema,
        gold=gold,
        to_llm_output=to_llm_output,
        **kwargs,
    ))
