"""Fase 2: ensambla data/train.jsonl y data/eval.jsonl (v1 a volumen reducido) a partir de
los seeds escritos a mano en data/raw/seeds/*.jsonl mas data/raw/replay.jsonl.

Corre localmente (no requiere GPU ni el modelo). Aplica variacion deterministica
(random.seed(42)) sobre los seeds -- sustitucion de valores (colores, ids, numeros) +
un set chico de parafraseos por bucket -- para llegar al volumen v1 objetivo por bucket
sin duplicar lineas exactas. Verifica que ningun seed mencione la skill held-out
(spark-ssh). Tagea cada ejemplo con meta.bucket, mezcla, y separa 90/10 train/eval
estratificado por bucket para que eval no quede dominado por un solo bucket.
"""
import json
import random
import re
import sys
from pathlib import Path

REPO_ROOT = Path(__file__).resolve().parent.parent
SEEDS_DIR = REPO_ROOT / "data" / "raw" / "seeds"
REPLAY_PATH = REPO_ROOT / "data" / "raw" / "replay.jsonl"
TRAIN_PATH = REPO_ROOT / "data" / "train.jsonl"
EVAL_PATH = REPO_ROOT / "data" / "eval.jsonl"

SEED = 42
EVAL_FRACTION = 0.10
HELD_OUT_SKILL = "spark-ssh"

# Volumen v1 objetivo por bucket (aprox. 800-1000 nuevos, ver PLAN.md seccion 3).
# El bucket "replay" no se varia -- se usa tal cual viene de la Fase 1.
BUCKET_TARGETS = {
    "penpot": 110,
    "otros_mcps": 290,
    "skills_adherencia": 160,
    "delegacion_subagentes": 65,
    "negativos": 70,
    "manejo_errores": 55,
}

HEX_COLOR_RE = re.compile(r"^#[0-9a-fA-F]{6}$")
TRAILING_NUMBER_RE = re.compile(r"^(.*?)(\d+)$")

PARAPHRASE_PREFIXES = [
    "",
    "Por favor, ",
    "Necesito que ",
    "¿Podés ",
    "Che, ",
    "Dale, ",
]

PALETTE = ["#1a73e8", "#e8710a", "#188038", "#d93025", "#9334e6", "#12b5cb", "#f9ab00", "#3c4043"]


