#!/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 build() -> dict:
    entries = []
    warnings: list[str] = []
    docs = store.documents_for_index()
    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 = embed(chunk)
            entries.append({"source": doc["name"], "chunk": i, "text": chunk, "embedding": vec})
    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)}
    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 -> {config.INDEX_FILE}",
        file=sys.stderr,
    )