"""Task sources over the pretokenized caches and a seeded mixture loader.

Every micro-batch holds a single task (so lengths are similar and padding is small) and is
drawn from an RNG seeded by (seed, step, micro_index): sampling is a pure function of the step,
which makes runs reproducible and resumable without storing data-loader state.
"""
from __future__ import annotations

import json
import queue
import threading
from pathlib import Path

import numpy as np
import torch

from ..tokenizers.multimodal import Sample, UnifiedVocab
from ..tokenizers.text import TextTokenizer
from .collator import collate

BASE_TASKS = ("text", "t2i", "i2t", "i2i", "t2v", "v2t", "vcont")
CHAT_TASKS = ("chat_text", "chat_mm")
TASK_MODALITY = {"text": "text", "chat_text": "text", "t2i": "image", "i2t": "image", "i2i": "image",
                 "t2v": "video", "v2t": "video", "vcont": "video", "chat_mm": "multimodal"}

CHAT_PROMPTS = {
    "t2i": ["Generate an image of {c}.", "Draw {c}.", "Create a picture of {c}.", "Show me {c}.", "Can you make an image of {c}?"],
    "t2v": ["Make a video of {c}.", "Generate a video of {c}.", "Create a short clip of {c}.", "Animate {c}."],
    "i2t": ["Describe this image.", "What do you see in this picture?", "What is in this image?"],
    "v2t": ["Describe this video.", "What is happening in this video?", "What happens in this clip?"],
    "vcont": ["Continue the movement in this video.", "Continue this video.", "What happens next? Show me."],
    "i2v": ["Animate this image: {c}.", "Bring this picture to life: {c}."],
}
CHAT_KINDS = ("t2i", "t2v", "i2t", "v2t", "vcont", "i2i", "i2v")


def _load_media(cache_dir: Path, name: str, split: str, vocab: UnifiedVocab, keys: tuple[str, ...]) -> tuple[dict, dict] | None:
    """Load <name>_<split>.npz/.json, dropping rows whose codes are out of range (corruption guard)."""
    npz, js = cache_dir / f"{name}_{split}.npz", cache_dir / f"{name}_{split}.json"
    if not npz.exists():
        return None
    try:
        with np.load(npz) as z:
            arrs = {k: z[k] for k in keys}
        meta = json.loads(js.read_text())
    except (OSError, ValueError, KeyError) as e:
        raise RuntimeError(f"cache {npz} is unreadable ({e}); regenerate with scripts/generate_dataset.py") from e
    ok = np.ones(len(arrs[keys[0]]), dtype=bool)
    for a in arrs.values():
        ok &= (a.reshape(len(a), -1) < vocab.visual_size).all(1)
    if not ok.all():
        print(f"[data] {npz.name}: dropping {int((~ok).sum())} corrupted rows")
    keep = np.nonzero(ok)[0]
    return {k: a[keep] for k, a in arrs.items()}, {k: [v[i] for i in keep] for k, v in meta.items()}


