"""Task packs: directories of tasks (tasks.jsonl or tasks.jsonl.gz, with their reference plans) plus a manifest.

The server serves several packs at once; splits come from each task's `split` field. By default that is the dock-v1
eval and train packs (the hard tasks); the first pack, berth-v1, stays available by setting BERTH_TASKS_DIR.

    BERTH_TASKS_DIR=/path/a:/path/b     colon-separated pack directories
"""

from __future__ import annotations

import gzip
import hashlib
import json
import os
from pathlib import Path

from .model import Task

TASKS_ROOT = Path(__file__).resolve().parents[2] / "tasks"
DEFAULT_PACKS = ["dock-v1-eval", "dock-v1-train"]
LEGACY_PACK = "berth-v1"


def default_roots() -> list[Path]:
    env = os.environ.get("BERTH_TASKS_DIR", "").strip()
    if env:
        return [Path(p) for p in env.split(":") if p]
    roots = [TASKS_ROOT / p for p in DEFAULT_PACKS if _tasks_file(TASKS_ROOT / p)]
    return roots or [TASKS_ROOT / LEGACY_PACK]


def _tasks_file(root: Path) -> Path | None:
    for name in ("tasks.jsonl", "tasks.jsonl.gz"):
        if (root / name).is_file():
            return root / name
    return None


def _read(root: Path) -> tuple[list[Task], dict]:
    path = _tasks_file(root)
    if path is None:
        raise FileNotFoundError(f"no tasks.jsonl(.gz) in {root}")
    raw = path.read_bytes()
    body = gzip.decompress(raw) if path.suffix == ".gz" else raw
    manifest = json.loads((root / "manifest.json").read_text()) if (root / "manifest.json").is_file() else {}
    digest = hashlib.sha256(body).hexdigest()
    if manifest.get("sha256") and manifest["sha256"] != digest:
        raise ValueError(f"{path} does not match its manifest.json (sha256 {digest[:12]} != {manifest['sha256'][:12]})")
    return [Task.from_dict(json.loads(line)) for line in body.decode().splitlines() if line.strip()], manifest


class TaskPack:
    def __init__(self, roots: Path | str | list[Path | str] | None = None):
        if roots is None:
            roots = default_roots()
        if isinstance(roots, (str, Path)):
            roots = [roots]
        self.roots = [Path(r) for r in roots]
        self.root = self.roots[0]
        self.tasks: list[Task] = []
        self.manifests: dict[str, dict] = {}
        for r in self.roots:
            tasks, manifest = _read(r)
            self.tasks += tasks
            self.manifests[r.name] = manifest
        self.manifest = self.manifests[self.root.name]
        self._by_id = {}
        for t in self.tasks:
            if t.task_id in self._by_id:
                raise ValueError(f"task id {t.task_id} appears in more than one pack")
            self._by_id[t.task_id] = t
        self._by_split: dict[str, list[Task]] = {}
        for t in self.tasks:
            self._by_split.setdefault(t.split, []).append(t)

    def splits(self) -> list[str]:
        return sorted(self._by_split)

    def count(self, split: str) -> int:
        return len(self._by_split.get(split, []))

    def at(self, split: str, index: int) -> Task:
        return self._by_split[split][index]

    def get(self, task_id: str) -> Task:
        return self._by_id[task_id]

    def public(self, task: Task) -> dict:
        """What anyone may see: everything but the reference plans and costs."""
        return task.to_dict(public=True)


_PACK: TaskPack | None = None


def load_pack(root: Path | str | list | None = None) -> TaskPack:
    global _PACK
    if root is not None:
        return TaskPack(root)
    if _PACK is None:
        _PACK = TaskPack()
    return _PACK