"""Fase 4 -- suite de evaluacion en 4 puertas.

Puerta 1 (--gate 1): eval-loss offline por bucket sobre el checkpoint MERGEADO
(no el adapter puro) -- no necesita servir el modelo. Corre DENTRO del
contenedor qwen-lora-train en spark:

    docker exec qwen-lora-train python3 \
        /workspace/ai-projects/qwen3-6-lora/.worktrees/agente-fase4-merge-eval/scripts/30_eval_suite.py --gate 1

Carga el checkpoint mergeado con AutoModelForCausalLM (para detectar bugs de
merge que un eval sobre el adapter puro no veria), le pisa en memoria el
chat_template con data/chat_template_train.jinja (igual que en training, para
poder generar assistant_masks), recorre EVAL_FILE (default: data/eval.jsonl)
agrupado por meta.bucket, y reporta loss promedio global y por bucket (aislando
bucket=="replay"), comparado contra BASELINE_EVAL_LOSS (default: el
eval_loss=0.2750 de Fase 3). Rutas y baselines son configurables por env var --
ver el bloque de constantes.

Las puertas 2-4 (tool-calls, adherencia, E2E) viven en scripts separados
(scripts/31_gate2_toolcalls.py, scripts/32_gate3_adherencia.py,
scripts/33_gate4_e2e.py) porque necesitan el contenedor de eval sirviendo el
checkpoint mergeado via HTTP, no solo lectura offline.

Fase 6: EVAL_FILE, BASELINE_EVAL_LOSS y BASELINE_BUCKET_LOSSES son env vars. Los
defaults son el eval.jsonl CONGELADO y el eval_loss=0.2750, ambos de Fase 3 -- de
modo que una corrida pelada reproduce exactamente la medicion de la puerta 1 y la
unica variable entre baseline y candidato sea el propio LoRA #2. El reporte imprime
la perdida ponderada global y el desglose por bucket con su delta contra el baseline,
y un veredicto PASS/FAIL explicito por cada uno de los dos umbrales:
  - MAX_GLOBAL_WEIGHTED_LOSS (default 0.35): perdida ponderada global.
  - MAX_BUCKET_REGRESSION (default 0.10): ningun bucket NO-penpot puede estar peor
    que SU PROPIO baseline por mas de ese margen.

PRINCIPIO RECTOR de todos los chequeos de este script: ausente no es OK; ausente es
"no se pudo verificar", y eso tiene que fallar o reportarse como SKIP explicito,
nunca colarse como verde. En concreto: un bucket sin baseline propio NO se compara
contra el numero global (eso produce falsos PASS y falsos FAIL por igual) -- se marca
SKIP y el veredicto reporta cuantos buckets quedaron sin verificar.

El script sale con codigo 1 si la puerta no pasa: el runbook la encadena con la
cuantizacion, y una puerta que no puede fallar no es una puerta.
"""
import argparse
import json
import os
import sys
import time
from collections import defaultdict
from pathlib import Path

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

REPO_ROOT = Path(__file__).resolve().parent.parent
OUTPUT_PATH = os.environ.get("OUTPUT_PATH", "/workspace/ft-models/Qwen3.6-35B-A3B-mcp-bf16")
TRAIN_CHAT_TEMPLATE_PATH = REPO_ROOT / "data" / "chat_template_train.jinja"
# Por env para poder apuntar a data/eval_lora2.jsonl sin tocar el script, pero el
# default sigue siendo el eval.jsonl CONGELADO de Fase 3: es la procedencia exacta
# del modelo en produccion y el unico archivo contra el que 0.2750 significa algo.
EVAL_FILE = Path(os.environ.get("EVAL_FILE", str(REPO_ROOT / "data" / "eval.jsonl")))
# eval_loss de Fase 3 (adapter puro, transformers.Trainer, ponderado por token).
BASELINE_EVAL_LOSS = float(os.environ.get("BASELINE_EVAL_LOSS", "0.2750"))
# Baselines por bucket (JSON: {"bucket": loss, ...}). El Trainer solo reporta el
# numero global, asi que los por-bucket hay que pasarlos a mano desde la corrida
# anterior de esta misma puerta.
#
# Un bucket SIN entrada aca NO se compara contra BASELINE_EVAL_LOSS: el numero
# global no es el baseline de ningun bucket en particular, y usarlo como tal produce
# tanto falsos PASS (un bucket que valia 0.12 y ahora vale 0.33 queda dentro del
# margen contra 0.2750) como falsos FAIL (un bucket que siempre valio 0.42 y sigue
# igual "regresiona" +0.145). Ausente no es OK; ausente es "no se pudo verificar":
# la fila se marca SKIP y el veredicto reporta cuantos buckets quedaron sin cubrir.


