"""Rubric model, markdown parser, scaffold writer, and validation."""

from __future__ import annotations

import logging
import os
import re
from dataclasses import dataclass, field
from typing import Any

logger = logging.getLogger(__name__)


@dataclass
class RubricExample:
    input: str
    output: str
    critique: str


@dataclass
class RubricMeta:
    description: str = ""


@dataclass
class Rubric:
    meta: RubricMeta = field(default_factory=RubricMeta)
    pass_criteria: str = ""
    fail_criteria: str = ""
    good_examples: list[RubricExample] = field(default_factory=list)
    bad_examples: list[RubricExample] = field(default_factory=list)

    def validate(self) -> list[str]:
        """Return a list of warnings (empty = valid)."""
        warnings = []
        if not self.pass_criteria.strip():
            warnings.append("Missing pass criteria (## What makes a good response)")
        if not self.fail_criteria.strip():
            warnings.append("Missing fail criteria (## What makes a bad response)")
        if not self.good_examples:
            warnings.append("No good examples — add at least 2 for reliable judging")
        elif len(self.good_examples) < 2:
            warnings.append("Only 1 good example — add at least 2 for reliable judging")
        if not self.bad_examples:
            warnings.append("No bad examples — add at least 2 for reliable judging")
        elif len(self.bad_examples) < 2:
            warnings.append("Only 1 bad example — add at least 2 for reliable judging")
        return warnings


# --- Parser ---

# Heading patterns — flexible matching
_GOOD_HEADING = re.compile(
    r"^##\s+(?:what\s+makes\s+a\s+)?good(?:\s+(?:response|examples?))?\s*$",
    re.IGNORECASE,
)
_BAD_HEADING = re.compile(
    r"^##\s+(?:what\s+makes\s+a\s+)?bad(?:\s+(?:response|examples?))?\s*$",
    re.IGNORECASE,
)
_EXAMPLES_HEADING = re.compile(r"^##\s+examples?\s*$", re.IGNORECASE)
_GOOD_EXAMPLE = re.compile(r"^###\s+good\s*$", re.IGNORECASE)
_BAD_EXAMPLE = re.compile(r"^###\s+bad\s*$", re.IGNORECASE)


def parse_rubric(path: str) -> Rubric:
    """Parse a rubric markdown file into a Rubric dataclass."""
    with open(path) as f:
        content = f.read()
    return parse_rubric_content(content)


def parse_rubric_content(content: str) -> Rubric:
    """Parse rubric markdown content (string) into a Rubric dataclass."""
    meta = _parse_frontmatter(content)
    body = _strip_frontmatter(content)

    rubric = Rubric(meta=meta)
    rubric.pass_criteria = _extract_section(body, _GOOD_HEADING)
    rubric.fail_criteria = _extract_section(body, _BAD_HEADING)

    # Parse examples section
    examples_text = _extract_section(body, _EXAMPLES_HEADING)
    if examples_text:
        rubric.good_examples, rubric.bad_examples = _parse_examples(examples_text)

    return rubric


def _parse_frontmatter(content: str) -> RubricMeta:
    """Extract YAML frontmatter between --- markers."""
    if not content.startswith("---"):
        return RubricMeta()
    end = content.find("---", 3)
    if end == -1:
        return RubricMeta()
    yaml_text = content[3:end].strip()

    meta = RubricMeta()
    for line in yaml_text.splitlines():
        line = line.strip()
        if ":" not in line:
            continue
        key, _, value = line.partition(":")
        key = key.strip()
        value = value.strip().strip('"').strip("'")
        if key == "description":
            meta.description = value
    return meta


def _strip_frontmatter(content: str) -> str:
    """Return content without YAML frontmatter."""
    if not content.startswith("---"):
        return content
    end = content.find("---", 3)
    if end == -1:
        return content
    return content[end + 3:]


def _extract_section(body: str, heading_pattern: re.Pattern) -> str:
    """Extract text between a matched heading and the next ## heading."""
    lines = body.splitlines()
    start = None
    for i, line in enumerate(lines):
        if heading_pattern.match(line.strip()):
            start = i + 1
            break
    if start is None:
        return ""

    end = len(lines)
    for i in range(start, len(lines)):
        stripped = lines[i].strip()
        if stripped.startswith("## ") and not stripped.startswith("### "):
            end = i
            break

    return "\n".join(lines[start:end]).strip()


