"""Dataset: a lazy handle over a registered dataset.

``get_dataset('gsm8k')`` costs nothing -- no download, no parsing. The first
actual use (iteration / len / indexing) triggers ``materialize()``:

    cache hit  -> read the cached samples.jsonl
    cache miss -> download raw -> record_to_sample -> atomic cache write -> read

    Cache directory name: ``{safe_name}-{md5(source+split+subset+version+params)}``
    so any config change yields a different cache entry (zero invalidation logic).

    Each cache entry is self-contained:
        raw/           native source data, exactly as downloaded (never converted)
        samples.jsonl  the unified Sample stream converted from raw/
        meta.json      spec + provenance
    """

import fcntl
import hashlib
import json
import os
import re
import string
import time
from pathlib import Path
from typing import Any, Callable, Dict, Iterator, List, Optional, Union

from .loader import get_cache_root, load_raw_records
from .sample import Sample
from .spec import DatasetSpec

# Backward-compatible alias: the canonical root lives in loader.get_cache_root()
# (a function, so runtime overrides via set_cache_root are always respected).
CACHE_ROOT = get_cache_root()


def set_cache_root(path) -> None:
    """Override the cache root at runtime (used by the CLI --cache-dir flag)."""
    from . import loader

    loader.set_cache_root(path)


def safe_filename(s: str, max_length: int = 255) -> str:
    safe_chars = string.ascii_letters + string.digits + '.-_'
    s = ''.join(c if c in safe_chars else '_' for c in s)
    s = re.sub(r'_+', '_', s).strip('._')
    return (s or 'untitled')[:max_length]


def gen_hash(s: str) -> str:
    return hashlib.md5(s.encode('utf-8')).hexdigest()