def load_jsonl(path):
    lines = []
    with open(path, encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if line:
                lines.append(json.loads(line))
    return lines


def load_seeds():
    by_bucket = {}
    if not SEEDS_DIR.exists():
        print(f"[ERROR] {SEEDS_DIR} no existe")
        sys.exit(1)
    for path in sorted(SEEDS_DIR.glob("*.jsonl")):
        examples = load_jsonl(path)
        for ex in examples:
            bucket = ex.get("meta", {}).get("bucket")
            if bucket is None:
                print(f"[ERROR] {path}: ejemplo sin meta.bucket")
                sys.exit(1)
            by_bucket.setdefault(bucket, []).append(ex)
        print(f"[INFO] {path.name}: {len(examples)} ejemplos cargados")
    return by_bucket


def assert_no_held_out_skill(by_bucket):
    for bucket, examples in by_bucket.items():
        for ex in examples:
            blob = json.dumps(ex, ensure_ascii=False)
            assert HELD_OUT_SKILL not in blob, (
                f"[ASSERT FAIL] bucket '{bucket}' menciona la skill held-out "
                f"'{HELD_OUT_SKILL}' -- esta prohibido en todos los buckets"
            )
    print(f"[OK] ninguna mencion de la skill held-out '{HELD_OUT_SKILL}' en los seeds")


def perturb_value(value, rng):
    if isinstance(value, str):
        if HEX_COLOR_RE.match(value):
            return rng.choice(PALETTE)
        m = TRAILING_NUMBER_RE.match(value)
        if m and len(m.group(2)) <= 3:
            prefix, number = m.group(1), m.group(2)
            new_number = rng.randint(1, 99)
            return f"{prefix}{new_number}"
        return value
    if isinstance(value, bool):
        return value
    if isinstance(value, int):
        jitter = rng.randint(-max(1, value // 4), max(1, value // 4))
        return max(0, value + jitter)
    if isinstance(value, float):
        jitter = rng.uniform(-0.25, 0.25) * value
        return round(value + jitter, 2)
    if isinstance(value, dict):
        return {k: perturb_value(v, rng) for k, v in value.items()}
    if isinstance(value, list):
        return [perturb_value(v, rng) for v in value]
    return value


def vary_example(example, rng, variant_idx):
    """Produce una variante determinista del ejemplo (variant_idx=0 devuelve el
    original sin tocar, para preservar los seeds tal cual como parte del volumen)."""
    if variant_idx == 0:
        return json.loads(json.dumps(example, ensure_ascii=False))

    varied = json.loads(json.dumps(example, ensure_ascii=False))

    for msg in varied.get("messages", []):
        if msg.get("role") == "user" and variant_idx > 0:
            prefix = rng.choice(PARAPHRASE_PREFIXES)
            if prefix and not msg["content"][:1].islower():
                msg["content"] = prefix + msg["content"][0].lower() + msg["content"][1:]
        for tool_call in msg.get("tool_calls", []) or []:
            args = tool_call.get("function", {}).get("arguments")
            if isinstance(args, dict):
                tool_call["function"]["arguments"] = perturb_value(args, rng)

    return varied


def build_bucket(bucket, seeds, target_count):
    rng = random.Random(f"{SEED}-{bucket}")
    n_seeds = len(seeds)
    if n_seeds == 0:
        print(f"[WARN] bucket '{bucket}' no tiene seeds, se omite")
        return []
    if target_count <= n_seeds:
        print(f"[WARN] bucket '{bucket}': target ({target_count}) <= seeds ({n_seeds}), se usan solo los seeds originales")
        return [vary_example(ex, rng, 0) for ex in seeds]

    result = []
    variant_idx = 0
    seed_order = list(range(n_seeds))
    while len(result) < target_count:
        if variant_idx > 0:
            rng.shuffle(seed_order)
        for i in seed_order:
            if len(result) >= target_count:
                break
            result.append(vary_example(seeds[i], rng, variant_idx))
        variant_idx += 1
    print(f"[INFO] bucket '{bucket}': {n_seeds} seeds -> {len(result)} ejemplos ({variant_idx} pasadas de variacion)")
    return result


def stratified_split(all_examples, rng):
    by_bucket = {}
    for ex in all_examples:
        by_bucket.setdefault(ex["meta"]["bucket"], []).append(ex)

    train, eval_ = [], []
    for bucket, examples in by_bucket.items():
        shuffled = examples[:]
        rng.shuffle(shuffled)
        n_eval = max(1, round(len(shuffled) * EVAL_FRACTION))
        eval_.extend(shuffled[:n_eval])
        train.extend(shuffled[n_eval:])
        print(f"[INFO] split '{bucket}': {len(shuffled)} total -> train={len(shuffled) - n_eval} eval={n_eval}")

    rng.shuffle(train)
    rng.shuffle(eval_)
    return train, eval_


def write_jsonl(path, examples):
    with open(path, "w", encoding="utf-8") as f:
        for ex in examples:
            f.write(json.dumps(ex, ensure_ascii=False))
            f.write("\n")
            f.flush()


def main():
    by_bucket = load_seeds()
    assert_no_held_out_skill(by_bucket)

    all_examples = []
    for bucket, target_count in BUCKET_TARGETS.items():
        seeds = by_bucket.get(bucket, [])
        all_examples.extend(build_bucket(bucket, seeds, target_count))

    replay_examples = load_jsonl(REPLAY_PATH)
    for ex in replay_examples:
        ex.setdefault("meta", {})["bucket"] = "replay"
    print(f"[INFO] bucket 'replay': {len(replay_examples)} ejemplos (sin variacion, tal cual Fase 1)")
    all_examples.extend(replay_examples)

    print(f"[INFO] total combinado: {len(all_examples)} ejemplos")

    split_rng = random.Random(SEED)
    train, eval_ = stratified_split(all_examples, split_rng)

    write_jsonl(TRAIN_PATH, train)
    write_jsonl(EVAL_PATH, eval_)

    print(f"[OK] {TRAIN_PATH} escrito: {len(train)} ejemplos")
    print(f"[OK] {EVAL_PATH} escrito: {len(eval_)} ejemplos")


if __name__ == "__main__":
    main()