"""Fase 1: genera ~700 ejemplos de "replay" (anti-forgetting) contra el vLLM de produccion.

Pensado para correr directo en el HOST de spark (no en un contenedor) via SSH,
usando solo la stdlib (urllib.request) -- no requiere ningun paquete Python
adicional. Apunta a http://localhost:8000/v1/chat/completions (vllm-qwen36).

Cada prompt semilla se repite variando la temperatura para llegar a ~700
ejemplos con salidas genuinamente distintas (muestreo estocastico), sin
necesitar un segundo modelo generador de prompts.
"""
import json
import sys
import time
import urllib.error
import urllib.request

API_URL = "http://localhost:8000/v1/chat/completions"
MODEL = "qwen3.6-35b"
OUTPUT_PATH = "/mnt/docker-nas/projects/ai-projects/qwen3-6-lora/data/raw/replay.jsonl"
MAX_TOKENS = 2048
TEMPERATURES = [0.3, 0.6, 0.9, 1.1]
N_PASSES = 9  # 80 prompts semilla x 9 pasadas ~= 720, aprox. 700 objetivo
MAX_RETRIES = 1
TIMEOUT_S = 180

CONVERSACION_GENERAL = [
    "¿Cuál es la diferencia entre clima y tiempo atmosférico?",
    "Explica brevemente por qué el cielo es azul.",
    "¿Qué causó la caída del Imperio Romano de Occidente?",
    "Recomiéndame tres libros de ciencia ficción para alguien que nunca ha leído el género.",
    "¿Cómo funciona la fotosíntesis, en términos simples?",
    "¿Cuáles son las diferencias principales entre un virus y una bacteria?",
    "Explica qué es la inflación económica y por qué ocurre.",
    "¿Qué es el efecto invernadero y cómo contribuye al cambio climático?",
    "Describe brevemente la Revolución Francesa y sus causas principales.",
    "¿Por qué la Luna siempre muestra la misma cara hacia la Tierra?",
    "Explica la diferencia entre memoria RAM y almacenamiento en disco.",
    "¿Qué es la democracia representativa y en qué se diferencia de la democracia directa?",
    "Dame consejos para mejorar la calidad del sueño.",
    "¿Cómo se forma un arcoíris?",
    "Explica qué es la meditación mindfulness y sus beneficios principales.",
    "¿Qué diferencia hay entre el café arábica y el robusta?",
    "Cuéntame brevemente la historia de la Torre Eiffel.",
    "¿Por qué el agua se expande al congelarse, a diferencia de la mayoría de líquidos?",
    "Explica qué son los agujeros negros de forma sencilla.",
    "¿Qué papel juega la Reserva Federal en la economía de Estados Unidos?",
    "Dame una receta simple para hacer pan casero.",
    "¿Cuáles son los beneficios de aprender un segundo idioma?",
    "Explica la diferencia entre un huracán, un tifón y un ciclón.",
    "¿Cómo funcionan las vacunas para generar inmunidad?",
    "Describe brevemente qué fue la Guerra Fría.",
    "¿Qué es la biodiversidad y por qué es importante conservarla?",
    "Explica cómo funciona el reciclaje de plástico.",
    "¿Cuáles son las principales teorías sobre el origen del universo?",
]

