"""The task pack: prepared tasks on disk, and the catalog the server reads them from.

Layout:

    <pack>/
      manifest.json                 version, source revision, split counts
      fonts/manifest.json + files   the 263 Crello families, extracted
      tasks/<split>.jsonl           one task per line (ground truth included)
      designs/<crello_id>/          assets/<aid>.png|webp and preview.jpg

A task is one Crello template. How it is played is chosen at reset: `mode`
(reference image or text description) and `surface` (tools, HTML or GUI) are
episode parameters, not separate tasks, so every split serves every variant and
the Task API stays a plain list of designs.
"""

from __future__ import annotations

import json
import os
import threading
import time
from collections import OrderedDict
from pathlib import Path

from .describe import describe
from .doc import DesignDoc
from .fonts import FontRegistry, extract_crello_fonts
from .render import Renderer

PACK_VERSION = 1
SPLITS = ("train", "validation", "test")
DEFAULT_PACK = Path(os.environ.get("DESIGN_CANVAS_PACK", Path.home() / ".cache/design_canvas/pack"))
CRELLO_REPO = "cyberagent/crello"
CRELLO_REVISION = "5.1.0"


def difficulty(n: int) -> str:
    return "easy" if n <= 6 else "medium" if n <= 12 else "hard"


class TaskPack:
    """Read-only view of a prepared pack. Safe to share across sessions and threads."""

    def __init__(self, root: Path | str):
        self.root = Path(root)
        manifest = self.root / "manifest.json"
        if not manifest.exists():
            raise FileNotFoundError(
                f"no task pack at {self.root}. Build one with `design-canvas prepare` "
                f"or point DESIGN_CANVAS_PACK at an existing pack."
            )
        self.manifest = json.loads(manifest.read_text())
        self.fonts = FontRegistry(self.root / "fonts")
        self.renderer = Renderer(self.fonts)
        self._tasks: dict[str, list[dict]] = {}
        self._by_id: dict[str, dict] = {}
        self._docs: dict[str, DesignDoc] = {}
        self._renders: OrderedDict[tuple, object] = OrderedDict()
        self._lock = threading.Lock()
        for split in self.splits():
            path = self.root / "tasks" / f"{split}.jsonl"
            tasks = [json.loads(line) for line in path.read_text().splitlines() if line.strip()]
            self._tasks[split] = tasks
            for t in tasks:
                self._by_id[t["task_id"]] = t

    # ------------------------------------------------------------ index

    def splits(self) -> list[str]:
        return [s for s in SPLITS if s in self.manifest.get("splits", {})]

    def count(self, split: str) -> int:
        if split not in self._tasks:
            raise KeyError(f"unknown split {split!r}; have {', '.join(self.splits())}")
        return len(self._tasks[split])

    def at(self, split: str, index: int) -> dict:
        tasks = self._tasks.get(split)
        if tasks is None:
            raise KeyError(f"unknown split {split!r}; have {', '.join(self.splits())}")
        if not 0 <= index < len(tasks):
            raise IndexError(f"{split} has {len(tasks)} tasks; index {index} is out of range")
        return tasks[index]

    def get(self, task_id: str) -> dict:
        try:
            return self._by_id[task_id]
        except KeyError:
            raise KeyError(f"unknown task_id {task_id!r}") from None

    @staticmethod
    def public(task: dict) -> dict:
        """The Task API view: enough to choose a task, nothing that gives it away."""
        return {
            "task_id": task["task_id"],
            "index": task["index"],
            "split": task["split"],
            "crello_id": task["crello_id"],
            "title": task["title"],
            "format": task["format"],
            "category": task["category"],
            "canvas": task["canvas"],
            "n_elements": task["stats"]["n_elements"],
            "n_text": task["stats"]["n_text"],
            "n_assets": len(task["assets"]),
            "difficulty": task["difficulty"],
            "modes": ["reference", "description"],
            "surfaces": ["tools", "html", "gui"],
        }

    def task_range(self, split: str, start=None, stop=None) -> list[dict]:
        tasks = self._tasks.get(split)
        if tasks is None:
            raise KeyError(f"unknown split {split!r}")
        return [self.public(t) for t in tasks[slice(start, stop)]]

    # --------------------------------------------------------- per task

    def gt_doc(self, task: dict) -> DesignDoc:
        tid = task["task_id"]
        with self._lock:
            doc = self._docs.get(tid)
            if doc is None:
                doc = DesignDoc.model_validate(task["doc"])
                self._docs[tid] = doc
        return doc

    def asset_paths(self, task: dict) -> dict[str, Path]:
        return {a["id"]: self.root / a["file"] for a in task["assets"]}

    def asset_ctx(self, task: dict) -> dict[str, dict]:
        return {a["id"]: {"aspect": a["aspect"], "kind": a["kind"]} for a in task["assets"]}

    def preview_path(self, task: dict) -> Path:
        return self.root / task["preview"]

    def _cached(self, key: tuple, make):
        with self._lock:
            if key in self._renders:
                self._renders.move_to_end(key)
                return self._renders[key]
        value = make()
        with self._lock:
            self._renders[key] = value
            while len(self._renders) > 256:
                self._renders.popitem(last=False)
        return value

    def target_png(self, task: dict, max_size: int = 512) -> bytes:
        return self._cached(
            ("png", task["task_id"], max_size),
            lambda: self.renderer.render_png(self.gt_doc(task), self.asset_paths(task), max_size),
        )

    def target_jpeg(self, task: dict, max_size: int = 768) -> bytes:
        return self._cached(
            ("jpg", task["task_id"], max_size),
            lambda: self.renderer.render_jpeg(self.gt_doc(task), self.asset_paths(task), max_size),
        )

    def target_pixels(self, task: dict, max_size: int):
        return self._cached(
            ("arr", task["task_id"], max_size),
            lambda: self.renderer.render_array(self.gt_doc(task), self.asset_paths(task), max_size),
        )

    def asset_similarity(self, task: dict):
        """`sim(a, b)` in [0, 1]: how alike two of this task's assets look.

        Compared on 32 px thumbnails composited over grey (so the cut-out shape
        counts), on average colour, and on aspect ratio. Squared, so only
        genuine look-alikes earn real credit.
        """

        def make():
            import numpy as np
            from PIL import Image

            from .score import ssim

            thumbs, colors, aspect = {}, {}, {}
            for a in task["assets"]:
                with Image.open(self.root / a["file"]) as im:
                    im = im.convert("RGBA")
                    bg = Image.new("RGBA", im.size, (128, 128, 128, 255))
                    flat = Image.alpha_composite(bg, im).convert("RGB")
                thumbs[a["id"]] = np.asarray(flat.resize((32, 32), Image.BOX), dtype=np.float64)
                colors[a["id"]] = np.asarray(flat.resize((4, 4), Image.BOX), dtype=np.float64)
                aspect[a["id"]] = a["aspect"]
            table = {}
            ids = list(thumbs)
            for i, x in enumerate(ids):
                for y in ids[i + 1:]:
                    shape = ssim(thumbs[x], thumbs[y])
                    color = float(np.exp(-np.abs(colors[x] - colors[y]).mean() / 255 / 0.1))
                    ratio = float(np.exp(-abs(np.log(aspect[x] / aspect[y])) / 0.25))
                    sim = (0.4 * shape + 0.3 * color + 0.3 * ratio) ** 2
                    table[(x, y)] = table[(y, x)] = round(sim, 4)
            return table

        table = self._cached(("assetsim", task["task_id"]), make)
        return lambda a, b: 1.0 if a == b else table.get((a, b), 0.0)

    def grade(self, task: dict, doc: DesignDoc, mode: str = "reference", version: int | None = None):
        """The one scoring entry point: episodes, tests and scripts all call this."""
        from .score import CONFIGS, grade

        size = CONFIGS[mode].pixel_size
        attempt = self.renderer.render_array(doc, self.asset_paths(task), size)
        return grade(doc, self.gt_doc(task), mode, attempt, self.target_pixels(task, size),
                     self.blank_pixels(task, size), fonts=self.fonts,
                     asset_sim=self.asset_similarity(task),
                     version=version or int(os.environ.get("DESIGN_CANVAS_REWARD", "1")))

    def blank_pixels(self, task: dict, max_size: int):
        def make():
            doc = self.gt_doc(task)
            return self.renderer.render_array(
                DesignDoc(width=doc.width, height=doc.height), {}, max_size
            )

        return self._cached(("blank", task["task_id"], max_size), make)


