"""Corpus download (pinned + licensed), text pretokenization, dataset validation and statistics.

Every external corpus is listed in SOURCES with a pinned revision and its license, and the
provenance record is written next to the downloaded files so caches can be audited later.
"""
from __future__ import annotations

import hashlib
import json
from pathlib import Path

import numpy as np
from PIL import Image

SOURCES: dict[str, dict] = {
    "tinystories": {
        "repo_id": "roneneldan/TinyStories",
        "repo_type": "dataset",
        "revision": "f54c09fd23315a6f9c86f9dc80f725de7d8f9c64",
        "files": [
            "data/train-00000-of-00004-2d5a1467fff1081b.parquet",
            "data/validation-00000-of-00001-869c898b519ad725.parquet",
        ],
        "license": "CDLA-Sharing-1.0",
        "url": "https://huggingface.co/datasets/roneneldan/TinyStories",
        "notes": "Synthetic children's stories (GPT-3.5/4 generated). Only 1 of 4 train shards is used.",
    },
    "dolly": {
        "repo_id": "databricks/databricks-dolly-15k",
        "repo_type": "dataset",
        "revision": "bdd27f4d94b9c1f951818a7da7fd7aeea5dbff1a",
        "files": ["databricks-dolly-15k.jsonl"],
        "license": "CC-BY-SA-3.0",
        "url": "https://huggingface.co/datasets/databricks/databricks-dolly-15k",
        "notes": "Human-written instruction/response pairs; used for instruction tuning only.",
    },
}


def sha256_file(path: str | Path) -> str:
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(1 << 20), b""):
            h.update(chunk)
    return h.hexdigest()


def download_sources(raw_dir: str | Path, names: list[str] | None = None) -> dict[str, list[Path]]:
    """Download the pinned corpus files and write a provenance.json record."""
    from huggingface_hub import hf_hub_download

    raw_dir = Path(raw_dir)
    out: dict[str, list[Path]] = {}
    for name in names or list(SOURCES):
        src = SOURCES[name]
        dest = raw_dir / name
        paths = [
            Path(hf_hub_download(src["repo_id"], f, repo_type=src["repo_type"], revision=src["revision"], local_dir=dest))
            for f in src["files"]
        ]
        record = {k: v for k, v in src.items()} | {"sha256": {p.name: sha256_file(p) for p in paths}}
        (dest / "provenance.json").write_text(json.dumps(record, indent=2))
        out[name] = paths
    return out


def load_tinystories(raw_dir: str | Path, split: str) -> list[str]:
    import pyarrow.parquet as pq

    files = sorted((Path(raw_dir) / "tinystories" / "data").glob(f"{split}-*.parquet"))
    if not files:
        raise FileNotFoundError(f"No TinyStories {split} parquet under {raw_dir}; run scripts/generate_dataset.py download")
    texts: list[str] = []
    for f in files:
        texts += [t.strip() for t in pq.read_table(f, columns=["text"]).column("text").to_pylist() if t and t.strip()]
    return texts


def load_dolly(raw_dir: str | Path) -> list[dict]:
    rows = []
    with open(Path(raw_dir) / "dolly" / "databricks-dolly-15k.jsonl", encoding="utf-8") as f:
        for line in f:
            r = json.loads(line)
            if r["instruction"].strip() and r["response"].strip():
                rows.append(r)
    return rows


def tokenize_stream(texts: list[str], encode_batch, prefix: list[int], eos_id: int, out_path: str | Path) -> int:
    """Concatenate documents as <prefix> doc [EOS] ... into one uint16 token stream (for packing)."""
    out_path = Path(out_path)
    out_path.parent.mkdir(parents=True, exist_ok=True)
    n = 0
    with open(out_path, "wb") as f:
        for i in range(0, len(texts), 4096):
            arrs = [np.array([*prefix, *ids, eos_id], dtype=np.uint16) for ids in encode_batch(texts[i : i + 4096])]
            buf = np.concatenate(arrs)
            buf.tofile(f)
            n += len(buf)
    return n


