"""Prompt construction for the agent LLM.

System prompt is stable across iterations (cacheable).
User message changes each iteration with .to_llm() output, gold, and history.
"""

from __future__ import annotations

from agentic.api_reference import API_REFERENCE
from agentic.config import AgentConfig
from agentic.models import ExtractionAttempt, FieldSchema, GoldStandard


def build_system_prompt(schema: list[FieldSchema], config: AgentConfig) -> str:
    """Build the stable system prompt with API reference, rules, and schema."""
    schema_lines = []
    for f in schema:
        parts = [f.name]
        if f.description:
            parts.append(f"— {f.description}")
        if f.field_type != "string":
            parts.append(f"({f.field_type})")
        schema_lines.append("  " + " ".join(parts))
    schema_block = "\n".join(schema_lines)

    return f"""\
You are writing Python extraction code using the Natural PDF library.

## Your Task
Write a function `def extract(pdf) -> dict` that extracts specific fields from a PDF.
The function receives a Natural PDF PDF object (with pdf.pages[N] for each page) and
returns a dictionary with these keys:

{schema_block}

## Natural PDF API Reference
{API_REFERENCE}

## Rules
1. No try/except — if something goes wrong, let it crash. The error message tells us
   what to fix. Don't mask errors with exception handling.
2. Only labels in selectors — NEVER put data values (names, dates, amounts, cities) in
   :contains(). Only use template text (labels, headers) that would be the same on every
   copy of this document type.
3. No absolute coordinates — use spatial navigation (.right(), .below(until=...)), not
   pixel positions. No .x0, .top, .x1, .bottom access. No hardcoded verticals=[...].
4. Find the RIGHT selector — if a selector doesn't work, try a fundamentally different
   approach rather than chaining alternatives.
5. Use named variables for selectors — e.g. TOTAL_LABEL = page.find('text:contains("Total")')
   This makes the code readable and debuggable.
6. Comment extensively — explain WHY you chose each approach, what you considered, what
   you see in the page description that led to your decision. More comments = better
   debugging when something goes wrong. Show your reasoning.
7. Keep extraction code simple and direct. If something can be done in one line, do it
   in one line. Avoid building complex intermediate data structures when a simple chain
   of Natural PDF calls would work.
8. Pandas rule: when you have a DataFrame from .extract_table().to_df(), use vectorized
   column operations and df.to_dict(orient="records"). Don't iterate rows with iterrows().
9. Table fields: when a schema field has type "table", return a list of row dicts.
   Example: result["enrollments"] = df.to_dict(orient="records")
   Use clean snake_case column names. The scorer compares tables cell-by-cell with
   column normalization and row reordering tolerance.
10. Return all schema keys. If a field crashes during extraction, that's fine — it tells
   us the approach is wrong and we'll fix it next iteration.

## How Spatial Navigation Works
In the .to_llm() layout preview, ┃ pipes show text element boundaries:
  Name: ┃ Jonathan ┃ David ┃ Smith Malago

**Default: use until='text' to grab one cell.**
  find("Name").right(until='text').extract_text()  → "Jonathan"
  find("Name").below(until='text').extract_text()  → same idea, downward
Think of until='text' as "give me the next text cell only."

**Use right() without until when you want everything:**
  find("Name").right().extract_text()  → "Jonathan David Smith Malago"
  find("Name").right(include_source=True).extract_text()  → "Name: Jonathan David Smith Malago"

.above()/.below() capture full page width by default. .left()/.right() capture element
height only. Switch with width='element'/'full' or height='element'/'full'.

Prefer bounded navigation (until=...) over collecting lists and indexing. Minimize code
surface area — the shortest correct Natural PDF idiom is usually the best one.

## Output Format
Return ONLY the Python code for the extract function. No markdown fences, no imports
(the wrapper handles imports). Just the function definition starting with `def extract(pdf):`.

## Approach
Explore the page structure described in the .to_llm() output. Try things out — use the
API features that seem right for what you see. If a field doesn't extract correctly,
try a different approach entirely in the next iteration. Even if all fields are correct,
look for ways to simplify or make the code more robust.
"""


