"""Fase 2 (revisado en Fase 6): valida los jsonl de dataset (por defecto data/train.jsonl y
data/eval.jsonl) contra el tokenizer/chat_template REAL del modelo de produccion.

Nota sobre emails de contacto en el copy de los seeds: `example.com` ya esta en
SAFE_EMAIL_DOMAINS, asi que una direccion como `reservas@example.com` en el texto de una
landing pasa el gate de secretos sin cambios. Cualquier otro dominio lo hace fallar.

Corre DENTRO del contenedor `qwen-lora-train` en spark (necesita `transformers` con el
chat_template.jinja real de Qwen3.6, no una version instalada localmente). Invocar via:

    docker exec qwen-lora-train python3 /workspace/ai-projects/qwen3-6-lora/scripts/06_validate_dataset.py

**Hallazgo de esta fase**: el chat_template.jinja real de produccion (7764 bytes, confirmado
identico al de Fase 0) NO tiene tags `{% generation %}/{% endgeneration %}` -- por diseno,
sirve solo para inferencia, no para masking de loss de entrenamiento. Con
`return_assistant_tokens_mask=True` sobre ese template, la mascara sale vacia para el 100%
de los ejemplos (excepcion real encontrada al correr este script por primera vez). Este es
exactamente el escenario de fallback anticipado en la Decision de diseno #4 del plan
principal ("si TRL no aplica el masking nativo, copiar el .jinja con
{% generation %}...{% endgeneration %} manual"). Se genero `data/chat_template_train.jinja`
-- copia exacta del template de produccion, con `{%- generation -%}` envolviendo unicamente
el contenido/tool_calls/<|im_end|> de cada turno assistant (nunca el texto de system/user/tool)
-- verificado que el texto renderizado es byte-identico al original (los tags de generation
no emiten caracteres, solo delimitan offsets para la mascara). Este script usa ese template
SOLO para la validacion/masking; el `chat_template.jinja` original (sin tags) es el que se
usa en inferencia/produccion y no se toca.

Por cada ejemplo: tokenizer.apply_chat_template(messages, tools=..., tokenize=True,
return_assistant_tokens_mask=True, return_dict=True) -- assert sin excepcion, mascara de
assistant no vacia. Re-corre un gate de secretos (regex explicitas, igual que 04_sanitize.py,
sin depender de detect-secrets -- puede no estar instalado en este contenedor) sobre los
archivos validados como ultima linea de defensa. Reporta un resumen final por bucket con el
histograma de longitudes (p50/p90/p99/max).

PARAMETROS POR ENV
------------------
    VALIDATE_FILES     lista de jsonl separados por coma  (def: data/train.jsonl,data/eval.jsonl)
    MAX_TOKENS         tope duro de longitud              (def: 8192)
    PRESERVE_THINKING  '1' para conservar el thinking de turnos previos (def: off)
    MODEL_PATH         ruta del tokenizer (o argv[1])

Las rutas relativas se resuelven contra la raiz del repo. Ejemplo de Fase 6:

    VALIDATE_FILES=data/train_lora2.jsonl,data/eval_lora2.jsonl \
    MAX_TOKENS=3000 PRESERVE_THINKING=1 \
    python3 scripts/06_validate_dataset.py

**El tope de longitud ABORTA, no filtra.** Hasta la Fase 6 este script imprimia `[FILTERED]`
e incrementaba un contador, pero la fila se quedaba en el archivo y `10_train.py` (que corre
sin `max_seq_length` y con batch 1) la entrenaba igual: era un falso sentido de seguridad que
podia reventar el presupuesto de memoria a mitad de corrida, horas adentro. Ahora un solo
ejemplo por encima de MAX_TOKENS termina en `sys.exit(1)`.

**PRESERVE_THINKING tiene que coincidir con el del training.** El template condiciona el
thinking de los turnos previos a `preserve_thinking`; si este script valida con el flag
apagado y el training entrena con el prendido, se esta midiendo otra cosa (otra longitud y
otra mascara).
"""
import json
import os
import re
import sys
from pathlib import Path

from transformers import AutoTokenizer

REPO_ROOT = Path(__file__).resolve().parent.parent


def _resolve(raw):
    """Las rutas relativas cuelgan de la raiz del repo."""
    p = Path(raw.strip())
    return p if p.is_absolute() else REPO_ROOT / p


MODEL_PATH = (
    sys.argv[1] if len(sys.argv) > 1
    else os.environ.get("MODEL_PATH", "/workspace/ft-models/Qwen--Qwen3.6-35B-A3B")
)
TRAIN_CHAT_TEMPLATE_PATH = REPO_ROOT / "data" / "chat_template_train.jinja"
MAX_TOKENS = int(os.environ.get("MAX_TOKENS", "8192"))
PRESERVE_THINKING = os.environ.get("PRESERVE_THINKING", "").lower() in ("1", "true", "yes")