def _parse_examples(text: str) -> tuple[list[RubricExample], list[RubricExample]]:
    """Parse ### Good and ### Bad example blocks."""
    good: list[RubricExample] = []
    bad: list[RubricExample] = []

    lines = text.splitlines()
    current_kind: str | None = None
    current_block: list[str] = []

    def _flush():
        if current_kind and current_block:
            example = _parse_single_example("\n".join(current_block))
            if example:
                if current_kind == "good":
                    good.append(example)
                else:
                    bad.append(example)
            else:
                logger.warning("Failed to parse example in %s block (missing Input/Output fields)", current_kind)

    for line in lines:
        stripped = line.strip()
        if _GOOD_EXAMPLE.match(stripped):
            _flush()
            current_kind = "good"
            current_block = []
        elif _BAD_EXAMPLE.match(stripped):
            _flush()
            current_kind = "bad"
            current_block = []
        elif current_kind is not None:
            current_block.append(line)

    _flush()
    return good, bad


def _parse_single_example(text: str) -> RubricExample | None:
    """Parse a single example block: **Input:**, **Output:**, **Critique:**."""
    input_val = _extract_field(text, "Input")
    output_val = _extract_field(text, "Output")
    critique_val = _extract_field(text, "Critique")

    if not input_val or not output_val:
        return None

    return RubricExample(
        input=input_val,
        output=output_val,
        critique=critique_val or "",
    )


def _extract_field(text: str, field_name: str) -> str | None:
    """Extract text after **FieldName:** until the next **field or end."""
    pattern = re.compile(
        rf"\*\*{field_name}:\*\*\s*(.*?)(?=\n\*\*\w|$)",
        re.DOTALL | re.IGNORECASE,
    )
    match = pattern.search(text)
    if not match:
        return None
    return match.group(1).strip()


# --- LLM-drafted rubric ---


_INJECTION_PATTERNS = [
    re.compile(r"(?i)^\s*ignore\s+(previous|all|prior)\s+(instructions?|rules?)", re.IGNORECASE),
    re.compile(r"(?i)^\s*disregard\s+(previous|all|prior)", re.IGNORECASE),
    re.compile(r"(?i)^\s*forget\s+(everything|all|your)", re.IGNORECASE),
    re.compile(r"(?i)^\s*new\s+instructions?:", re.IGNORECASE),
]


def _sanitize_for_prompt(value: str) -> str:
    """Basic sanitization to reduce prompt injection risk."""
    if not value:
        return value
    sanitized = value
    for pattern in _INJECTION_PATTERNS:
        sanitized = pattern.sub(r"[FILTERED]", sanitized)
    return sanitized


def _prepare_traces(traces: list[dict[str, Any]]) -> list[dict[str, Any]]:
    """Trim traces to the fields the draft API needs."""
    prepared = []
    for trace in traces[:10]:
        spans = []
        for span in trace.get("spans", []):
            s: dict[str, Any] = {
                "kind": span.get("kind", "unknown"),
                "span_id": span.get("span_id", ""),
            }
            if span.get("name"):
                s["name"] = span["name"]
            if span.get("model"):
                s["model"] = span["model"]
            if span.get("input") is not None:
                s["input"] = _sanitize_for_prompt(span["input"])
            if span.get("output") is not None:
                s["output"] = _sanitize_for_prompt(span["output"])
            spans.append(s)
        prepared.append({
            "trace_id": trace.get("trace_id", ""),
            "input": _sanitize_for_prompt(trace.get("input", "")),
            "output": _sanitize_for_prompt(trace.get("output", "")),
            "spans": spans,
        })
    return prepared


def draft_rubric(project: str, traces: list[dict[str, Any]], client: Any) -> str:
    """Draft a rubric via the Bandito API draft agent.

    Sends structured trace data to the backend. The backend agent handles
    system prompt extraction, example seeding from evaluations, prompt
    formatting, and LLM calls.
    """
    if not traces:
        raise RuntimeError("No traces available to draft rubric")

    return client.draft_rubric(project, _prepare_traces(traces))
