"""Prepare a bounded, reproducible text-only dataset on the cloud machine."""
import argparse
from collections import Counter
import hashlib
import json
from pathlib import Path
import random
import re

from .runtime import CONFIG
from .articles import behavior_examples


def normalize_messages(row):
    messages = row.get("messages")
    if not isinstance(messages, list) or not messages:
        raise ValueError("Expected a nonempty messages list.")
    clean = []
    for index, item in enumerate(messages):
        role, content = item.get("role"), item.get("content")
        if role not in {"system", "user", "assistant"} or not isinstance(content, str) or not content.strip():
            raise ValueError("Only nonempty text messages are supported.")
        if role == "system" and index != 0:
            raise ValueError("System messages must come first.")
        clean.append({"role": role, "content": content.strip()})
    turns = clean[1:] if clean[0]["role"] == "system" else clean
    if len(turns) < 2 or len(turns) % 2:
        raise ValueError("Expected complete conversation turns.")
    if any(m["role"] != ("user" if i % 2 == 0 else "assistant") for i, m in enumerate(turns)):
        raise ValueError("Expected alternating user and assistant turns.")
    if clean[0]["role"] != "system":
        clean.insert(0, {"role": "system", "content": CONFIG["system_prompt"]})
    return clean


def prompt_key(messages):
    # Group identical opening questions, even when answers differ.
    first = next(m["content"] for m in messages if m["role"] == "user")
    normalized = re.sub(r"\s+", " ", first).strip().casefold()
    return hashlib.sha256(normalized.encode()).hexdigest()


def encode_example(tokenizer, messages, max_length):
    """Supervise assistant content and end-of-turn tokens; mask every prompt."""
    ids = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=False)
    if len(ids) > max_length:
        raise ValueError("Overlength example; skip rather than truncate code or its answer.")
    labels = [-100] * len(ids)
    for index, message in enumerate(messages):
        if message["role"] != "assistant":
            continue
        prefix = tokenizer.apply_chat_template(messages[:index], tokenize=True, add_generation_prompt=True)
        end = tokenizer.apply_chat_template(messages[:index + 1], tokenize=True, add_generation_prompt=False)
        if ids[:len(prefix)] != prefix or ids[:len(end)] != end:
            raise ValueError("Chat template does not preserve assistant boundaries.")
        labels[len(prefix):len(end)] = ids[len(prefix):len(end)]
    if not any(x != -100 for x in labels[1:]):
        raise ValueError("No assistant target tokens.")
    return {"input_ids": ids, "attention_mask": [1] * len(ids), "labels": labels}


def split_for(key):
    bucket = int(key[:8], 16) % 100
    return "validation" if bucket < 2 else "test" if bucket < 4 else "train"


def prepare(args):
    from datasets import load_dataset
    from transformers import AutoTokenizer

    output = Path(args.output)
    if output.exists() and any(output.iterdir()):
        raise FileExistsError("Use an empty data output directory.")
    tokenizer = AutoTokenizer.from_pretrained(
        CONFIG["base_model"], revision=CONFIG["base_revision"], trust_remote_code=False)
    result = {name: [] for name in ["train", "validation", "test"]}
    seen = set()
    stats = {}
    for source in CONFIG["sources"]:
        targets = {
            "train": round(args.train * source["weight"]),
            "validation": args.holdout_per_source,
            "test": args.holdout_per_source,
        }
        counts, skipped = Counter(), Counter()
        stream = load_dataset(
            CONFIG["dataset"], source["config"], split="train",
            revision=CONFIG["dataset_revision"], streaming=True)
        stream = stream.shuffle(seed=CONFIG["seed"], buffer_size=4096)
        for row in stream:
            try:
                messages = normalize_messages(row)
                key = prompt_key(messages)
                if key in seen:
                    skipped["duplicate_prompt"] += 1
                    continue
                encode_example(tokenizer, messages, CONFIG["max_sequence_length"])
            except (ValueError, TypeError, KeyError):
                skipped["invalid_or_overlength"] += 1
                continue
            seen.add(key)
            split = split_for(key)
            if counts[split] >= targets[split]:
                continue
            content_hash = hashlib.sha256(json.dumps(messages, ensure_ascii=False, sort_keys=True).encode()).hexdigest()
            result[split].append({"id": content_hash, "prompt_hash": key,
                                  "source": source["config"], "messages": messages})
            counts[split] += 1
            if all(counts[s] >= targets[s] for s in targets):
                break
        if any(counts[s] < targets[s] for s in targets):
            raise RuntimeError(f"Not enough usable data for {source['config']}: {dict(counts)}")
        stats[source["config"]] = {"selected": dict(counts), "skipped": dict(skipped)}
        print(source["config"], dict(counts), flush=True)
    behavior_path = Path(__file__).with_name("article_examples.json")
    for split, rows in behavior_examples().items():
        for row in rows:
            encode_example(tokenizer, row["messages"], CONFIG["max_sequence_length"])
            key = prompt_key(row["messages"])
            if key in seen:
                raise ValueError("Article behavior prompt duplicates a selected Hub prompt.")
            seen.add(key)
            row.update({"id": hashlib.sha256(json.dumps(row["messages"], sort_keys=True).encode()).hexdigest(),
                        "prompt_hash": key})
            result[split].append(row)
    stats["authored-article-behavior"] = {"selected": {k: len(v) for k, v in behavior_examples().items()},
                                          "license": "GPL-3.0-only", "fictional": True}
    output.mkdir(parents=True, exist_ok=True)
    hashes = {}
    for split, rows in result.items():
        random.Random(CONFIG["seed"]).shuffle(rows)
        content = "".join(json.dumps(row, ensure_ascii=False) + "\n" for row in rows)
        (output / f"{split}.jsonl").write_bytes(content.encode("utf-8"))
        hashes[split] = hashlib.sha256(content.encode()).hexdigest()
    manifest = {"config": CONFIG, "counts": {k: len(v) for k, v in result.items()},
                "article_examples_sha256": hashlib.sha256(behavior_path.read_bytes()).hexdigest(),
                "sources": stats, "sha256": hashes,
                "split_policy": "Hub: hash of normalized first user question. Article behavior: source-packet groups. Official source test splits unused.",
                "limitations": ["Synthetic data may contain mistakes.",
                    "Exact prompt deduplication is not semantic or repository-level decontamination.",
                    "The pretrained base may already have seen these data.",
                    "This pilot's code source focuses on Python; other languages need additional data and evaluation."]}
    (output / "manifest.json").write_text(json.dumps(manifest, indent=2), encoding="utf-8")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Prepare Chat & Code data on the cloud machine")
    parser.add_argument("--output", default="chatcode-data")
    parser.add_argument("--train", type=int, default=CONFIG["train_examples"])
    parser.add_argument("--holdout-per-source", type=int, default=128)
    args = parser.parse_args()
    if args.train < 10 or args.train % 10 or args.holdout_per_source < 1:
        parser.error("--train must be a positive multiple of 10; holdouts must be positive.")
    prepare(args)