class TaskSources:
    """Builds one training Sample for any task from the caches of a given split."""

    def __init__(self, cache_dir: str | Path, split: str, vocab: UnifiedVocab, tok: TextTokenizer, seq_len: int, caption_dropout: float = 0.1):
        cache_dir = Path(cache_dir)
        self.vocab, self.tok, self.seq_len = vocab, tok, seq_len
        # Blank the t2i/t2v caption this often so the model also learns p(media) -> classifier-free guidance.
        self.caption_dropout = caption_dropout if split == "train" else 0.0
        tp = cache_dir / f"text_{split}.bin"
        self.text = np.memmap(tp, dtype=np.uint16, mode="r") if tp.exists() else None
        self.images = _load_media(cache_dir, "images", split, vocab, ("codes",))
        self.edits = _load_media(cache_dir, "edits", split, vocab, ("src", "tgt"))
        self.videos = _load_media(cache_dir, "videos", split, vocab, ("codes",))
        cp = cache_dir / f"chat_{split}.json"
        self.chat = json.loads(cp.read_text()) if cp.exists() else None
        # indices with a caption (user videos may be uncaptioned: continuation-only)
        cap_idx = lambda m: np.array([i for i, c in enumerate(m[1]["captions"]) if c], dtype=np.int64) if m else np.zeros(0, np.int64)
        self.img_cap, self.vid_cap = cap_idx(self.images), cap_idx(self.videos)
        need = {"t2i": len(self.img_cap), "i2t": len(self.img_cap), "t2v": len(self.vid_cap), "v2t": len(self.vid_cap),
                "vcont": self.videos is not None, "i2i": self.edits is not None, "i2v": len(self.vid_cap)}
        self.chat_kinds = [k for k in CHAT_KINDS if need[k]]

    def available(self) -> list[str]:
        have = {"text": self.text is not None and len(self.text) > self.seq_len + 1, "chat_text": bool(self.chat),
                "t2i": len(self.img_cap) > 0, "i2t": len(self.img_cap) > 0, "i2i": self.edits is not None,
                "t2v": len(self.vid_cap) > 0, "v2t": len(self.vid_cap) > 0, "vcont": self.videos is not None,
                "chat_mm": bool(self.chat_kinds)}
        return [t for t, ok in have.items() if ok]

    def enc(self, s: str) -> list[int]:
        return self.tok.encode(s)

    def _prompt(self, caps: list[str], rng: np.random.Generator) -> list[int]:
        return [] if rng.random() < self.caption_dropout else self.enc(str(rng.choice(caps)))

    def sample(self, task: str, rng: np.random.Generator) -> Sample:
        v = self.vocab
        if task == "text":
            off = int(rng.integers(0, len(self.text) - self.seq_len - 1))
            return Sample(np.asarray(self.text[off : off + self.seq_len + 1], dtype=np.int32), 0)
        if task in ("t2i", "i2t"):
            (arr, meta), i = self.images, int(self.img_cap[rng.integers(len(self.img_cap))])
            caps = meta["captions"][i]
            if task == "t2i":
                return v.build("t2i", self._prompt(caps, rng), image=arr["codes"][i])
            return v.build("i2t", self.enc(caps[0]), image=arr["codes"][i])
        if task == "i2i":
            (arr, meta), i = self.edits, int(rng.integers(len(self.edits[0]["src"])))
            return v.build("i2i", self.enc(meta["instructions"][i]), image=arr["src"][i], image2=arr["tgt"][i])
        if task in ("t2v", "v2t", "vcont"):
            pool = self.vid_cap if task != "vcont" else np.arange(len(self.videos[0]["codes"]))
            (arr, meta), i = self.videos, int(pool[rng.integers(len(pool))])
            caps = meta["captions"][i]
            if task == "t2v":
                return v.build("t2v", self._prompt(caps, rng), video=arr["codes"][i])
            if task == "v2t":
                return v.build("v2t", self.enc(caps[0]), video=arr["codes"][i])
            return v.build("vcont", video=arr["codes"][i])
        if task == "chat_text":
            r = self.chat[int(rng.integers(len(self.chat)))]
            s = v.chat(self.enc(r["user"]), self.enc(r["assistant"]))
            return Sample(s.tokens[: self.seq_len + 1], s.target_start)
        if task == "chat_mm":
            return self.chat_mm(str(rng.choice(self.chat_kinds)), rng)
        raise KeyError(task)

    def chat_mm(self, kind: str, rng: np.random.Generator) -> Sample:
        """Instruction-formatted multimodal conversation built from the media caches."""
        v, enc = self.vocab, self.enc
        pick = lambda xs: str(rng.choice(xs))
        if kind in ("t2i", "i2t"):
            (arr, meta), i = self.images, int(self.img_cap[rng.integers(len(self.img_cap))])
            cap, img = meta["captions"][i][0], arr["codes"][i]
            if kind == "t2i":
                return v.chat(enc(pick(CHAT_PROMPTS["t2i"]).format(c=pick(meta["captions"][i]))), enc(f"Here is an image of {cap}.") + v.image_tokens(img))
            return v.chat(v.image_tokens(img) + enc(" " + pick(CHAT_PROMPTS["i2t"])), enc(f"The image shows {cap}."))
        if kind == "i2i":
            (arr, meta), i = self.edits, int(rng.integers(len(self.edits[0]["src"])))
            instr = meta["instructions"][i]
            return v.chat(v.image_tokens(arr["src"][i]) + enc(" " + instr[0].upper() + instr[1:] + "."), enc("Here is the edited image.") + v.image_tokens(arr["tgt"][i]))
        pool = self.vid_cap if kind != "vcont" else np.arange(len(self.videos[0]["codes"]))
        (arr, meta), i = self.videos, int(pool[rng.integers(len(pool))])
        cap, vid = (meta["captions"][i] or [""])[0], arr["codes"][i]
        if kind == "t2v":
            return v.chat(enc(pick(CHAT_PROMPTS["t2v"]).format(c=pick(meta["captions"][i]))), enc(f"Here is a video of {cap}.") + v.video_tokens(vid))
        if kind == "v2t":
            return v.chat(v.video_tokens(vid) + enc(" " + pick(CHAT_PROMPTS["v2t"])), enc(f"The video shows {cap}."))
        if kind == "vcont":
            k = len(vid) // 2
            return v.chat(v.video_tokens(vid[:k]) + enc(" " + pick(CHAT_PROMPTS["vcont"])), enc("Here is the continuation.") + v.video_tokens(vid[k:], first_index=k))
        # i2v: the image is frame 0; the answer is the rest of the clip
        return v.chat(v.image_tokens(vid[0]) + enc(" " + pick(CHAT_PROMPTS["i2v"]).format(c=cap)), enc("Here is the animation.") + v.video_tokens(vid[1:], first_index=1))


