from __future__ import annotations

import re
from typing import get_args

from pydantic import BaseModel, Field, model_validator
import yaml

from .deck_models import DeckSpec, SlideSpec, SlideType
from .parsing import _select_yaml_candidate

WEAK_PLAN_SLIDE_TYPES = {"section", "quote"}
ALLOWED_PLAN_SLIDE_TYPES = get_args(SlideType)
STANDARD_GENERATION_SLIDE_TYPES = tuple(slide_type for slide_type in ALLOWED_PLAN_SLIDE_TYPES if slide_type != "custom")


class PlannedSlide(BaseModel):
    """Compact internal definition of one planned slide before deck generation."""

    index: int
    type: SlideType
    goal: str
    key_data_points: list[str] = Field(default_factory=list)
    must_be_informative: bool = True


class DeckPlan(BaseModel):
    """Internal plan used to lock deck shape before generating the final YAML deck."""

    target_slide_count: int
    slides: list[PlannedSlide]

    @model_validator(mode="after")
    def validate_plan(self) -> DeckPlan:
        if self.target_slide_count < 1:
            raise ValueError("Deck plan target_slide_count must be positive")
        if len(self.slides) != self.target_slide_count:
            raise ValueError("Deck plan slide count must match target_slide_count")
        expected_indices = list(range(1, self.target_slide_count + 1))
        actual_indices = [slide.index for slide in self.slides]
        if actual_indices != expected_indices:
            raise ValueError("Deck plan slide indices must be sequential and start at 1")
        return self


class SlidePatch(BaseModel):
    """Replacement payload for one slide during targeted repair."""

    index: int
    slide: SlideSpec


class SlidePatchSet(BaseModel):
    """Collection of slide replacements returned by a targeted repair completion."""

    slides: list[SlidePatch]

    @model_validator(mode="after")
    def validate_patches(self) -> SlidePatchSet:
        if not self.slides:
            raise SlidePatchContentError("Targeted repair must return at least one slide patch")
        indices = [patch.index for patch in self.slides]
        if len(indices) != len(set(indices)):
            raise SlidePatchContentError("Targeted repair slide indices must be unique")
        return self


class DeckConstraintFailure(ValueError):
    """Raised when a generated plan or deck violates global constraints outside DeckSpec."""


class SlidePatchContentError(ValueError):
    """Raised when targeted repair payload content is structurally invalid but retryable."""


def parse_deck_plan(raw_text: str) -> DeckPlan:
    """Parse one YAML deck plan completion into a validated DeckPlan."""

    payload = _select_yaml_candidate(raw_text).payload
    if not isinstance(payload, dict):
        raise ValueError("Deck plan root must be a mapping")
    return DeckPlan.model_validate(payload)


def parse_slide_patch_set(raw_text: str) -> SlidePatchSet:
    """Parse one targeted repair completion into slide replacements."""

    payload = _select_yaml_candidate(raw_text).payload
    if not isinstance(payload, dict):
        raise SlidePatchContentError("Slide patch root must be a mapping")
    return SlidePatchSet.model_validate(payload)


def detect_requested_slide_count(user_prompt: str) -> int | None:
    """Extract an exact slide-count request from the user prompt when one is clearly stated."""

    patterns = [
        r"\bexactement\s+(\d+)\s+slides?\b",
        r"\bexactly\s+(\d+)\s+slides?\b",
        r"\bje veux\s+exactement\s+(\d+)\s+slides?\b",
    ]
    normalized = user_prompt.lower()
    for pattern in patterns:
        match = re.search(pattern, normalized)
        if match:
            return int(match.group(1))
    return None


def validate_deck_plan(deck_plan: DeckPlan, user_prompt: str) -> None:
    """Reject invalid plans before deck generation starts."""

    requested_slide_count = detect_requested_slide_count(user_prompt)
    if requested_slide_count is not None and deck_plan.target_slide_count != requested_slide_count:
        raise DeckConstraintFailure(
            f"Deck plan must target exactly {requested_slide_count} slides, got {deck_plan.target_slide_count}"
        )

    weak_slide_count = sum(1 for slide in deck_plan.slides if slide.type in WEAK_PLAN_SLIDE_TYPES)
    max_weak_slide_count = max(2, deck_plan.target_slide_count // 4)
    if weak_slide_count > max_weak_slide_count:
        raise DeckConstraintFailure("Deck plan overuses section or quote slides for the requested deck length")
    if any(slide.type == "custom" for slide in deck_plan.slides):
        raise DeckConstraintFailure("Deck plan uses custom slides, which are not allowed in standard generation")


def validate_deck_against_plan(spec: DeckSpec, deck_plan: DeckPlan, user_prompt: str) -> None:
    """Reject generated decks that violate validated plan-level constraints."""

    requested_slide_count = detect_requested_slide_count(user_prompt)
    expected_slide_count = requested_slide_count or deck_plan.target_slide_count

    if len(spec.slides) != expected_slide_count:
        raise DeckConstraintFailure(f"Deck must contain exactly {expected_slide_count} slides, got {len(spec.slides)}")
    if len(spec.slides) != len(deck_plan.slides):
        raise DeckConstraintFailure("Generated deck slide count must stay aligned with the approved deck plan")
    if any(slide.type == "custom" for slide in spec.slides):
        raise DeckConstraintFailure("Generated deck uses custom slides, which are not allowed in standard generation")


def apply_slide_patches(spec: DeckSpec, patch_set: SlidePatchSet) -> DeckSpec:
    """Apply targeted slide replacements and return a freshly validated deck spec."""

    deck_payload = spec.model_dump(mode="python")
    for patch in patch_set.slides:
        if patch.index < 1 or patch.index > len(deck_payload["slides"]):
            raise DeckConstraintFailure(f"Targeted repair references slide index {patch.index}, which is out of range")
        deck_payload["slides"][patch.index - 1] = patch.slide.model_dump(
            mode="python",
            exclude_none=True,
        )
    return DeckSpec.model_validate(deck_payload)


def format_deck_plan_for_prompt(deck_plan: DeckPlan) -> str:
    """Serialize a validated deck plan into stable YAML for prompts and logs."""

    return yaml.safe_dump(
        deck_plan.model_dump(mode="python"),
        allow_unicode=True,
        sort_keys=False,
        default_flow_style=False,
    ).strip()


def format_targeted_slides_for_prompt(spec: DeckSpec, slide_indices: list[int]) -> str:
    """Serialize only the flagged slides so targeted repair stays compact."""

    payload = {
        "slides": [
            {
                "index": slide_index,
                "slide": spec.slides[slide_index - 1].model_dump(
                    mode="python",
                    exclude_none=True,
                ),
            }
            for slide_index in slide_indices
        ]
    }
    return yaml.safe_dump(
        payload,
        allow_unicode=True,
        sort_keys=False,
        default_flow_style=False,
    ).strip()


def format_deck_spec_for_prompt(spec: DeckSpec) -> str:
    """Serialize a full deck spec to YAML for reduced-context repair prompts."""

    return yaml.safe_dump(
        spec.model_dump(mode="python", exclude_none=True),
        allow_unicode=True,
        sort_keys=False,
        default_flow_style=False,
    ).strip()