_PACKS: dict[str, TaskPack] = {}
_PACKS_LOCK = threading.Lock()


def load_pack(root: Path | str | None = None) -> TaskPack:
    """The task catalog the server plays from.

    `DESIGN_CANVAS_CORPUS` (a corpus manifest: local path or `hf://buckets/...`)
    selects the full bucket-backed corpus; otherwise a prepared pack directory
    (`root`, `DESIGN_CANVAS_PACK`, or the default) is served.
    """
    corpus = None if root else os.environ.get("DESIGN_CANVAS_CORPUS")
    key = corpus or str(Path(root or DEFAULT_PACK).expanduser().resolve())
    with _PACKS_LOCK:
        pack = _PACKS.get(key)
        if pack is None:
            if corpus:
                from .corpus import CorpusPack

                pack = CorpusPack(corpus)
            else:
                pack = TaskPack(key)
            _PACKS[key] = pack
        return pack


# ------------------------------------------------------------ building


def accept(design, max_elements: int) -> str | None:
    """Why a design is rejected, or None to keep it."""
    n = len(design.doc.elements)
    if n < 2:
        return "fewer than two elements"
    if n > max_elements:
        return f"more than {max_elements} elements"
    if not any(e.kind == "text" for e in design.doc.elements):
        return "no text"
    if any(w.startswith("font fallback") for w in design.warnings):
        return "font not in the Crello font set"
    if max(design.doc.width, design.doc.height) > 4000:
        return "canvas too large"
    return None


