"""Utilities for discovering and loading precomputed patch metadata."""

from __future__ import annotations

import json
import os
from dataclasses import dataclass
from pathlib import Path
from typing import Dict, Iterable, Mapping, Optional


APP_DIR = Path(__file__).resolve().parent
DEFAULT_DATA_CANDIDATES = (
    APP_DIR / "data" / "test-10",
    Path("data/test-10"),
)
ENV_DATA_ROOT = "PATCH_VIEW_DATA_ROOT"


def resolve_data_root(candidates: Iterable[Path] = DEFAULT_DATA_CANDIDATES) -> Path:
    """Locate the dataset root, honoring environment override when present."""
    override = os.getenv(ENV_DATA_ROOT)
    if override:
        candidate = Path(override).expanduser().resolve()
        if not candidate.exists():
            raise FileNotFoundError(f"PATCH_VIEW_DATA_ROOT={candidate} not found")
        return candidate

    for candidate in candidates:
        resolved = candidate.expanduser()
        if (resolved / "precomputed").exists():
            return resolved

    # Fall back to the last candidate even if missing, to keep error messages consistent
    return Path(list(candidates)[-1])


@dataclass(frozen=True)
class PrecomputedRecord:
    """Loaded metadata for a single audio file."""

    file_id: str
    data: Dict[str, object]

    @property
    def audio_path(self) -> str:
        return str(self.data["audio_path"])

    @property
    def reconstructed_audio_path(self) -> str:
        return str(self.data["reconstructed_audio_path"])


class PrecomputedRepository:
    """In-memory index of precomputed entropy/perplexity artifacts."""

    def __init__(self, data_root: Optional[Path] = None):
        self.data_root = data_root or resolve_data_root()
        self.precomputed_dir = self.data_root / "precomputed"
        self._records: Dict[str, PrecomputedRecord] = {}

    @property
    def records(self) -> Mapping[str, PrecomputedRecord]:
        return self._records

    def load(self) -> None:
        """Load all JSON metadata into memory."""
        if not self.precomputed_dir.exists():
            raise FileNotFoundError(
                f"Expected precomputed dir at {self.precomputed_dir}, but it does not exist."
            )

        for json_path in sorted(self.precomputed_dir.glob("*.json")):
            with open(json_path, "r") as handle:
                raw = json.load(handle)

            enriched = dict(raw)
            # Rewrite paths to prefer colocated dataset snapshot
            enriched["audio_path"] = self._resolve_audio_path(raw["audio_path"])
            enriched["reconstructed_audio_path"] = str(
                self.precomputed_dir / Path(raw["reconstructed_audio_path"]).name
            )

            record = PrecomputedRecord(file_id=enriched["file_id"], data=enriched)
            self._records[record.file_id] = record

    def ids(self) -> Iterable[str]:
        return sorted(self._records.keys())

    def get(self, file_id: str) -> Dict[str, object]:
        if file_id not in self._records:
            raise KeyError(f"Unknown file id '{file_id}'")
        return self._records[file_id].data

    def _resolve_audio_path(self, original_path: str) -> str:
        original_name = Path(original_path).name
        local_candidate = self.data_root / original_name
        if local_candidate.exists():
            return str(local_candidate)
        return str(Path(original_path))