DATASET_FILES = [
    _resolve(part)
    for part in os.environ.get("VALIDATE_FILES", "data/train.jsonl,data/eval.jsonl").split(",")
    if part.strip()
]

# Mismos patrones explicitos que 04_sanitize.py (subset sin dependencia de detect-secrets,
# que puede no estar instalado en este contenedor de training) -- ultima linea de defensa
# sobre la salida YA sanitizada y ensamblada. El local-part exige 2+ caracteres (no 1+)
# para no matchear falsos positivos de codigo como "\n@app.route" (decorador Flask en un
# seed) leido como si "n" fuera el local-part de un email.
EMAIL_RE = re.compile(r"\b[A-Za-z0-9._%+-]{2,}@[A-Za-z0-9.-]+\.[A-Za-z]{2,}\b")

# Dominios de ejemplo/placeholder de uso convencional en contenido sintetico de
# entrenamiento (RFC 2606 reserva example.com/.org/.net exactamente para esto) -- un match
# de EMAIL en uno de estos dominios no es un secreto real, es contenido de ejemplo
# intencional (direcciones de Jira/Confluence ficticias, snippets de validacion de email,
# etc.), asi que no debe hacer fallar el gate.
SAFE_EMAIL_DOMAINS = {"example.com", "example.org", "example.net", "email.com", "ejemplo.com", "test.com", "anthropic.com"}

EXPLICIT_PATTERNS = [
    ("SPARK_PASSWORD", re.compile(r"\b01140102Alb\?")),
    ("BEARER_TOKEN", re.compile(r"\b7c1f76a62391a47941d7aab8369eb8f20334daf136ba88080815ff3070773a1f\b")),
    ("DB_PASSWORD", re.compile(r"\bsarh21234\b")),
    ("SPARK_IP", re.compile(r"\b10\.212\.133\.200\b")),
    ("INTERNAL_IP", re.compile(r"\b10\.212\.133\.\d{1,3}\b")),
    ("INTERNAL_SUBNET", re.compile(r"\b10\.212\.133\.0/24\b")),
]


def find_unsafe_emails(line):
    unsafe = []
    for match in EMAIL_RE.finditer(line):
        email = match.group(0)
        domain = email.split("@", 1)[1].lower()
        if domain not in SAFE_EMAIL_DOMAINS:
            unsafe.append(email)
    return unsafe


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


def secrets_gate(path):
    problems = []
    with open(path, encoding="utf-8") as f:
        for lineno, line in enumerate(f, start=1):
            for semantic_name, pattern in EXPLICIT_PATTERNS:
                if pattern.search(line):
                    problems.append(f"{path.name}:{lineno}: patron '{semantic_name}' sobrevivio")
            for email in find_unsafe_emails(line):
                problems.append(f"{path.name}:{lineno}: email fuera de dominios placeholder conocidos: {email}")
    return problems


def strip_tools_for_check(tools):
    if not tools:
        return None
    return tools


def percentile(sorted_values, q):
    """Percentil por el metodo del vecino mas cercano (sin numpy: este contenedor puede no
    tenerlo y no vale la pena una dependencia por esto)."""
    if not sorted_values:
        return 0
    idx = min(len(sorted_values) - 1, max(0, int(round(q * (len(sorted_values) - 1)))))
    return sorted_values[idx]


def validate_example(tokenizer, example):
    messages = example["messages"]
    tools = strip_tools_for_check(example.get("tools"))

    template_kwargs = {}
    if PRESERVE_THINKING:
        template_kwargs["preserve_thinking"] = True

    rendered = tokenizer.apply_chat_template(
        messages,
        tools=tools,
        tokenize=True,
        return_assistant_tokens_mask=True,
        return_dict=True,
        add_generation_prompt=False,
        **template_kwargs,
    )

    input_ids = rendered["input_ids"]
    assistant_masks = rendered.get("assistant_masks")

    n_tokens = len(input_ids)
    mask_sum = sum(assistant_masks) if assistant_masks is not None else 0

    if assistant_masks is None:
        raise AssertionError("apply_chat_template no devolvio assistant_masks")
    if mask_sum == 0:
        raise AssertionError("assistant_masks esta vacia (0 tokens de assistant marcados)")

    return n_tokens, mask_sum


