"""
tasks/registry.py
=================
Deterministic task registry for InformationIntegrityEnv.

Loads pre-built scenario datasets from tasks/{difficulty}/scenarios.json.
NO external API calls — all scenarios are static and reproducible.

Usage
-----
    from tasks.registry import TASKS

    task     = TASKS["easy"]
    scenario = task.get_scenario(0)          # by index
    scenario = task.get_scenario("easy_001") # by id
    all_sc   = task.get_all_scenarios()

    # Random selection (seeded for reproducibility)
    import random
    rng      = random.Random(42)
    scenario = task.sample(rng)
"""

from __future__ import annotations

import json
import os
import random
from typing import Dict, List, Optional


# ---------------------------------------------------------------------------
# Path resolution
# ---------------------------------------------------------------------------

_TASKS_DIR = os.path.dirname(os.path.abspath(__file__))


def _load_scenarios(difficulty: str) -> List[dict]:
    path = os.path.join(_TASKS_DIR, difficulty, "scenarios.json")
    if not os.path.exists(path):
        raise FileNotFoundError(
            f"Scenario dataset not found: {path}\n"
            f"Run data_collector.py to build datasets, or check tasks/{difficulty}/scenarios.json"
        )
    with open(path, "r", encoding="utf-8") as fh:
        data = json.load(fh)
    if not isinstance(data, list) or not data:
        raise ValueError(f"scenarios.json for '{difficulty}' must be a non-empty JSON array")
    return data


# ---------------------------------------------------------------------------
# Task class
# ---------------------------------------------------------------------------

class Task:
    """
    Wraps one difficulty-level dataset.

    Parameters
    ----------
    difficulty : str
        One of 'easy', 'medium', 'hard'.
    """

    def __init__(self, difficulty: str) -> None:
        self.difficulty = difficulty
        self._scenarios: List[dict] = _load_scenarios(difficulty)
        self._index_by_id: Dict[str, dict] = {s["id"]: s for s in self._scenarios}

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def get_scenario(self, scenario_id: str | int) -> dict:
        """
        Return a scenario by index (int) or scenario id (str).
        Returns a fresh copy so the env cannot mutate the dataset.
        """
        if isinstance(scenario_id, int):
            if scenario_id < 0 or scenario_id >= len(self._scenarios):
                raise IndexError(
                    f"Scenario index {scenario_id} out of range "
                    f"(dataset has {len(self._scenarios)} scenarios)"
                )
            return dict(self._scenarios[scenario_id])
        # String id lookup
        if scenario_id not in self._index_by_id:
            raise KeyError(
                f"Scenario id '{scenario_id}' not found in {self.difficulty} dataset. "
                f"Available ids: {list(self._index_by_id.keys())}"
            )
        return dict(self._index_by_id[scenario_id])

    def get_all_scenarios(self) -> List[dict]:
        """Return all scenarios as a list of dicts (copies)."""
        return [dict(s) for s in self._scenarios]

    def sample(self, rng: Optional[random.Random] = None) -> dict:
        """Return a uniformly random scenario (optionally seeded)."""
        rng = rng or random.Random()
        return dict(rng.choice(self._scenarios))

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

    def __repr__(self) -> str:
        return f"Task(difficulty={self.difficulty!r}, n_scenarios={len(self)})"


# ---------------------------------------------------------------------------
# Lazy-loaded registry singleton
# ---------------------------------------------------------------------------

class _TaskRegistry:
    """
    Lazy-loading registry.  Datasets are only read from disk on first access,
    so importing this module does not cause IO at import time.
    """

    _cache: Dict[str, Task] = {}

    _SUPPORTED = ("easy", "medium", "hard")

    def __getitem__(self, difficulty: str) -> Task:
        key = difficulty.lower().strip()
        if key not in self._SUPPORTED:
            raise KeyError(
                f"Unknown difficulty '{difficulty}'. "
                f"Supported values: {list(self._SUPPORTED)}"
            )
        if key not in self._cache:
            self._cache[key] = Task(key)
        return self._cache[key]

    def __contains__(self, difficulty: str) -> bool:
        return difficulty.lower().strip() in self._SUPPORTED

    def keys(self):
        return self._SUPPORTED

    def items(self):
        return [(d, self[d]) for d in self._SUPPORTED]

    def reload(self) -> None:
        """Force reload all datasets from disk (useful for tests)."""
        self._cache.clear()

    def __repr__(self) -> str:
        loaded = list(self._cache.keys())
        return f"TaskRegistry(supported={list(self._SUPPORTED)}, loaded={loaded})"


# Singleton — import and use directly:  from tasks.registry import TASKS
TASKS = _TaskRegistry()