#!/usr/bin/env python3
"""(Re)construieste rag_index.json din toate documentele din depozit (store.py),
cu embeddings Ollama. Rulat manual sau declansat din dashboard (`/api/reindex`).

Doua strategii de taiere in chunk-uri:
- **.xml** — un chunk per element de nivel 1 (o problema = un chunk). Sursele sunt
  structurate `<eroare_x><mesaj_eroare>…</mesaj_eroare><rezolvare>…</rezolvare></eroare_x>`,
  iar taierea pe linii goale ar rupe mesajul de eroare de rezolvarea lui: cautarea ar
  gasi eroarea si ar returna un chunk fara raspuns.
- **restul** — paragrafe unite pana la ~200 de caractere, ca inainte.
"""

from __future__ import annotations

import json
import os
import re
import sys
import xml.etree.ElementTree as ET

import requests

import config
import store


def chunk_text(text: str) -> list[str]:
    # imparte pe linii goale in paragrafe, uneste bucatile mici cu urmatoarea
    raw_parts = re.split(r"\n\s*\n", text.strip())
    chunks: list[str] = []
    buffer = ""
    for part in raw_parts:
        part = part.strip()
        if not part:
            continue
        buffer = f"{buffer}\n\n{part}" if buffer else part
        if len(buffer) >= 200:
            chunks.append(buffer)
            buffer = ""
    if buffer:
        chunks.append(buffer)
    return chunks


# Peste atat, un singur element XML se taie mai departe ca text obisnuit.
MAX_XML_CHUNK = 4000


def _render_element(el, depth: int = 0) -> str:
    """Elementul, ca text cu etichete pe post de etichete: `mesaj_eroare: ...`.

    Numele etichetelor poarta intelesul (care bucata e eroarea, care e rezolvarea),
    deci raman in text — dar fara paranteze unghiulare, care doar dilueaza embedding-ul.
    """
    lines = []
    own = (el.text or "").strip()
    label = el.tag.replace("_", " ")
    children = list(el)
    if own:
        lines.append(f"{label}: {own}" if children or depth else own)
    elif children and depth == 0:
        lines.append(f"[{label}]")
    for child in children:
        rendered = _render_element(child, depth + 1)
        if rendered:
            lines.append(rendered)
        tail = (child.tail or "").strip()
        if tail:
            lines.append(tail)
    return "\n".join(lines)


def _split_long(text: str, limit: int = MAX_XML_CHUNK) -> list[str]:
    """Taie o bucata prea lunga pe granite de cuvant.

    `chunk_text` imparte pe linii goale, deci un element XML scris ca un singur
    paragraf lung nu s-ar imparti deloc — de aici taierea asta, pe lungime.
    """
    out: list[str] = []
    for part in chunk_text(text) or [text]:
        while len(part) > limit:
            cut = part.rfind(" ", 0, limit)
            if cut <= 0:
                cut = limit
            out.append(part[:cut].strip())
            part = part[cut:].lstrip()
        if part.strip():
            out.append(part.strip())
    return out


def chunk_xml(text: str) -> list[str] | None:
    """Un chunk per element de nivel 1. `None` daca XML-ul nu se poate parsa."""
    try:
        root = ET.fromstring(text)
    except ET.ParseError:
        return None

    chunks: list[str] = []
    for el in root:
        rendered = _render_element(el).strip()
        if not rendered:
            continue
        if len(rendered) > MAX_XML_CHUNK:
            chunks.extend(_split_long(rendered))
        else:
            chunks.append(rendered)
    if not chunks:
        # radacina fara copii (sau doar text): trateaz-o ca document obisnuit
        whole = _render_element(root).strip()
        return chunk_text(whole) if whole else []
    return chunks


def chunk_document(name: str, text: str) -> list[str]:
    if name.endswith(".xml"):
        chunks = chunk_xml(text)
        if chunks is not None:
            return chunks
        print(f"[indexer] {name}: XML invalid, il tai ca text simplu", file=sys.stderr)
    return chunk_text(text)