CODIGO = [
    "Escribe una función en Python que determine si un número es primo.",
    "¿Cómo invierto una lista en Python sin usar reversed()?",
    "Escribe una función en JavaScript que calcule el factorial de un número usando recursión.",
    "Explica la diferencia entre una lista y una tupla en Python.",
    "¿Cómo implemento una búsqueda binaria en Python?",
    "Escribe una consulta SQL que obtenga los 5 empleados con mayor salario de una tabla 'employees'.",
    "¿Qué es un decorador en Python y para qué sirve? Da un ejemplo simple.",
    "Escribe una función que detecte si una cadena de texto es un palíndromo.",
    "Explica la diferencia entre == y === en JavaScript.",
    "¿Cómo se implementa una pila (stack) usando una lista en Python?",
    "Escribe una función en Python que elimine duplicados de una lista preservando el orden.",
    "¿Qué es la complejidad Big-O y por qué importa al escribir código?",
    "Escribe un ejemplo de una clase en Python con herencia simple.",
    "¿Cómo funciona el garbage collector en Python?",
    "Escribe una función que combine dos listas ordenadas en una sola lista ordenada.",
    "Explica qué es una promesa (Promise) en JavaScript y cómo se usa con async/await.",
    "¿Cómo se hace una petición HTTP GET en Python usando la librería requests?",
    "Escribe una función recursiva que calcule los primeros N números de Fibonacci.",
    "¿Qué diferencia hay entre un array y un objeto en JavaScript?",
    "Escribe una función en Python que cuente la frecuencia de cada palabra en un texto.",
    "Explica qué es un índice en una base de datos y cuándo conviene usarlo.",
    "¿Cómo se maneja una excepción en Python usando try/except?",
    "Escribe una función que ordene una lista usando el algoritmo bubble sort.",
    "¿Qué es Docker y para qué se usa en desarrollo de software?",
    "Escribe un ejemplo de una API REST simple usando Flask en Python.",
    "¿Cómo se implementa un patrón singleton en Python?",
    "Explica la diferencia entre programación síncrona y asíncrona.",
    "Escribe una función que verifique si un árbol binario está balanceado.",
]

RAZONAMIENTO = [
    "Si un tren viaja a 80 km/h y otro a 100 km/h en direcciones opuestas partiendo del mismo punto, ¿qué distancia los separa después de 2 horas?",
    "Tengo 3 cajas: una con manzanas, una con naranjas, y una con ambas, pero las 3 etiquetas están mal puestas. ¿Cómo puedo identificar el contenido correcto sacando una sola fruta?",
    "¿Cuál es la suma de todos los números del 1 al 100?",
    "Si hoy es miércoles, ¿qué día será dentro de 100 días?",
    "Un granjero tiene gallinas y conejos, en total 35 cabezas y 94 patas. ¿Cuántas gallinas y cuántos conejos tiene?",
    "Explica la paradoja del cumpleaños: ¿por qué en un grupo de 23 personas hay más del 50% de probabilidad de que dos compartan cumpleaños?",
    "Resuelve: si 5 máquinas hacen 5 productos en 5 minutos, ¿cuánto tardan 100 máquinas en hacer 100 productos?",
    "¿Cuántos cuadrados hay en un tablero de ajedrez de 8x8, contando todos los tamaños posibles?",
    "Tres personas se reparten una cuenta de 30 dólares, pagando 10 cada una. El camarero se da cuenta que debía ser 25 y les devuelve 5 dólares (1 cada una, sobra 2). Cada uno pagó 9, total 27, más los 2 que se guardó el camarero, son 29. ¿Dónde está el dólar faltante?",
    "Si A es mayor que B, y B es mayor que C, ¿qué se puede concluir sobre la relación entre A y C?",
    "¿Cuál es la probabilidad de obtener al menos un 6 al lanzar un dado tres veces?",
    "Un caracol sube 3 metros de día por un pozo de 10 metros, pero resbala 2 metros de noche. ¿Cuántos días tarda en salir del pozo?",
    "Demuestra por qué la raíz cuadrada de 2 es un número irracional.",
    "Si todos los gatos son mamíferos, y algunos mamíferos son negros, ¿se puede concluir que algunos gatos son negros?",
    "¿Cuál es el siguiente número en la secuencia: 2, 6, 12, 20, 30, ...?",
    "Resuelve la ecuación: 2x + 5 = 17.",
    "Si lanzo dos monedas, ¿cuál es la probabilidad de obtener exactamente una cara?",
    "Explica por qué 0.999... (periódico) es igual a 1.",
    "Un reloj se atrasa 3 minutos cada hora. Si lo pongo en hora a las 12:00, ¿qué hora marcará cuando el reloj real marque las 18:00?",
    "¿Cuántas formas distintas hay de ordenar las letras de la palabra 'CASA'?",
    "Si el área de un círculo es 78.5 cm², ¿cuál es aproximadamente su radio? (usa pi=3.14)",
    "Tengo el doble de la edad que tenía hace 10 años más 5 años. Si tengo 45 años, ¿qué edad tenía hace 10 años?",
    "Explica la diferencia entre correlación y causalidad con un ejemplo.",
    "¿Cuál es la probabilidad de que, al repartir una baraja de 52 cartas, las primeras 2 cartas sean ases?",
]

