from __future__ import annotations

import json
import random
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Iterable

from .utils.reward_models import SourceDocument, SourcePack, TaskConstraints

DEFAULT_THEME = {
    "bg": "#F8FAFC",
    "surface": "#FFFFFF",
    "accent": "#2563EB",
    "primary": "#0F172A",
    "secondary": "#475569",
    "font": "Aptos",
    "title_size": 28,
    "body_size": 16,
    "caption_size": 10,
}

_DATA_PATH = Path(__file__).with_name("data.json")


@dataclass(frozen=True, slots=True)
class TaskScenario:
    task_id: str
    difficulty: str
    prompt_text: str
    source_pack: SourcePack
    task_constraints: TaskConstraints
    theme: dict[str, Any]


class TaskRegistry:
    def __init__(self, scenarios: Iterable[TaskScenario]):
        ordered_scenarios = tuple(scenarios)
        if not ordered_scenarios:
            raise ValueError("TaskRegistry requires at least one scenario")

        scenarios_by_id = {scenario.task_id: scenario for scenario in ordered_scenarios}
        if len(scenarios_by_id) != len(ordered_scenarios):
            raise ValueError("TaskRegistry task ids must be unique")

        self._ordered_scenarios = ordered_scenarios
        self._scenarios_by_id = scenarios_by_id

    def __len__(self) -> int:
        return len(self._ordered_scenarios)

    def all(self) -> tuple[TaskScenario, ...]:
        return self._ordered_scenarios

    def get(self, task_id: str) -> TaskScenario:
        try:
            return self._scenarios_by_id[task_id]
        except KeyError as error:
            raise KeyError(f"Unknown task id '{task_id}'") from error

    def by_difficulty(self, difficulty: str) -> tuple[TaskScenario, ...]:
        normalized = difficulty.strip().lower()
        matches = tuple(
            scenario
            for scenario in self._ordered_scenarios
            if scenario.difficulty.lower() == normalized
        )
        if not matches:
            raise KeyError(f"Unknown difficulty '{difficulty}'")
        return matches

    def sample(
        self,
        rng: random.Random,
        *,
        difficulty: str | None = None,
    ) -> TaskScenario:
        candidates = (
            self.by_difficulty(difficulty)
            if difficulty is not None
            else self._ordered_scenarios
        )
        return rng.choice(candidates)


def _load_raw_scenarios() -> list[dict[str, Any]]:
    payload = json.loads(_DATA_PATH.read_text(encoding="utf-8"))
    scenarios = payload.get("scenarios")
    if not isinstance(scenarios, list) or not scenarios:
        raise ValueError("server/data.json must contain a non-empty 'scenarios' list")
    return scenarios


def _build_source_document(payload: dict[str, Any]) -> SourceDocument:
    pages_payload = payload.get("pages")
    pages = None
    if isinstance(pages_payload, list):
        pages = [str(item).strip() for item in pages_payload if str(item).strip()]
        if not pages:
            pages = None

    return SourceDocument(
        doc_id=str(payload["doc_id"]),
        title=str(payload["title"]),
        path=payload.get("path"),
        mime_type=str(payload.get("mime_type", "text/plain")),
        text=str(payload["text"]).strip() if payload.get("text") is not None else None,
        pages=pages,
        images=payload.get("images"),
        metadata=(
            payload.get("metadata") if isinstance(payload.get("metadata"), dict) else {}
        ),
    )


def _build_task_constraints(payload: dict[str, Any]) -> TaskConstraints:
    return TaskConstraints(
        min_slides=payload.get("min_slides"),
        max_slides=payload.get("max_slides"),
        target_audience=payload.get("target_audience"),
        tone=payload.get("tone"),
        extra_constraints=(
            payload.get("extra_constraints")
            if isinstance(payload.get("extra_constraints"), dict)
            else {}
        ),
    )


def _build_scenario(payload: dict[str, Any]) -> TaskScenario:
    task_id = str(payload["task_id"])
    difficulty = str(payload["difficulty"]).lower()
    source_pack_payload = payload.get("source_pack")
    source_documents = None
    source_pack_brief = None
    if isinstance(source_pack_payload, dict):
        source_documents = source_pack_payload.get("documents")
        brief = source_pack_payload.get("brief")
        if isinstance(brief, str) and brief.strip():
            source_pack_brief = brief.strip()
    else:
        source_documents = payload.get("source_documents")

    if not isinstance(source_documents, list) or not source_documents:
        raise ValueError(
            f"Scenario '{task_id}' must define source_pack.documents or source_documents"
        )

    source_pack = SourcePack(
        task_id=task_id,
        documents=[_build_source_document(item) for item in source_documents],
        brief=source_pack_brief,
        metadata={
            **(
                payload.get("metadata")
                if isinstance(payload.get("metadata"), dict)
                else {}
            ),
            "difficulty": difficulty,
        },
    )

    return TaskScenario(
        task_id=task_id,
        difficulty=difficulty,
        prompt_text=str(payload["prompt_text"]),
        source_pack=source_pack,
        task_constraints=_build_task_constraints(payload.get("task_constraints") or {}),
        theme={**DEFAULT_THEME, **dict(payload.get("theme") or {})},
    )


def _load_default_task_registry() -> TaskRegistry:
    return TaskRegistry(_build_scenario(item) for item in _load_raw_scenarios())


DEFAULT_TASK_REGISTRY = _load_default_task_registry()


__all__ = [
    "DEFAULT_TASK_REGISTRY",
    "DEFAULT_THEME",
    "TaskRegistry",
    "TaskScenario",
]