def _parse_baseline_bucket_losses():
    """Parsea BASELINE_BUCKET_LOSSES. Se llama en tiempo de import, asi que un JSON
    mal formado tiene que salir con un mensaje claro (y no con un traceback crudo
    hasta corriendo --help)."""
    raw = os.environ.get("BASELINE_BUCKET_LOSSES", "{}")
    try:
        parsed = json.loads(raw)
    except json.JSONDecodeError as exc:
        raise SystemExit(
            f"[ERROR] BASELINE_BUCKET_LOSSES no es JSON valido ({exc}). "
            f'Valor recibido: {raw!r}. Formato esperado: {{"bucket": 0.1234, ...}}'
        )
    if not isinstance(parsed, dict):
        raise SystemExit(
            f"[ERROR] BASELINE_BUCKET_LOSSES debe ser un objeto JSON, se recibio {type(parsed).__name__}: {raw!r}"
        )
    out = {}
    for bucket, value in parsed.items():
        if not isinstance(value, (int, float)) or isinstance(value, bool):
            raise SystemExit(
                f"[ERROR] BASELINE_BUCKET_LOSSES[{bucket!r}] = {value!r} no es un numero"
            )
        out[bucket] = float(value)
    return out


BASELINE_BUCKET_LOSSES = _parse_baseline_bucket_losses()
# Umbrales de la puerta de olvido (Fase 6).
MAX_GLOBAL_WEIGHTED_LOSS = float(os.environ.get("MAX_GLOBAL_WEIGHTED_LOSS", "0.35"))
MAX_BUCKET_REGRESSION = float(os.environ.get("MAX_BUCKET_REGRESSION", "0.10"))
# Bucket exento del umbral de regresion por bucket: es donde el LoRA #2 debe moverse.
# OJO con lo que esta exencion NO significa: en data/eval.jsonl las 11 filas de este
# bucket son PREEXISTENTES del LoRA #1 (tool-calling del MCP de Penpot, ensenado en
# Fase 3), no capacidad nueva. Exentarlas cumple el plan al pie de la letra, pero deja
# a la puerta 1 sin cobertura de regresion sobre la capacidad Penpot VIEJA -- que es
# justo la que el LoRA #2 mas va a pisar. Si se quiere esa cobertura, hay que dejar el
# bucket bajo umbral (PENPOT_BUCKET="") y separar la capacidad nueva en otro archivo
# de eval.
PENPOT_BUCKET = os.environ.get("PENPOT_BUCKET", "penpot")


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--gate", type=int, required=True, choices=[1])
    return parser.parse_args()