def _asset_ext(kind: str) -> str:
    """Photos as WebP (they are most of the bytes), graphics as PNG (crisp edges)."""
    return ".webp" if kind in ("photo", "background") else ".png"


def write_design(out: Path, design) -> None:
    """Write a converted design's assets and Crello preview under `out/designs/<id>/`."""
    ddir = Path(out) / "designs" / design.crello_id
    (ddir / "assets").mkdir(parents=True, exist_ok=True)
    for aid, img in design.assets.items():
        ext = _asset_ext(design.asset_meta[aid]["kind"])
        tmp = ddir / "assets" / f".{aid}{ext}.part"
        if ext == ".webp":
            img.save(tmp, "WEBP", quality=90, method=4)
        else:
            img.save(tmp, "PNG", optimize=True)
        tmp.replace(ddir / "assets" / f"{aid}{ext}")
    if design.preview is not None:
        tmp = ddir / ".preview.jpg.part"
        design.preview.convert("RGB").save(tmp, "JPEG", quality=88)
        tmp.replace(ddir / "preview.jpg")


def make_task(design, split: str, index: int) -> dict:
    """The task record for a converted design: ground truth, asset list, brief, stats."""
    cid = design.crello_id
    assets = sorted(
        ({**meta, "file": f"designs/{cid}/assets/{aid}{_asset_ext(meta['kind'])}"}
         for aid, meta in design.asset_meta.items()),
        key=lambda a: int(a["id"][1:]),
    )
    n_text = sum(e.kind == "text" for e in design.doc.elements)
    return {
        "task_id": f"{split}-{index:05d}",
        "index": index,
        "split": split,
        **design.meta,
        "doc": design.doc.model_dump(mode="json", exclude_defaults=False),
        "assets": assets,
        "fonts": sorted({e.font_family for e in design.doc.elements if e.kind == "text"}),
        "preview": f"designs/{cid}/preview.jpg",
        "description": describe(design.doc, design.meta, design.asset_meta),
        "stats": {
            "n_elements": len(design.doc.elements),
            "n_text": n_text,
            "n_images": len(design.doc.elements) - n_text,
        },
        "difficulty": difficulty(len(design.doc.elements)),
    }


def build_pack(
    out: Path | str,
    sources: dict[str, list[str]],
    per_split: dict[str, int],
    fonts_pickle: str,
    max_elements: int = 20,
    log=print,
) -> dict:
    """Convert Crello parquet shards into a task pack.

    `sources` maps split -> local parquet paths; `per_split` caps how many
    accepted designs each split keeps, taken in file order.
    """
    import datasets
    import pyarrow.parquet as pq

    from .crello import convert_row

    out = Path(out)
    (out / "tasks").mkdir(parents=True, exist_ok=True)
    if not (out / "fonts" / "manifest.json").exists():
        extract_crello_fonts(fonts_pickle, out / "fonts")
    fonts = FontRegistry(out / "fonts")
    features = datasets.load_dataset_builder(CRELLO_REPO, revision=CRELLO_REVISION).info.features

    counts = {}
    rejected: dict[str, int] = {}
    for split, paths in sources.items():
        want = per_split.get(split, 0)
        tasks = []
        for path in paths:
            if len(tasks) >= want:
                break
            pf = pq.ParquetFile(path)
            for rg in range(pf.num_row_groups):
                if len(tasks) >= want:
                    break
                for row in pf.read_row_group(rg).to_pylist():
                    if len(tasks) >= want:
                        break
                    design = convert_row(row, features, fonts)
                    reason = accept(design, max_elements)
                    if reason:
                        rejected[reason] = rejected.get(reason, 0) + 1
                        continue
                    write_design(out, design)
                    tasks.append(make_task(design, split, len(tasks)))
            log(f"{split}: {len(tasks)}/{want} after {Path(path).name}")
        with open(out / "tasks" / f"{split}.jsonl", "w") as fh:
            for t in tasks:
                fh.write(json.dumps(t, separators=(",", ":")) + "\n")
        counts[split] = len(tasks)

    manifest = {
        "version": PACK_VERSION,
        "created": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        "source": {"repo": CRELLO_REPO, "revision": CRELLO_REVISION,
                   "license": "CDLA-Permissive-2.0; templates (c) VistaCreate, see dataset card"},
        "splits": counts,
        "max_elements": max_elements,
        "rejected": rejected,
    }
    (out / "manifest.json").write_text(json.dumps(manifest, indent=2))
    return manifest