"""Fase 2: valida data/train.jsonl y data/eval.jsonl contra el tokenizer/chat_template REAL
del modelo de produccion.

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. Filtra (no trunca) ejemplos que excedan MAX_TOKENS. 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 train.jsonl/eval.jsonl como ultima linea
de defensa. Reporta un resumen final por bucket.
"""
import json
import re
import sys
from pathlib import Path

from transformers import AutoTokenizer

MODEL_PATH = sys.argv[1] if len(sys.argv) > 1 else "/workspace/ft-models/Qwen--Qwen3.6-35B-A3B"
REPO_ROOT = Path(__file__).resolve().parent.parent
TRAIN_CHAT_TEMPLATE_PATH = REPO_ROOT / "data" / "chat_template_train.jinja"
MAX_TOKENS = 8192

DATASET_FILES = [
    REPO_ROOT / "data" / "train.jsonl",
    REPO_ROOT / "data" / "eval.jsonl",
]

# 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 validate_example(tokenizer, example):
    messages = example["messages"]
    tools = strip_tools_for_check(example.get("tools"))

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

    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")

    summary = {}
    total_exceptions = 0
    total_filtered_length = 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, "filtered_length": 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

            if n_tokens > MAX_TOKENS:
                stats["filtered_length"] += 1
                total_filtered_length += 1
                print(f"[FILTERED] {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 train.jsonl/eval.jsonl")
    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']} filtrados_por_longitud={stats['filtered_length']} excepciones={stats['exceptions']}")

    print(f"\n=== TOTAL: ok={total_ok} filtrados_por_longitud={total_filtered_length} excepciones={total_exceptions} ===")

    if secret_problems:
        print(f"\n[GATE FAIL] {len(secret_problems)} secretos sobrevivientes en train/eval:")
        for problem in secret_problems:
            print(f"  - {problem}")
    else:
        print("\n[GATE OK] 0 secretos sobrevivientes en train.jsonl/eval.jsonl")

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

    sys.exit(0)


if __name__ == "__main__":
    main()