#!/usr/bin/env python3
"""Convierte COFRE en algo que JARVIS puede consultar antes de contestar.

    ./indexa.py            reconstruye el indice
    ./indexa.py --cuenta   solo dice que encontraria, sin indexar

## El problema que resuelve

COFRE tiene 35.000 documentos y 30 GB. Indexarlo entero seria a la vez inutil y
lentisimo: el 95 % es codigo fuente de terceros —node_modules, .git, dist— que
no dice nada de COMO USAR una herramienta. Una busqueda sobre eso devuelve
ruido con total confianza, que es peor que no tener nada.

Lo que de verdad vale es el SABER, y son tres cosas:

  1. Los apuntes propios del usuario: 26.000 palabras de flujos de trabajo
     escritos probandolos. Es lo mas valioso porque es correcto y es suyo.
  2. Los READMEs de primer nivel de cada herramienta: el "para que sirve y como
     se llama", una vez por herramienta.
  3. Una ficha por herramienta: nombre, perfil, ruta, una linea. Asi JARVIS
     sabe QUE tiene aunque no haya un README.

## Como se trocea

En fragmentos por seccion (los encabezados markdown) y no en trozos ciegos de N
caracteres: un apunte partido por la mitad de una frase recupera peor que uno
partido por donde el autor ya separo las ideas. Cada fragmento recuerda de que
fichero y que herramienta viene, para poder citar la fuente.
"""
import argparse
import json
import os
import re
import sys

AQUI = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
COFRE = os.path.expanduser("~/COFRE")
INDICE = os.path.join(AQUI, "datos", "saber.jsonl")

# Los perfiles que se indexan. Cada uno es una carpeta de ~/COFRE.
PERFILES = ["PENTESTERS", "PHISHERS", "REVERSERS", "PROTECTORS", "HARD-WARERS",
            "AUTOMATAS", "DEFACERS", "SNEAKERS", "ZAPPERS"]

# Lo que NO se mira: dependencias, control de versiones, builds. Aqui esta la
# diferencia entre un indice util y 30 GB de ruido.
IGNORA = re.compile(
    r"(^|/)(node_modules|\.git|vendor|dist|build|__pycache__|\.venv|venv|"
    r"site-packages|target|\.cache|test|tests|fixtures|examples?|locale|"
    r"i18n|translations|\.github)(/|$)", re.I)

# Documentos que son saber, no plantillas de proyecto ni codigo de conducta.
DOC_UTIL = re.compile(r"(readme|install|usage|apuntes|notas|trucos|howto|"
                      r"documentation|guide|guia|cheat|manual)", re.I)
DOC_RUIDO = re.compile(r"(code_of_conduct|contributing|changelog|license|"
                       r"copying|pull_request|issue_template|security\.md|"
                       r"hall_of_fame|redistributed|historical)", re.I)

MIN_FRAGMENTO = 80        # menos de esto no es un parrafo, es un titulo suelto
MAX_FRAGMENTO = 1200      # mas de esto recupera con demasiado ruido alrededor


def _profundidad(ruta):
    return ruta.rstrip("/").count("/")


def herramientas(perfil):
    """Cada subcarpeta de primer nivel es una herramienta. Una ficha por cada."""
    base = os.path.join(COFRE, perfil)
    if not os.path.isdir(base):
        return
    for nombre in sorted(os.listdir(base)):
        ruta = os.path.join(base, nombre)
        if not os.path.isdir(ruta) or nombre.startswith("."):
            continue
        yield {"herramienta": nombre, "perfil": perfil, "ruta": ruta}


def _resumen_readme(ruta_herramienta):
    """El primer parrafo con sustancia de su README, si tiene."""
    for nombre in ("README.md", "README", "README.txt", "readme.md"):
        p = os.path.join(ruta_herramienta, nombre)
        if not os.path.exists(p):
            continue
        try:
            with open(p, encoding="utf-8", errors="ignore") as f:
                texto = f.read(4000)
        except OSError:
            continue
        # saltar titulos, insignias y lineas vacias hasta el primer parrafo real
        for parrafo in re.split(r"\n\s*\n", texto):
            limpio = re.sub(r"[#>*`\[\]!]|\(https?://[^)]+\)|<[^>]+>", "", parrafo).strip()
            limpio = re.sub(r"\s+", " ", limpio)
            if len(limpio) >= 60 and not limpio.lower().startswith(("http", "===")):
                return limpio[:400]
    return ""


def documentos_utiles(perfil):
    """Los .md/.txt que son saber, no codigo ni plantillas."""
    base = os.path.join(COFRE, perfil)
    for raiz, dirs, ficheros in os.walk(base):
        if IGNORA.search(raiz):
            dirs[:] = []
            continue
        # no bajar mas de 3 niveles: el saber vive arriba, el codigo abajo
        if _profundidad(raiz) - _profundidad(base) > 3:
            dirs[:] = []
            continue
        for fich in ficheros:
            if not fich.lower().endswith((".md", ".txt")):
                continue
            if DOC_RUIDO.search(fich):
                continue
            # los apuntes propios entran siempre; de terceros, solo los utiles
            propio = re.search(r"apuntes|trucos|notas|cheat|guia", fich, re.I)
            if not propio and not DOC_UTIL.search(fich):
                continue
            yield os.path.join(raiz, fich)


