from __future__ import annotations

import re
from dataclasses import dataclass
from typing import Any

import yaml

from .document_models import DocumentSpec

YAML_TAGGED_CODE_BLOCK_RE = re.compile(r"```(?:yaml|yml)[ \t]*\n?(.*?)```", re.IGNORECASE | re.DOTALL)
YAML_UNMARKED_CODE_BLOCK_RE = re.compile(r"```[ \t]*\n?(.*?)```", re.DOTALL)
YAML_FENCE_LINE_RE = re.compile(r"(?m)^[ \t]*```(?:yaml|yml)?[ \t]*$|^[ \t]*```[ \t]*$")
YAML_ROOT_RE = re.compile(r"(?m)^[ \t]*(?:schema_version|document|blocks):")


@dataclass(frozen=True)
class _YamlCandidate:
    text: str
    order: int
    source: str


@dataclass(frozen=True)
class _ParsedYamlCandidate:
    text: str
    payload: Any
    order: int
    source: str


class DocxDslSyntaxError(ValueError):
    """Raised when an LLM response contains YAML-like content that cannot be parsed."""


class DocxDslContentError(ValueError):
    """Raised when an LLM response parses but is not a usable document specification."""


def _collect_yaml_candidates(raw_text: str) -> list[_YamlCandidate]:
    if not raw_text or not raw_text.strip():
        raise DocxDslContentError("DOCX specification response is empty")

    candidates: list[_YamlCandidate] = []
    seen: set[str] = set()

    def add_candidate(text: str, source: str) -> None:
        normalized = text.strip()
        if not normalized or normalized in seen:
            return
        seen.add(normalized)
        candidates.append(_YamlCandidate(text=normalized, order=len(candidates), source=source))

    for block in YAML_TAGGED_CODE_BLOCK_RE.findall(raw_text):
        add_candidate(block, "tagged_fence")
    for block in YAML_UNMARKED_CODE_BLOCK_RE.findall(raw_text):
        add_candidate(block, "unmarked_fence")

    sanitized = YAML_FENCE_LINE_RE.sub("", raw_text).strip()
    add_candidate(sanitized, "sanitized")

    root_match = YAML_ROOT_RE.search(sanitized)
    if root_match:
        add_candidate(sanitized[root_match.start() :], "anchored_root")

    return candidates


def _select_yaml_candidate(
    raw_text: str,
    *,
    expected_root_keys: tuple[str, ...] = (),
) -> _ParsedYamlCandidate:
    syntax_errors: list[yaml.YAMLError] = []
    parsed_candidates: list[_ParsedYamlCandidate] = []

    for candidate in _collect_yaml_candidates(raw_text):
        try:
            payload = yaml.safe_load(candidate.text)
        except yaml.YAMLError as exc:
            syntax_errors.append(exc)
            continue
        parsed_candidates.append(
            _ParsedYamlCandidate(
                text=candidate.text,
                payload=payload,
                order=candidate.order,
                source=candidate.source,
            )
        )

    mapping_candidates = [candidate for candidate in parsed_candidates if isinstance(candidate.payload, dict)]
    if mapping_candidates:
        preferred = [
            candidate
            for candidate in mapping_candidates
            if all(root_key in candidate.payload for root_key in expected_root_keys)
        ]
        return min(preferred or mapping_candidates, key=lambda candidate: candidate.order)

    if syntax_errors:
        raise DocxDslSyntaxError(f"DOCX specification YAML parsing failed: {syntax_errors[0]}") from syntax_errors[0]
    raise DocxDslContentError("DOCX specification response did not contain exploitable YAML")


def parse_docx_spec(raw_text: str) -> DocumentSpec:
    candidate = _select_yaml_candidate(raw_text, expected_root_keys=("document", "blocks"))
    payload = candidate.payload
    if "document" not in payload:
        raise DocxDslContentError("DOCX specification root must contain document")
    if "blocks" not in payload:
        raise DocxDslContentError("DOCX specification root must contain blocks")
    return DocumentSpec.model_validate(payload)


def format_document_spec_for_prompt(spec: DocumentSpec) -> str:
    return yaml.safe_dump(
        spec.model_dump(mode="python", exclude_none=True),
        allow_unicode=True,
        sort_keys=False,
        default_flow_style=False,
    ).strip()