def split_task(name: str) -> tuple[str, str]:
    """'user/t2i' -> ('user', 't2i'); 't2i' -> ('', 't2i')  (source prefix selects the cache)."""
    src, _, task = name.rpartition("/")
    return src, task


def available_tasks(sources: dict[str, TaskSources]) -> list[str]:
    return [f"{k}/{t}" if k else t for k, src in sources.items() for t in src.available()]


def compile_bucket_length(length: int, context: int) -> int:
    """At most four fixed sequence sizes at context 1024; no content truncation."""
    if not 0 < length <= context:
        raise ValueError("sample length is outside the model context")
    return min(context, max(128, 1 << (length - 1).bit_length()))


class MixtureLoader:
    """Yields (task, input_ids, labels) micro-batches; one task per micro-batch, ~`micro_tokens` tokens each.

    `sources` maps a source name to its TaskSources ('' = built-in cache, 'user' = imported data);
    mixture keys address them as 't2i' or 'user/t2i', so every source/task share is an explicit weight."""

    def __init__(self, sources: dict[str, TaskSources] | TaskSources, weights: dict[str, float], micro_tokens: int, seed: int, prefetch: int = 6,
                 compile_buckets: bool = False):
        sources = sources if isinstance(sources, dict) else {"": sources}
        missing = [t for t, w in weights.items() if w > 0 and t not in available_tasks(sources)]
        if missing:
            raise ValueError(f"mixture uses tasks with no cache: {missing}; available: {available_tasks(sources)}")
        self.sources, self.seed, self.micro_tokens = sources, seed, micro_tokens
        self.vocab = next(iter(sources.values())).vocab
        self.tasks = [t for t, w in weights.items() if w > 0]
        p = np.array([weights[t] for t in self.tasks], dtype=np.float64)
        self.p = p / p.sum()
        self.prefetch = prefetch
        self.compile_buckets = compile_buckets

    def batch(self, step: int, micro: int) -> tuple[str, torch.Tensor, torch.Tensor]:
        rng = np.random.default_rng([self.seed, step, micro])
        task = self.tasks[int(rng.choice(len(self.tasks), p=self.p))]
        samples: list[Sample] = []
        longest = 0
        src, base = split_task(task)
        while True:
            s = self.sources[src].sample(base, rng)
            proposed = max(longest, len(s.tokens) - 1 if self.compile_buckets else len(s.tokens))
            padded = compile_bucket_length(proposed, self.sources[src].seq_len) if self.compile_buckets else proposed
            if self.compile_buckets and padded > self.micro_tokens:
                raise ValueError("compiled batch token budget is smaller than its padded context")
            if samples and (len(samples) + 1) * padded > self.micro_tokens:
                break
            longest = proposed
            samples.append(s)
        if self.compile_buckets:
            length = compile_bucket_length(longest, self.sources[src].seq_len)
            # Empty rows have no target labels and do not change the supervised objective.
            inp, lab = collate(samples, self.vocab.pad, target_length=length,
                               batch_size=self.micro_tokens // length)
        else:
            inp, lab = collate(samples, self.vocab.pad)
        return task, inp, lab

    def iterate(self, start_step: int, grad_accum: int):
        """Background-thread prefetching iterator of (step, micro, task, inp, lab)."""
        q: queue.Queue = queue.Queue(maxsize=self.prefetch)
        stop = threading.Event()

        def work():
            step = start_step
            while not stop.is_set():
                for m in range(grad_accum):
                    item = (step, m, *self.batch(step, m))
                    while not stop.is_set():
                        try:
                            q.put(item, timeout=0.5)
                            break
                        except queue.Full:
                            continue
                step += 1

        th = threading.Thread(target=work, daemon=True)
        th.start()
        try:
            while True:
                yield q.get()
        finally:
            stop.set()


def fixed_eval_batches(sources: dict[str, TaskSources] | TaskSources, tasks: list[str], micro_tokens: int, n_batches: int, seed: int = 999) -> dict[str, list]:
    """Deterministic held-out batches per task (same every evaluation)."""
    sources = sources if isinstance(sources, dict) else {"": sources}
    out = {}
    for ti, task in enumerate(tasks):
        if task not in available_tasks(sources):
            continue
        ml = MixtureLoader(sources, {task: 1.0}, micro_tokens, seed)
        out[task] = [ml.batch(10_000 + ti, i)[1:] for i in range(n_batches)]
    return out