def trocea(ruta):
    """Un fichero markdown en fragmentos por seccion."""
    try:
        with open(ruta, encoding="utf-8", errors="ignore") as f:
            texto = f.read()
    except OSError:
        return
    # partir por encabezados, conservando el titulo con su seccion
    partes = re.split(r"\n(?=#{1,3}\s)", texto)
    for parte in partes:
        limpio = parte.strip()
        if len(limpio) < MIN_FRAGMENTO:
            continue
        # si una seccion es enorme, partirla por parrafos sin pasar del tope
        if len(limpio) <= MAX_FRAGMENTO:
            yield limpio
        else:
            buffer = ""
            for parrafo in re.split(r"\n\s*\n", limpio):
                if len(buffer) + len(parrafo) > MAX_FRAGMENTO and buffer:
                    yield buffer.strip()
                    buffer = ""
                buffer += parrafo + "\n\n"
            if len(buffer.strip()) >= MIN_FRAGMENTO:
                yield buffer.strip()


def construye(solo_cuenta=False, solo_cofre=False):
    fichas, fragmentos = [], []

    for perfil in PERFILES:
        n_herr = 0
        for h in herramientas(perfil):
            resumen = _resumen_readme(h["ruta"])
            texto = f"{h['herramienta']} ({perfil}). {resumen}".strip()
            fichas.append({
                "tipo": "ficha", "herramienta": h["herramienta"],
                "perfil": perfil, "ruta": h["ruta"], "texto": texto,
            })
            n_herr += 1

        n_frag = 0
        for doc in documentos_utiles(perfil):
            herr = _herramienta_de(doc, perfil)
            for trozo in trocea(doc):
                fragmentos.append({
                    "tipo": "apunte", "herramienta": herr, "perfil": perfil,
                    "fichero": os.path.relpath(doc, COFRE), "texto": trozo,
                })
                n_frag += 1
        print(f"  {perfil:14s} {n_herr:4d} herramientas   {n_frag:5d} fragmentos de apuntes")

    # El conocimiento de la propia maquina: comandos de Linux, servidores SSH,
    # extensiones. Es lo que convierte "sabe de mis herramientas" en "sabe hacer
    # cualquier cosa en Linux". Va salvo que se pida solo COFRE.
    sistema = []
    if not solo_cofre:
        try:
            from saber import sistema as sis
        except ImportError:
            import sistema as sis
        print("  extrayendo la documentacion del sistema (man pages, ssh, firefox)...",
              flush=True)
        n_cmd = n_srv = n_nav = 0
        for e in sis.todo():
            sistema.append(e)
            n_cmd += e["tipo"] == "comando"
            n_srv += e["tipo"] == "servidor"
            n_nav += e["tipo"] == "navegador"
            if len(sistema) % 500 == 0:
                print(f"    {len(sistema)} entradas de sistema...", flush=True)
        print(f"  LINUX          {n_cmd:4d} comandos   {n_srv} servidores   {n_nav} navegador")

    todo = fichas + fragmentos + sistema
    print(f"\n  total: {len(fichas)} fichas + {len(fragmentos)} apuntes + "
          f"{len(sistema)} sistema = {len(todo)} entradas")
    if solo_cuenta:
        return 0

    print("  cargando el modelo de embeddings...", flush=True)
    from model2vec import StaticModel
    modelo = StaticModel.from_pretrained("minishlab/potion-multilingual-128M")

    print("  calculando vectores...", flush=True)
    vectores = modelo.encode([e["texto"] for e in todo], show_progress_bar=False)

    os.makedirs(os.path.dirname(INDICE), exist_ok=True)
    with open(INDICE, "w", encoding="utf-8") as f:
        for entrada, vec in zip(todo, vectores):
            entrada["v"] = [round(float(x), 5) for x in vec]
            f.write(json.dumps(entrada, ensure_ascii=False) + "\n")
    print(f"  indice guardado: {INDICE} ({os.path.getsize(INDICE)//1024} KB)")
    return 0


def _herramienta_de(ruta, perfil):
    partes = os.path.relpath(ruta, os.path.join(COFRE, perfil)).split(os.sep)
    return partes[0] if len(partes) > 1 else perfil


def main() -> int:
    p = argparse.ArgumentParser()
    p.add_argument("--cuenta", action="store_true",
                   help="di que encontrarias, sin indexar")
    p.add_argument("--solo-cofre", action="store_true",
                   help="indexar solo COFRE, sin la documentacion del sistema")
    a = p.parse_args()
    return construye(solo_cuenta=a.cuenta, solo_cofre=a.solo_cofre)


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