def validate_sequence(tokens: np.ndarray, vocab) -> None:
    """Raise ValueError if a token sequence violates the unified-vocabulary structure rules."""
    vocab.check_structure(tokens)


def dataset_stats(cache_dir: str | Path) -> dict:
    """Summarise every cache file (token counts / sample counts) for docs and the dashboard."""
    cache_dir = Path(cache_dir)
    stats: dict[str, dict] = {}
    for p in sorted(cache_dir.glob("*.bin")):
        stats[p.name] = {"tokens": p.stat().st_size // 2}
    for p in sorted(cache_dir.glob("*.npz")):
        with np.load(p) as z:
            stats[p.name] = {k: list(z[k].shape) for k in z.files}
    return stats


# ----------------------------------------------------------------------------- media loading
def prepare_image(img, res: int = 64) -> np.ndarray:
    """PIL / array / path -> uint8 [res,res,3] (centre square crop, area resize)."""
    if isinstance(img, (str, Path)):
        img = Image.open(img)
    if isinstance(img, np.ndarray):
        img = Image.fromarray(img.astype(np.uint8))
    img = img.convert("RGB")
    w, h = img.size
    s = min(w, h)
    img = img.crop(((w - s) // 2, (h - s) // 2, (w - s) // 2 + s, (h - s) // 2 + s))
    return np.asarray(img.resize((res, res), Image.BOX if s >= res else Image.BICUBIC))


def load_video(src, max_frames: int = 8, fps: float = 4.0, res: int = 64) -> np.ndarray:
    """Video file (mp4/gif/webm...) or frame list -> uint8 [T,res,res,3], resampled to ~`fps`."""
    if not isinstance(src, (str, Path)):
        return np.stack([prepare_image(f, res) for f in list(src)[:max_frames]])
    import imageio_ffmpeg

    reader = imageio_ffmpeg.read_frames(str(src))
    meta = next(reader)
    w, h = meta["size"]
    stride = max(1, round(float(meta.get("fps") or fps) / fps))
    frames = []
    for i, raw in enumerate(reader):
        if i % stride == 0:
            frames.append(prepare_image(np.frombuffer(raw, np.uint8).reshape(h, w, 3), res))
        if len(frames) == max_frames:
            break
    reader.close()
    if not frames:
        raise ValueError(f"could not decode any frames from {src}")
    return np.stack(frames)


def video_windows(path: str | Path, n_frames: int = 8, fps: float = 4.0, max_windows: int = 4, res: int = 64) -> list[np.ndarray]:
    """Cut a video file into up to `max_windows` consecutive `n_frames`-frame clips at ~`fps`.
    Clips too short for one window are sampled evenly at native frame rate; < n_frames frames -> []."""
    clip = load_video(path, max_frames=n_frames * max_windows, fps=fps, res=res)
    if len(clip) >= n_frames:
        return [clip[i : i + n_frames] for i in range(0, len(clip) - n_frames + 1, n_frames)]
    native = load_video(path, max_frames=10_000, fps=1e9, res=res)  # every frame
    if len(native) < n_frames:
        return []
    return [native[np.linspace(0, len(native) - 1, n_frames).round().astype(int)]]


# ----------------------------------------------------------------------------- user datasets
IMAGE_EXT = {".png", ".jpg", ".jpeg", ".webp", ".bmp"}
VIDEO_EXT = {".mp4", ".mov", ".webm", ".mkv", ".avi", ".gif"}
USER_KEYS = ("user", "instruction", "prompt", "question", "input")
ASSISTANT_KEYS = ("assistant", "response", "output", "answer", "completion")


def scan_user_dataset(root: str | Path) -> dict[str, list]:
    """Collect samples from a user-provided folder or .jsonl file, routing records by their keys.

    JSONL records (paths relative to the .jsonl file):
        {"text": ...}                                             -> text corpus
        {"image": p, "caption": ...}                              -> image-caption (t2i + i2t)
        {"image": p, "target_image": p, "instruction": ...}       -> image editing (i2i)
        {"video": p, "caption": ...}  (caption optional)          -> video-caption (t2v, v2t) / continuation
        {"instruction"|"prompt"|"user": ..., "response"|"output"|"assistant": ...}  -> instruction tuning
        {"messages": [{"role": "user", ...}, {"role": "assistant", ...}, ...]}       -> instruction tuning
    Folders: *.txt -> text documents; image/video files with a same-name .txt sidecar -> captioned
    media (images without a caption are skipped; videos without one are kept for continuation).
    """
    root = Path(root)
    out: dict[str, list] = {"texts": [], "images": [], "edits": [], "videos": [], "chats": [], "skipped": []}
    files = [root] if root.is_file() else sorted(p for p in root.rglob("*") if p.is_file())
    sidecar = {p.with_suffix(".txt") for p in files if p.suffix.lower() in IMAGE_EXT | VIDEO_EXT}
    for p in files:
        if p.suffix.lower() == ".jsonl":
            for ln, line in enumerate(p.read_text(encoding="utf-8", errors="replace").splitlines()):
                if not line.strip():
                    continue
                try:
                    _route(json.loads(line), p.parent, out)
                except (ValueError, KeyError, TypeError) as e:
                    out["skipped"].append(f"{p.name}:{ln + 1}: {e}")
    # media referenced by a manifest record is not also scanned as a loose file
    referenced = {Path(x).resolve() for k in ("images", "videos") for x, _ in out[k]} | {Path(x).resolve() for e in out["edits"] for x in e[:2]}
    for p in files:
        ext = p.suffix.lower()
        if ext == ".jsonl" or p.resolve() in referenced:
            continue
        if ext == ".txt" and p not in sidecar:
            t = p.read_text(encoding="utf-8", errors="replace").strip()
            if t:
                out["texts"].append(t)
        elif ext in IMAGE_EXT | VIDEO_EXT:
            cap_p = p.with_suffix(".txt")
            cap = cap_p.read_text(encoding="utf-8", errors="replace").strip() if cap_p.exists() else ""
            if ext in VIDEO_EXT:
                out["videos"].append((p, cap))
            elif cap:
                out["images"].append((p, cap))
            else:
                out["skipped"].append(f"{p.name}: image without caption")
    return out


def _route(r: dict, base: Path, out: dict) -> None:
    get = lambda keys: next((str(r[k]).strip() for k in keys if r.get(k)), "")
    if "messages" in r:
        msgs = r["messages"]
        for a, b in zip(msgs, msgs[1:]):
            if a.get("role") == "user" and b.get("role") == "assistant" and a.get("content") and b.get("content"):
                out["chats"].append((str(a["content"]).strip(), str(b["content"]).strip()))
    elif "video" in r:
        out["videos"].append((base / r["video"], get(("caption", "text", "prompt"))))
    elif "image" in r and ("target_image" in r or "edited_image" in r):
        out["edits"].append((base / r["image"], base / (r.get("target_image") or r["edited_image"]), get(("instruction", "edit", "prompt"))))
    elif "image" in r:
        cap = get(("caption", "text", "prompt"))
        if not cap:
            raise ValueError("image record without caption")
        out["images"].append((base / r["image"], cap))
    elif get(USER_KEYS) and get(ASSISTANT_KEYS):
        user = get(USER_KEYS) + (("\n\n" + str(r["context"]).strip()) if r.get("context") else "")
        out["chats"].append((user, get(ASSISTANT_KEYS)))
    elif get(("text", "content", "story")):
        out["texts"].append(get(("text", "content", "story")))
    else:
        raise ValueError(f"unrecognised record keys {sorted(r)}")


def split_of(key: str, val_pct: float) -> str:
    """Deterministic train/val assignment from a content hash (stable across runs and machines)."""
    return "val" if int(hashlib.sha1(key.encode("utf-8", "replace")).hexdigest()[:8], 16) % 10_000 < val_pct * 100 else "train"