def embed(text: str) -> list[float]:
    resp = requests.post(
        f"{config.get('OLLAMA_URL')}/api/embeddings",
        json={"model": config.get("EMBED_MODEL"), "prompt": text},
        timeout=60,
    )
    resp.raise_for_status()
    return resp.json()["embedding"]


def _vectori_existenti() -> dict[str, list[float]]:
    """Embeddings-urile din indexul curent, dupa textul chunk-ului.

    Un embedding pe CPU costa ~7 secunde. Fara asta, adaugarea unui singur
    document reface toate embeddings-urile — la 170 de chunk-uri inseamna 20 de
    minute ca sa se schimbe unul. Textul chunk-ului e cheia potrivita: daca nu s-a
    schimbat niciun caracter, vectorul e acelasi.

    Sursa e chiar indexul precedent, nu un al doilea fisier de cache: nu are ce sa
    se desincronizeze, si daca indexul lipseste se recalculeaza tot, ca inainte.

    Se refolosesc doar intrarile scrise cu modelul curent. Doua modele dau vectori
    diferiti (adesea si de alta lungime) pentru acelasi text, iar un index amestecat
    ar strica ordonarea fara nici un mesaj de eroare. Intrarile vechi, dinainte de
    campul `model`, se trateaza ca necunoscute: se recalculeaza o singura data.
    """
    try:
        vechi = json.loads(config.INDEX_FILE.read_text(encoding="utf-8"))
    except (OSError, ValueError):
        return {}
    model = config.get("EMBED_MODEL")
    return {
        e["text"]: e["embedding"]
        for e in vechi
        if e.get("text") and e.get("embedding") and e.get("model") == model
    }


def build() -> dict:
    entries = []
    warnings: list[str] = []
    docs = store.documents_for_index()
    refolosite = _vectori_existenti()
    model = config.get("EMBED_MODEL")
    reutilizari = 0
    for doc in docs:
        text = store.read_document(doc["name"])
        if doc["name"].endswith(".xml") and chunk_xml(text) is None:
            warnings.append(f"{doc['name']}: XML invalid, indexat ca text simplu")
        for i, chunk in enumerate(chunk_document(doc["name"], text)):
            vec = refolosite.get(chunk)
            if vec is None:
                vec = embed(chunk)
            else:
                reutilizari += 1
            entries.append({"source": doc["name"], "chunk": i, "text": chunk,
                            "embedding": vec, "model": model})
    config.STATE_DIR.mkdir(parents=True, exist_ok=True)
    # Scriere atomica: consumer-ul reciteste fisierul la 30s si ar putea prinde
    # un JSON pe jumatate scris daca am scrie direct peste el.
    tmp = config.INDEX_FILE.with_suffix(".json.tmp")
    tmp.write_text(json.dumps(entries, ensure_ascii=False), encoding="utf-8")
    os.replace(tmp, config.INDEX_FILE)
    out = {"documents": len(docs), "chunks": len(entries),
           "embeddings_refolosite": reutilizari,
           "embeddings_noi": len(entries) - reutilizari}
    if warnings:
        out["warnings"] = warnings
    return out


if __name__ == "__main__":
    try:
        with config.exclusive():
            result = build()
    except config.Busy as exc:
        print(f"[indexer] {exc}, ies fara sa fac nimic", file=sys.stderr)
        raise SystemExit(0)
    for w in result.get("warnings", []):
        print(f"[indexer] ATENTIE {w}", file=sys.stderr)
    print(
        f"[indexer] {result['documents']} documente, {result['chunks']} chunk-uri "
        f"({result['embeddings_noi']} embeddings noi, {result['embeddings_refolosite']} refolosite)"
        f" -> {config.INDEX_FILE}",
        file=sys.stderr,
    )