def load_eval_examples():
    examples = []
    with open(EVAL_FILE, encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            examples.append(json.loads(line))
    return examples


def compute_loss_per_example(model, tokenizer, example):
    rendered = tokenizer.apply_chat_template(
        example["messages"],
        tools=example.get("tools"),
        tokenize=True,
        return_assistant_tokens_mask=True,
        return_dict=True,
        add_generation_prompt=False,
    )
    input_ids = rendered["input_ids"]
    assistant_masks = rendered["assistant_masks"]
    if sum(assistant_masks) == 0:
        raise AssertionError("assistant_masks vacia para un ejemplo de eval.jsonl")
    labels = [tok if mask == 1 else -100 for tok, mask in zip(input_ids, assistant_masks)]

    input_ids_t = torch.tensor([input_ids], dtype=torch.long, device=model.device)
    labels_t = torch.tensor([labels], dtype=torch.long, device=model.device)
    with torch.no_grad():
        out = model(input_ids=input_ids_t, labels=labels_t)
    n_assistant_tokens = sum(assistant_masks)
    return out.loss.item(), n_assistant_tokens


def run_gate1():
    print(f"[INFO] cargando checkpoint mergeado desde {OUTPUT_PATH}")
    tokenizer = AutoTokenizer.from_pretrained(OUTPUT_PATH)
    tokenizer.chat_template = TRAIN_CHAT_TEMPLATE_PATH.read_text(encoding="utf-8")
    if tokenizer.pad_token_id is None:
        tokenizer.pad_token = tokenizer.eos_token

    t0 = time.time()
    model = AutoModelForCausalLM.from_pretrained(
        OUTPUT_PATH,
        dtype=torch.bfloat16,
        attn_implementation="flash_attention_2",
    )
    model = model.to("cuda")
    model.eval()
    load_time = time.time() - t0
    print(f"[INFO] modelo cargado en {load_time:.1f}s")

    examples = load_eval_examples()
    print(f"[INFO] {len(examples)} ejemplos en {EVAL_FILE}")

    torch.cuda.reset_peak_memory_stats()
    t0 = time.time()
    # Cada entrada es (loss_del_ejemplo, n_tokens_assistant_del_ejemplo) -- se necesitan
    # ambos para poder reportar tanto el promedio simple por ejemplo (util para comparar
    # buckets entre si) como el promedio ponderado por token (comparable directamente
    # contra el eval_loss que reporta transformers.Trainer, que pondera por cantidad de
    # tokens validos y no por cantidad de ejemplos -- un bucket con pocos ejemplos pero
    # secuencias largas/dificiles no debe pesar igual que uno con muchos ejemplos cortos).
    losses_by_bucket = defaultdict(list)
    for i, example in enumerate(examples):
        bucket = example.get("meta", {}).get("bucket", "sin_bucket")
        loss, n_tokens = compute_loss_per_example(model, tokenizer, example)
        losses_by_bucket[bucket].append((loss, n_tokens))
        if (i + 1) % 25 == 0:
            print(f"[INFO] {i + 1}/{len(examples)} ejemplos evaluados")
    eval_time = time.time() - t0
    peak_mem_gb = torch.cuda.max_memory_allocated() / (1024 ** 3)

    def weighted_avg(pairs):
        total_tokens = sum(n for _, n in pairs)
        return sum(loss * n for loss, n in pairs) / total_tokens

    def simple_avg(pairs):
        return sum(loss for loss, _ in pairs) / len(pairs)

    all_pairs = [pair for pairs in losses_by_bucket.values() for pair in pairs]
    global_avg_simple = simple_avg(all_pairs)
    global_avg_weighted = weighted_avg(all_pairs)

    print("\n=== Puerta 1 -- eval-loss offline por bucket (checkpoint mergeado) ===")
    print(f"[INFO] EVAL_FILE={EVAL_FILE}")
    print(f"[INFO] BASELINE_EVAL_LOSS={BASELINE_EVAL_LOSS:.4f}")
    print(f"[INFO] tiempo de eval: {eval_time:.1f}s, memoria pico: {peak_mem_gb:.2f} GB")

    # Toda clave de BASELINE_BUCKET_LOSSES tiene que existir entre los buckets
    # encontrados: un typo ("otros_mcp" por "otros_mcps") no matchea nada, no dice
    # nada, y deja ese bucket sin verificar mientras el operador cree que lo cubrio.
    buckets_encontrados = set(losses_by_bucket)
    baselines_sobrantes = sorted(set(BASELINE_BUCKET_LOSSES) - buckets_encontrados)
    if baselines_sobrantes:
        print(
            f"[ERROR] BASELINE_BUCKET_LOSSES tiene claves que no existen en {EVAL_FILE}: "
            f"{baselines_sobrantes} (buckets encontrados: {sorted(buckets_encontrados)}). "
            "Probablemente un typo -- ese baseline no se estaria aplicando a nada."
        )
        raise SystemExit(1)

    # Desglose por bucket con el delta contra SU PROPIO baseline. El delta es lo que
    # decide la puerta de olvido: un promedio global sano puede esconder un bucket
    # no-penpot que se degrado, compensado por la mejora del bucket penpot.
    print(
        f"\n  {'bucket':22s} {'n':>4s} {'tokens':>7s} {'simple':>8s} {'ponderado':>10s} "
        f"{'baseline':>9s} {'delta':>8s}"
    )
    bucket_regressions = []
    buckets_sin_baseline = []
    for bucket in sorted(losses_by_bucket):
        pairs = losses_by_bucket[bucket]
        n_tokens_total = sum(n for _, n in pairs)
        w = weighted_avg(pairs)
        # Nunca se inventa un baseline: si no hay uno propio para este bucket, el
        # umbral simplemente NO se evalua y la fila se marca SKIP.
        baseline = BASELINE_BUCKET_LOSSES.get(bucket)
        es_penpot = bucket == PENPOT_BUCKET
        if es_penpot:
            nota = " (exento: tool-calling Penpot PREEXISTENTE del LoRA #1, sin cobertura de regresion aca)"
        elif baseline is None:
            nota = " SKIP (sin baseline propio -- no se pudo verificar)"
            buckets_sin_baseline.append(bucket)
        else:
            nota = ""
        baseline_txt = f"{baseline:9.4f}" if baseline is not None else f"{'n/d':>9s}"
        delta_txt = f"{w - baseline:+8.4f}" if baseline is not None else f"{'--':>8s}"
        print(
            f"  {bucket:22s} {len(pairs):4d} {n_tokens_total:7d} {simple_avg(pairs):8.4f} "
            f"{w:10.4f} {baseline_txt} {delta_txt}{nota}"
        )
        if not es_penpot and baseline is not None and (w - baseline) > MAX_BUCKET_REGRESSION:
            bucket_regressions.append((bucket, w, baseline, w - baseline))

    replay_pairs = losses_by_bucket.get("replay")
    if replay_pairs:
        print(
            f"\n  bucket=replay (aislado)  n={len(replay_pairs):4d} "
            f"loss_avg_simple={simple_avg(replay_pairs):.4f} loss_avg_weighted={weighted_avg(replay_pairs):.4f}"
        )

    print(f"\n  loss_avg GLOBAL simple (por ejemplo)     = {global_avg_simple:.4f}")
    print(f"  loss_avg GLOBAL ponderado (por token)     = {global_avg_weighted:.4f}")
    print(f"  baseline (ponderado por token)           = {BASELINE_EVAL_LOSS:.4f}")
    diff = abs(global_avg_weighted - BASELINE_EVAL_LOSS)
    print(f"  diferencia absoluta (ponderado vs baseline) = {diff:.4f}")
    if diff > 0.05:
        print(
            "  [WARN] diferencia > 0.05 -- senal posible de bug real en el merge, "
            "revisar antes de continuar a la puerta 2"
        )
    else:
        print("  [OK] loss del checkpoint mergeado consistente con el baseline -- merge probablemente correcto")

    # Veredicto explicito, uno por umbral. Son dos condiciones independientes: la
    # puerta pasa solo si las dos pasan.
    print("\n=== Veredicto de la puerta de olvido ===")
    global_pass = global_avg_weighted <= MAX_GLOBAL_WEIGHTED_LOSS
    print(
        f"  [{'PASS' if global_pass else 'FAIL'}] perdida ponderada global "
        f"{global_avg_weighted:.4f} <= {MAX_GLOBAL_WEIGHTED_LOSS:.4f} (MAX_GLOBAL_WEIGHTED_LOSS)"
    )
    buckets_pass = not bucket_regressions
    # Cuantos buckets pasaron REALMENTE por el umbral: ni los exentos ni los que
    # quedaron en SKIP cuentan como verificados.
    n_verificados = len(losses_by_bucket) - len(buckets_sin_baseline) - (1 if PENPOT_BUCKET in losses_by_bucket else 0)
    print(
        f"  [{'PASS' if buckets_pass else 'FAIL'}] ningun bucket no-{PENPOT_BUCKET} peor que su "
        f"baseline por mas de {MAX_BUCKET_REGRESSION:.4f} (MAX_BUCKET_REGRESSION) "
        f"-- umbral evaluado sobre {n_verificados}/{len(losses_by_bucket)} buckets"
    )
    for bucket, w, baseline, delta in bucket_regressions:
        print(f"      bucket={bucket}: {w:.4f} vs baseline {baseline:.4f} (delta {delta:+.4f})")
    if buckets_sin_baseline:
        print(
            f"  [SKIP] {len(buckets_sin_baseline)} bucket(s) SIN VERIFICAR por falta de baseline "
            f"propio en BASELINE_BUCKET_LOSSES: {buckets_sin_baseline}"
        )
        print(
            "         Ausente no es OK: es 'no se pudo verificar'. Estos buckets NO pasaron por "
            "el umbral de regresion -- pasarles su baseline medido para que la puerta los cubra."
        )

    todo_pasa = global_pass and buckets_pass
    sufijo = f" (con {len(buckets_sin_baseline)} bucket(s) SIN VERIFICAR)" if buckets_sin_baseline else ""
    print(f"\n  VEREDICTO PUERTA 1: {'PASS' if todo_pasa else 'FAIL'}{sufijo}")
    return todo_pasa


def main():
    args = parse_args()
    if args.gate == 1:
        # Codigo de salida != 0 cuando la puerta no pasa: el runbook la encadena con
        # la cuantizacion, y una puerta que no puede fallar automaticamente no es una
        # puerta.
        sys.exit(0 if run_gate1() else 1)


if __name__ == "__main__":
    main()