#!/usr/bin/env python3
"""Genera el dataset de fine-tuning a partir de los apuntes de COFRE.

    ./dataset.py --muestra 5     enseña 5 ejemplos y para (para mirar la calidad)
    ./dataset.py                 genera el dataset entero (horas, en segundo plano)

## Por que esto es el paso que faltaba

Fine-tuning no es magia: es enseñarle al modelo con EJEMPLOS de como quieres que
conteste. Para eso hacen falta pares pregunta-respuesta, y no los teniamos. Los
apuntes de COFRE son conocimiento, pero en prosa; hay que convertirlos en la
forma "alguien pregunta X, JARVIS responde Y".

Se generan con el modelo local a partir de cada fragmento de los apuntes: se le
da el trozo y se le pide que invente las preguntas que ese trozo contesta, y la
respuesta al estilo de JARVIS. El fragmento es la VERDAD —lo escribio el
usuario— asi que las respuestas salen de ahi y no de lo que el modelo recuerde.

## Por que el dataset vale aunque no se entrene aqui

La T1200 de 4 GB va justa para entrenar un 4B. El sitio natural es el servidor
con la 3060 de 12 GB. Pero el DATASET es lo caro de conseguir y no depende de la
maquina: se genera una vez, es un fichero, y se entrena donde haga falta —aqui
si cabe, o en el servidor cuando haya acceso—.

## El formato

JSONL en el formato de mensajes, que es lo que comen trl/axolotl/unsloth:

    {"messages": [{"role":"user","content":"..."},
                  {"role":"assistant","content":"..."}]}
"""
import argparse
import json
import os
import re
import sys
import urllib.request

AQUI = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
INDICE = os.path.join(AQUI, "nucleo", "datos", "saber.jsonl")
SALIDA = os.path.join(AQUI, "entrena", "dataset-cofre.jsonl")
ENDPOINT = os.environ.get("JARVIS_OLLAMA", "http://127.0.0.1:11434")
MODELO = os.environ.get("JARVIS_MODELO", "qwen3.5:4b-jarvis")

# Al generador se le da MAS libertad que a JARVIS en produccion: aqui quereis
# variedad de preguntas, no una respuesta conservadora.
GENERADOR = """Eres un generador de datos de entrenamiento para un asistente de
pentesting llamado JARVIS. Te doy un fragmento de los apuntes tecnicos del
usuario. Devuelve de UNA a TRES preguntas que una persona haria y que ese
fragmento contesta, cada una con su respuesta.

Reglas:
- Las preguntas, naturales y variadas: como las diria alguien por voz.
- Las respuestas SALEN DEL FRAGMENTO. Si el fragmento tiene un comando exacto,
  la respuesta lo lleva tal cual. No te inventes flags que no aparezcan.
- Respuestas al estilo de JARVIS: directas, en castellano, tratando de "señor".
- Si el fragmento no da para una pregunta util (es un titulo, una lista de
  enlaces, ruido), devuelve una lista vacia.

Responde SOLO con JSON, un objeto con la clave "pares":
{"pares": [{"pregunta": "...", "respuesta": "..."}]}"""


def _pide(fragmento: str) -> list:
    cuerpo = json.dumps({
        "model": MODELO,
        "messages": [{"role": "system", "content": GENERADOR},
                     {"role": "user", "content": fragmento[:1500]}],
        "stream": False, "think": False,
        "options": {"temperature": 0.7, "num_predict": 500},
        "format": "json",
    }).encode()
    req = urllib.request.Request(ENDPOINT + "/api/chat", data=cuerpo,
                                 headers={"Content-Type": "application/json"})
    try:
        with urllib.request.urlopen(req, timeout=120) as r:
            contenido = json.loads(r.read())["message"]["content"]
    except Exception:
        return []
    # el modelo a veces envuelve el array en un objeto {"pares": [...]}
    try:
        d = json.loads(contenido)
    except json.JSONDecodeError:
        m = re.search(r"\[.*\]", contenido, re.S)
        if not m:
            return []
        try:
            d = json.loads(m.group(0))
        except json.JSONDecodeError:
            return []
    if isinstance(d, dict):
        d = next((v for v in d.values() if isinstance(v, list)), [])
    salida = []
    for par in d if isinstance(d, list) else []:
        if isinstance(par, dict) and par.get("pregunta") and par.get("respuesta"):
            salida.append((par["pregunta"].strip(), par["respuesta"].strip()))
    return salida


def apuntes():
    """Los fragmentos de apuntes del indice, los mas sustanciosos primero."""
    fuera = []
    with open(INDICE, encoding="utf-8") as f:
        for l in f:
            try:
                d = json.loads(l)
            except json.JSONDecodeError:
                continue
            if d.get("tipo") == "apunte" and len(d["texto"].split()) >= 20:
                d.pop("v", None)
                fuera.append(d)
    return fuera


def main() -> int:
    p = argparse.ArgumentParser()
    p.add_argument("--muestra", type=int, default=0,
                   help="genera solo N y los enseña, sin guardar")
    a = p.parse_args()

    if not os.path.exists(INDICE):
        print("  no hay indice: nucleo/saber/indexa.py primero")
        return 1
    fragmentos = apuntes()
    print(f"  {len(fragmentos)} fragmentos con sustancia", flush=True)

    if a.muestra:
        for frag in fragmentos[:a.muestra]:
            pares = _pide(frag["texto"])
            print(f"\n  ── {frag.get('fichero', frag['perfil'])} ──")
            for q, r in pares:
                print(f"    P: {q}")
                print(f"    R: {r[:120]}")
        return 0

    # Reanudable: un trabajo de 10 h no puede perder todo si se corta. Se apunta
    # en un .hecho por que fragmento se iba, y al arrancar se salta hasta ahi y
    # se AÑADE al dataset en vez de reescribirlo.
    marca = SALIDA + ".hecho"
    desde = 0
    if os.path.exists(marca) and os.path.exists(SALIDA):
        try:
            desde = int(open(marca).read().strip()) + 1
        except (ValueError, OSError):
            desde = 0
    modo = "a" if desde else "w"
    hechos = sum(1 for _ in open(SALIDA)) if desde and os.path.exists(SALIDA) else 0
    if desde:
        print(f"  reanudando desde el fragmento {desde} ({hechos} ejemplos hechos)",
              flush=True)

    with open(SALIDA, modo, encoding="utf-8", buffering=1) as out:
        for i in range(desde, len(fragmentos)):
            for q, r in _pide(fragmentos[i]["texto"]):
                out.write(json.dumps({"messages": [
                    {"role": "user", "content": q},
                    {"role": "assistant", "content": r},
                ]}, ensure_ascii=False) + "\n")
                hechos += 1
            with open(marca, "w") as m:
                m.write(str(i))
            if (i + 1) % 50 == 0:
                print(f"    {i+1}/{len(fragmentos)} fragmentos, {hechos} ejemplos",
                      flush=True)
    print(f"\n  dataset guardado: {SALIDA} ({hechos} ejemplos)")
    return 0


if __name__ == "__main__":
    sys.exit(main())