SEED_PROMPTS = CONVERSACION_GENERAL + CODIGO + RAZONAMIENTO


def build_jobs():
    jobs = []
    for pass_idx in range(N_PASSES):
        temp = TEMPERATURES[pass_idx % len(TEMPERATURES)]
        for prompt in SEED_PROMPTS:
            jobs.append((prompt, temp))
    return jobs


def call_model(prompt, temperature):
    payload = {
        "model": MODEL,
        "messages": [{"role": "user", "content": prompt}],
        "max_tokens": MAX_TOKENS,
        "temperature": temperature,
        "chat_template_kwargs": {"preserve_thinking": True},
    }
    data = json.dumps(payload).encode("utf-8")
    req = urllib.request.Request(
        API_URL, data=data, headers={"Content-Type": "application/json"}, method="POST"
    )
    with urllib.request.urlopen(req, timeout=TIMEOUT_S) as resp:
        return json.loads(resp.read().decode("utf-8"))


def main():
    jobs = build_jobs()
    total = len(jobs)
    print(f"Total de ejemplos a generar: {total} ({len(SEED_PROMPTS)} prompts semilla x {len(TEMPERATURES)} temperaturas)")

    stats = {"ok": 0, "truncated": 0, "failed": 0}
    start = time.time()

    with open(OUTPUT_PATH, "a", encoding="utf-8") as out:
        for i, (prompt, temp) in enumerate(jobs, start=1):
            attempt = 0
            while True:
                attempt += 1
                try:
                    resp = call_model(prompt, temp)
                    break
                except (urllib.error.URLError, TimeoutError, json.JSONDecodeError) as e:
                    if attempt > MAX_RETRIES:
                        print(f"[{i}/{total}] FALLO definitivo en prompt {prompt[:50]!r}: {e}")
                        stats["failed"] += 1
                        resp = None
                        break
                    print(f"[{i}/{total}] error ({e}), reintentando...")
                    time.sleep(2)

            if resp is not None:
                choice = resp["choices"][0]
                msg = choice["message"]
                finish_reason = choice["finish_reason"]
                if finish_reason == "length":
                    stats["truncated"] += 1
                else:
                    stats["ok"] += 1

                example = {
                    "messages": [
                        {"role": "user", "content": prompt},
                        {
                            "role": "assistant",
                            "content": msg.get("content"),
                            "reasoning_content": msg.get("reasoning"),
                        },
                    ],
                    "meta": {"temperature": temp, "finish_reason": finish_reason},
                }
                out.write(json.dumps(example, ensure_ascii=False) + "\n")
                out.flush()

            if i % 50 == 0 or i == total:
                elapsed = time.time() - start
                print(
                    f"[{i}/{total}] progreso -- ok={stats['ok']} truncated={stats['truncated']} "
                    f"failed={stats['failed']} -- {elapsed:.0f}s transcurridos"
                )

    elapsed = time.time() - start
    print(f"\n=== Terminado en {elapsed:.0f}s ({elapsed/60:.1f} min) ===")
    print(f"ok={stats['ok']} truncated={stats['truncated']} failed={stats['failed']} total={total}")


if __name__ == "__main__":
    main()