def main():
    print(f"[INFO] cargando tokenizer real desde {MODEL_PATH}")
    tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)

    print(f"[INFO] reemplazando chat_template por la variante de training con masking: {TRAIN_CHAT_TEMPLATE_PATH}")
    tokenizer.chat_template = TRAIN_CHAT_TEMPLATE_PATH.read_text(encoding="utf-8")

    print(f"[INFO] archivos: {', '.join(str(p) for p in DATASET_FILES)}")
    print(f"[INFO] MAX_TOKENS={MAX_TOKENS} (aborta, no filtra)  PRESERVE_THINKING={PRESERVE_THINKING}")

    summary = {}
    lengths_by_bucket = {}
    too_long = []
    total_exceptions = 0
    total_ok = 0

    for path in DATASET_FILES:
        if not path.exists():
            print(f"[ERROR] {path} no existe")
            sys.exit(1)

        examples = load_jsonl(path)
        print(f"[INFO] {path.name}: {len(examples)} ejemplos")

        for lineno, example in examples:
            bucket = example.get("meta", {}).get("bucket", "sin_bucket")
            stats = summary.setdefault(bucket, {"ok": 0, "exceptions": 0, "too_long": 0})

            try:
                n_tokens, mask_sum = validate_example(tokenizer, example)
            except Exception as e:
                stats["exceptions"] += 1
                total_exceptions += 1
                print(f"[EXCEPTION] {path.name}:{lineno} (bucket={bucket}): {e}")
                continue

            lengths_by_bucket.setdefault(bucket, []).append(n_tokens)

            if n_tokens > MAX_TOKENS:
                stats["too_long"] += 1
                too_long.append((path.name, lineno, bucket, n_tokens))
                print(f"[TOO LONG] {path.name}:{lineno} (bucket={bucket}): {n_tokens} tokens > {MAX_TOKENS}")
                continue

            stats["ok"] += 1
            total_ok += 1

    print("\n[INFO] re-corriendo gate de secretos sobre los archivos validados")
    secret_problems = []
    for path in DATASET_FILES:
        secret_problems.extend(secrets_gate(path))

    print("\n=== RESUMEN POR BUCKET ===")
    for bucket, stats in sorted(summary.items()):
        print(f"  {bucket}: ok={stats['ok']} sobre_max_tokens={stats['too_long']} excepciones={stats['exceptions']}")

    # Histograma de longitudes: es lo que permite decidir si MAX_TOKENS esta bien puesto ANTES
    # de quemar una corrida de entrenamiento. Un p99 pegado al tope significa que el proximo
    # seed largo aborta la build; un maximo muy por debajo significa que se puede bajar el tope
    # (y con el, el pico de memoria).
    print("\n=== HISTOGRAMA DE TOKENS POR BUCKET (p50 / p90 / p99 / max) ===")
    all_lengths = []
    for bucket, lengths in sorted(lengths_by_bucket.items()):
        ordered = sorted(lengths)
        all_lengths.extend(ordered)
        print(
            f"  {bucket:24} n={len(ordered):5}  p50={percentile(ordered, 0.50):6}  "
            f"p90={percentile(ordered, 0.90):6}  p99={percentile(ordered, 0.99):6}  "
            f"max={ordered[-1]:6}"
        )
    if all_lengths:
        ordered = sorted(all_lengths)
        print(
            f"  {'TOTAL':24} n={len(ordered):5}  p50={percentile(ordered, 0.50):6}  "
            f"p90={percentile(ordered, 0.90):6}  p99={percentile(ordered, 0.99):6}  "
            f"max={ordered[-1]:6}   (MAX_TOKENS={MAX_TOKENS})"
        )

    print(f"\n=== TOTAL: ok={total_ok} sobre_max_tokens={len(too_long)} excepciones={total_exceptions} ===")

    if secret_problems:
        print(f"\n[GATE FAIL] {len(secret_problems)} secretos sobrevivientes:")
        for problem in secret_problems:
            print(f"  - {problem}")
    else:
        print("\n[GATE OK] 0 secretos sobrevivientes en los archivos validados")

    if too_long:
        print(f"\n[GATE FAIL] {len(too_long)} ejemplo(s) superan MAX_TOKENS={MAX_TOKENS}:")
        for fname, lineno, bucket, n_tokens in too_long:
            print(f"  - {fname}:{lineno} (bucket={bucket}): {n_tokens} tokens")
        print(
            "  Estas filas NO se filtran solas: 10_train.py corre sin max_seq_length y con\n"
            "  batch 1, asi que las entrenaria enteras y podria reventar el presupuesto de\n"
            "  memoria a mitad de corrida. Hay que arreglar el dataset, no el validador.\n"
            "  Como arreglarlo: partir la trayectoria larga en DOS ejemplos, persistiendo los\n"
            "  ids que el segundo necesita en `storage` (penpot.local_storage) al final del\n"
            "  primero y leyendolos al principio del segundo. Ademas de bajar la longitud, es\n"
            "  mejor pedagogia: es exactamente lo que la descripcion de execute_code pide."
        )

    if total_exceptions > 0 or secret_problems or too_long:
        sys.exit(1)

    sys.exit(0)


if __name__ == "__main__":
    main()