def build_user_message(
    to_llm_output: str,
    schema: list[FieldSchema],
    gold: GoldStandard,
    attempts: list[ExtractionAttempt],
    iteration: int,
) -> str:
    """Build the per-iteration user message."""
    parts = []

    # ── Page description ─────────────────────────────────────────────────
    parts.append("## Page Description")
    parts.append(to_llm_output)
    parts.append("")

    # ── Expected values ──────────────────────────────────────────────────
    parts.append("## Expected Values (Gold Standard)")
    for field_def in schema:
        value = gold.fields.get(field_def.name, "???")
        parts.append(f"  {field_def.name}: {value}")
    if gold.warnings:
        parts.append("")
        parts.append("Warnings:")
        for w in gold.warnings:
            parts.append(f"  ⚠ {w}")
    parts.append("")

    # ── Previous attempts ────────────────────────────────────────────────
    if attempts:
        parts.append("## Previous Attempts")
        for attempt in attempts:
            parts.append(f"### Attempt {attempt.iteration + 1}: Score {attempt.overall_score:.0%}")

            # Per-field results with causal linking
            correct = [c for c in attempt.field_comparisons if c.match]
            wrong = [c for c in attempt.field_comparisons if not c.match and c.match_type == "wrong"]
            missing = [c for c in attempt.field_comparisons if c.match_type == "missing"]

            if correct:
                parts.append(f"  Correct: {', '.join(c.field_name for c in correct)}")
            for c in wrong:
                parts.append(f"  Wrong: {c.field_name} — extracted '{c.raw_actual}', expected '{c.raw_expected}'")
                if c.notes:
                    parts.append(f"    {c.notes}")
            for c in missing:
                parts.append(f"  Missing: {c.field_name}")

            # Diagnostics
            if attempt.diagnostics.failure_category:
                parts.append(f"  Failure category: {attempt.diagnostics.failure_category}")

            # Selector probes
            for probe in attempt.diagnostics.selector_probes:
                if probe.match_count == 0:
                    parts.append(
                        f"  Selector '{probe.selector}' (line {probe.line_number}): matched 0 elements"
                    )
                elif probe.match_count > 1:
                    texts = [c.get("text", "")[:40] for c in probe.candidates[:3]]
                    parts.append(
                        f"  Selector '{probe.selector}' (line {probe.line_number}): "
                        f"matched {probe.match_count} elements: {texts}"
                    )

            # Lint violations
            errors = [v for v in attempt.lint_violations if v.severity == "error"]
            if errors:
                parts.append("  Lint errors:")
                for v in errors:
                    parts.append(f"    line {v.line}: [{v.rule}] {v.message}")

            parts.append("")

        # Lessons learned
        all_lessons = []
        for attempt in attempts:
            all_lessons.extend(attempt.diagnostics.lessons)
        if all_lessons:
            # Deduplicate while preserving order
            seen = set()
            unique_lessons = []
            for lesson in all_lessons:
                if lesson not in seen:
                    seen.add(lesson)
                    unique_lessons.append(lesson)
            parts.append("## Lessons Learned")
            for lesson in unique_lessons:
                parts.append(f"- {lesson}")
            parts.append("")

        # Best-so-far script (for iteration > 0)
        best_attempt = max(attempts, key=lambda a: a.overall_score)
        if best_attempt.overall_score > 0:
            parts.append("## Best Script So Far")
            parts.append(f"Score: {best_attempt.overall_score:.0%}")
            parts.append("```python")
            parts.append(best_attempt.script)
            parts.append("```")
            parts.append("")

    # ── Instruction ──────────────────────────────────────────────────────
    all_correct = attempts and all(c.match for a in attempts[-1:] for c in a.field_comparisons)
    if iteration == 0:
        parts.append("Write the extract function.")
    elif all_correct:
        parts.append("All fields matched! Write a refined version — simpler, more robust, or using a cleaner approach. Explore whether a different technique would work better.")
    else:
        parts.append("Write a complete replacement extract function that fixes the issues above.")

    return "\n".join(parts)
