"""Task packs: a directory with tasks.jsonl and manifest.json. Each task carries its references (the exact optimum and
the baselines); `public()` strips them, with the record's releases and flows, before anything reaches an agent."""

from __future__ import annotations

import hashlib
import json
import os
from pathlib import Path

HIDDEN = ("reference",)
HIDDEN_DAM = ("record_release_cfs",)
HIDDEN_POINT = ("record_flow_cfs", "local_cfs")


def tasks_root() -> Path:
    here = Path(__file__).resolve()
    for p in here.parents:
        if (p / "tasks").is_dir() and any((p / "tasks").glob("*/tasks.jsonl")):
            return p / "tasks"
    return here.parents[2] / "tasks"


class TaskPack:
    def __init__(self, tasks: list[dict]):
        self.tasks = tasks
        self._by_id = {t["task_id"]: t for t in tasks}
        self._by_split: dict[str, list[dict]] = {}
        for t in tasks:
            self._by_split.setdefault(t["split"], []).append(t)

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

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

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

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

    @staticmethod
    def public(task: dict) -> dict:
        t = {k: v for k, v in task.items() if k not in HIDDEN}
        t["dams"] = [{k: v for k, v in d.items() if k not in HIDDEN_DAM} for d in task["dams"]]
        t["points"] = [{k: v for k, v in p.items() if k not in HIDDEN_POINT} for p in task["points"]]
        return t


def read_pack(d: Path) -> list[dict]:
    return [json.loads(line) for line in (Path(d) / "tasks.jsonl").read_text().splitlines() if line.strip()]


def write_pack(d: Path, tasks: list[dict], meta: dict):
    d = Path(d)
    d.mkdir(parents=True, exist_ok=True)
    body = "".join(json.dumps(t, separators=(",", ":")) + "\n" for t in tasks)
    (d / "tasks.jsonl").write_text(body)
    manifest = {**meta, "n": len(tasks), "sha256": hashlib.sha256(body.encode()).hexdigest(),
                "tasks": [{"task_id": t["task_id"], "split": t["split"], "tier": t["tier"],
                           "opt_cost": t["reference"]["opt_cost"], "lazy_cost": t["reference"]["lazy_cost"]}
                          for t in tasks]}
    (d / "manifest.json").write_text(json.dumps(manifest, indent=1))


def load_pack(dirs: str | None = None) -> TaskPack:
    dirs = dirs or os.environ.get("FLOOD_TASKS_DIR")
    paths = [Path(p) for p in dirs.split(":")] if dirs else sorted(p.parent for p in tasks_root().glob("*/tasks.jsonl"))
    tasks = []
    for p in paths:
        tasks += read_pack(p)
    return TaskPack(tasks)