class Dataset:
    """Sequence-like lazy dataset. Nothing is downloaded until first use."""

    def __init__(self, spec: DatasetSpec, record_fn: Callable, samples: Optional[List[Sample]] = None):
        self.spec = spec
        self._record_fn = record_fn
        self._samples = samples  # None => not materialized yet
        self.lineage: Dict[str, Any] = {}

    # ---------------- materialization ----------------

    @property
    def cache_dir(self) -> Path:
        key = f'{self.spec.source}{self.spec.split}{self.spec.subset}{self.spec.version}{self.spec.params}'
        # Layout: datasets/<benchmark_name>/<subset>_<split>[-<version>]-<hash6>/
        # Readable benchmark folder + readable subset/split; the 6-char hash
        # suffix disambiguates different sources/params that would otherwise
        # collide on the same subset_split name (correctness requirement).
        parts = f'{safe_filename(self.spec.subset)}_{safe_filename(self.spec.split)}'
        if self.spec.version:
            parts += f'_{safe_filename(self.spec.version)}'
        subdir = f'{parts}-{gen_hash(key)[:6]}'
        return get_cache_root() / 'datasets' / safe_filename(self.spec.name) / subdir

    @property
    def is_materialized(self) -> bool:
        return self._samples is not None

    def materialize(self, force: bool = False) -> 'Dataset':
        if self._samples is not None and not force:
            return self
        cache_dir = self.cache_dir
        cache_file = cache_dir / 'samples.jsonl'
        if cache_file.exists() and not force:
            self._samples = self._read_cache(cache_file)
            self.lineage = {'from': 'cache', 'cache_dir': str(cache_dir)}
            return self

        # mkdir+lock with retries: SOMETHING reaps freshly created dataset
        # dirs during heavy concurrent runs; retry a few times before giving up
        lock_f = None
        last_err = None
        for _ in range(5):
            try:
                cache_dir.parent.mkdir(parents=True, exist_ok=True)
                lock_path = cache_dir.with_suffix('.lock')
                lock_f = open(lock_path, 'w')  # noqa: PTH123
                break
            except (FileExistsError, FileNotFoundError) as e:
                last_err = e
                import time as _t

                _t.sleep(1.0)
        if lock_f is None:
            raise last_err
        with lock_f:
            fcntl.flock(lock_f, fcntl.LOCK_EX)
            try:
                if cache_file.exists() and not force:  # double-check under lock
                    self._samples = self._read_cache(cache_file)
                    self.lineage = {'from': 'cache', 'cache_dir': str(cache_dir)}
                    return self
                cache_dir.parent.mkdir(parents=True, exist_ok=True)
                tmp_dir = cache_dir.with_name(cache_dir.name + f'.tmp-{os.getpid()}')
                if tmp_dir.exists():
                    import shutil

                    shutil.rmtree(tmp_dir)
                # concurrent materialization (parallel stage runners) can race
                # on the parent-chain mkdir; retry once -- the dir existing is
                # always harmless for a scratchpad
                try:
                    tmp_dir.mkdir(parents=True, exist_ok=True)
                except FileExistsError:
                    tmp_dir.mkdir(parents=True, exist_ok=True)
                records = load_raw_records(self.spec, raw_dir=tmp_dir / 'raw')
                samples = [self._to_sample(r) for r in records]
                self._assign_ids(samples)
                self._write_cache_atomic(samples, cache_dir, tmp_dir)
                self._samples = samples
                self.lineage = {'from': 'source', 'cache_dir': str(cache_dir)}
            finally:
                fcntl.flock(lock_f, fcntl.LOCK_UN)
        return self

    def _to_sample(self, record: Dict[str, Any]) -> Sample:
        sample = self._record_fn(record)
        if not sample.task_type:
            sample.task_type = self.spec.task_type
        # expose the loaded subset on every sample: per-subtask eval dispatch
        # (e.g. bbh MC vs free-form) reads metadata['subset']
        if self.spec.subset not in ('default', 'all') and sample.metadata.get('subset') is None:
            sample.metadata.setdefault('subset', self.spec.subset)
        return sample

    @staticmethod
    def _assign_ids(samples: List[Sample]) -> None:
        """Assign sequential ids to samples that don't carry one."""
        for i, s in enumerate(samples):
            if s.id is None:
                s.id = i

    def _read_cache(self, cache_file: Path) -> List[Sample]:
        with open(cache_file, encoding='utf-8') as f:
            samples = [Sample.model_validate(json.loads(line)) for line in f if line.strip()]
        return samples

    def _write_cache_atomic(self, samples: List[Sample], cache_dir: Path, tmp_dir: Path) -> None:
        """Finalize the tmp dir (samples + meta next to raw/) and swap it in.

        ``raw/`` was already populated by ``load_raw_records`` inside tmp_dir;
        the tmp+rename swap makes the whole entry (raw + samples + meta)
        appear atomically.
        """
        with open(tmp_dir / 'samples.jsonl', 'w', encoding='utf-8') as f:
            for s in samples:
                f.write(json.dumps(s.model_dump(), ensure_ascii=False) + '\n')
        raw_dir = tmp_dir / 'raw'
        meta = {
            'spec': {k: v for k, v in vars(self.spec).items()},
            'num_samples': len(samples),
            'raw_files': sorted(p.name for p in raw_dir.iterdir()) if raw_dir.exists() else [],
            'note': 'raw/ holds the native source data exactly as downloaded; samples.jsonl is the converted view',
            'created_at': time.strftime('%Y-%m-%d %H:%M:%S'),
        }
        with open(tmp_dir / 'meta.json', 'w', encoding='utf-8') as f:
            json.dump(meta, f, ensure_ascii=False, indent=2)
        final = cache_dir
        if final.exists():
            import shutil

            # retire the old entry, then swap the new one in
            old = cache_dir.with_name(cache_dir.name + f'.old-{os.getpid()}')
            os.rename(final, old)
            try:
                os.rename(tmp_dir, final)
            except OSError:
                os.rename(old, final)
                raise
            shutil.rmtree(old, ignore_errors=True)
        else:
            os.rename(tmp_dir, final)

    def unload(self) -> bool:
        """Drop this dataset's cache entry (raw/ + samples.jsonl + meta.json).

        Pure cache management: in-memory samples (if any) stay usable; the
        next materialize rebuilds from source. Execution-environment
        resources (docker images declared by ``Sample.sandbox``) are NOT
        touched -- those belong to the sandbox layer's lifecycle.
        """
        import shutil

        removed = False
        if self.cache_dir.exists():
            shutil.rmtree(self.cache_dir)
            removed = True
        lock_path = self.cache_dir.with_suffix('.lock')
        if lock_path.exists():
            lock_path.unlink()
        if self._samples is not None or self.lineage:
            self._samples = None
            self.lineage = {}
        return removed

    # ---------------- sequence protocol (triggers materialize) ----------------

    def _require(self) -> List[Sample]:
        self.materialize()
        return self._samples

    def __iter__(self) -> Iterator[Sample]:
        return iter(self._require())

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

    def __getitem__(self, i: Union[int, slice]) -> Union[Sample, List[Sample]]:
        return self._require()[i]

    # ---------------- views / derived data ----------------

    def view(self, samples: List[Sample], lineage: Optional[Dict[str, Any]] = None) -> 'Dataset':
        """An in-memory derived dataset (filter/sample/dedup results).

        Same class, same interface; ``lineage`` records how it was produced.
        """
        derived = Dataset(self.spec, self._record_fn, samples=samples)
        derived.lineage = {'from': 'derived', 'parent': self.spec.name, **(lineage or {})}
        return derived

    # ---------------- introspection ----------------

    def stats(self) -> Dict[str, Any]:
        samples = self._require()
        lengths = [len(s.input_text) for s in samples]
        targets = [s.target if isinstance(s.target, str) else ','.join(s.target) for s in samples]
        return {
            'name': self.spec.name,
            'task_type': self.spec.task_type,
            'num_samples': len(samples),
            'input_len': {
                'min': min(lengths) if lengths else 0,
                'max': max(lengths) if lengths else 0,
                'mean': round(sum(lengths) / len(lengths), 1) if lengths else 0,
            },
            'target_top': sorted({t: targets.count(t) for t in set(targets)}.items(), key=lambda kv: -kv[1])[:10],
            'cache_dir': str(self.cache_dir),
        }

    def __repr__(self) -> str:
        state = 'materialized' if self.is_materialized else 'lazy'
        return f"Dataset(name={self.spec.name!r}, type={self.spec.task_type!r}, {state})"