"""
Data Loader for LLM Subject Extraction Demo

Handles loading of LLM evaluation run results from the results/llm_evaluation directory.
"""

import json
import re
from pathlib import Path
from typing import Dict, List, Any, Optional

# ---------------------------------------------------------------------------
# Colour palette for topic labels
# ---------------------------------------------------------------------------
TOPIC_COLORS = [
    "#4C72B0", "#DD8452", "#55A868", "#C44E52", "#8172B2",
    "#937860", "#DA8BC3", "#8C8C8C", "#CCB974", "#64B5CD",
    "#3B7ABF", "#E07B39", "#3D8F5C", "#A63030", "#6B5FA6",
    "#7A6550", "#C66FAD", "#6E6E6E", "#B39C5A", "#4EA3B8",
]


class LLMEvalDataLoader:
    """Loads and caches LLM evaluation run data from disk."""

    def __init__(self, results_root: str):
        self._run_cache: Dict[str, Dict[str, Any]] = {}
        example_data_path = Path(__file__).parent.parent / "example_data.json"
        if example_data_path.exists():
            try:
                self._run_cache = json.loads(example_data_path.read_text(encoding="utf-8"))
            except Exception as e:
                print(f"Error loading example data: {e}")
        self._topic_color_map: Dict[str, str] = {}
        self._color_idx = 0

    # ------------------------------------------------------------------
    # Run discovery
    # ------------------------------------------------------------------

    def get_available_runs(self) -> List[Dict[str, str]]:
        """Return metadata for every evaluation run found in the results root."""
        runs: List[Dict[str, str]] = []
        for run_id, run in self._run_cache.items():
            config = run.get("config", {})
            meta = {
                "run_id": run_id,
                "path": "",
                "backend": config.get("pipeline_config", {}).get("backend", "unknown"),
                "model_name": config.get("pipeline_config", {}).get("model_name", "unknown"),
                "timestamp": config.get("timestamp", ""),
                "split": config.get("split", ""),
            }
            runs.append(meta)
        return runs

    # ------------------------------------------------------------------
    # Per-run data
    # ------------------------------------------------------------------

    def load_run(self, run_id: str) -> Optional[Dict[str, Any]]:
        """Load the full data for a single evaluation run."""
        return self._run_cache.get(run_id)

    # ------------------------------------------------------------------
    # Helper accessors
    # ------------------------------------------------------------------

    def get_document(self, run_id: str, doc_id: str) -> Optional[Dict[str, Any]]:
        run = self.load_run(run_id)
        if run is None:
            return None
        return run.get("documents", {}).get(doc_id)

    def get_aggregate(self, run_id: str) -> Optional[Dict[str, Any]]:
        run = self.load_run(run_id)
        if run is None:
            return None
        return run.get("aggregate")

    def get_config(self, run_id: str) -> Optional[Dict[str, Any]]:
        run = self.load_run(run_id)
        if run is None:
            return None
        return run.get("config")

    # ------------------------------------------------------------------
    # Topic colour management
    # ------------------------------------------------------------------

    def get_topic_color(self, topic: str) -> str:
        if topic not in self._topic_color_map:
            self._topic_color_map[topic] = TOPIC_COLORS[
                self._color_idx % len(TOPIC_COLORS)
            ]
            self._color_idx += 1
        return self._topic_color_map[topic]

    def get_all_topics_in_run(self, run_id: str) -> List[str]:
        """Return the sorted list of all unique GT topics seen in a run."""
        run = self.load_run(run_id)
        if not run:
            return []
        topics: set = set()
        for doc_data in run.get("documents", {}).values():
            for alignment in doc_data.get("aligned_comparison", {}).get("subjects_alignment", []):
                gt = alignment.get("matched_gt", {})
                topics.update(gt.get("topics", []))
        return sorted(topics)

    # ------------------------------------------------------------------
    # Municipality extraction helper
    # ------------------------------------------------------------------

    @staticmethod
    def parse_municipality(doc_id: str) -> str:
        return doc_id.split("_")[0] if "_" in doc_id else doc_id

    @staticmethod
    def parse_date(doc_id: str) -> str:
        m = re.search(r"(\d{4}-\d{2}-\d{2})", doc_id)
        return m.group(1) if m else ""