import os, io, re
os.environ["GRADIO_ANALYTICS_ENABLED"] = "False"
os.environ["HF_HUB_DISABLE_TELEMETRY"] = "1"

import gradio as gr
from pdfminer.high_level import extract_text
from sentence_transformers import SentenceTransformer
import numpy as np

EMBED_MODEL_ID = "intfloat/multilingual-e5-small"  # malý, multi-lingual
_embedder = None

# In-memory knižnica
library_chunks = []          # list[str]
library_meta = []            # list[{"source":..., "idx":...}]
library_embeddings = None    # np.ndarray [N, D] (normované)

def get_embedder():
    global _embedder
    if _embedder is None:
        _embedder = SentenceTransformer(EMBED_MODEL_ID)
    return _embedder

def clean_text(t):
    return re.sub(r"\s+", " ", t).strip()

def chunk_text(text, size=1500, overlap=200):
    text = clean_text(text)
    chunks, i = [], 0
    step = max(1, size - overlap)
    L = len(text)
    while i < L:
        seg = text[i:i+size]
        if seg.strip():
            chunks.append(seg)
        i += step
    return chunks

def embed_passages(texts):
    emb = get_embedder()
    vecs = emb.encode([f"passage: {t}" for t in texts], normalize_embeddings=True)
    return np.asarray(vecs, dtype=np.float32)

def embed_query(q):
    emb = get_embedder()
    vec = emb.encode([f"query: {q}"], normalize_embeddings=True)[0]
    return np.asarray(vec, dtype=np.float32)

def add_to_library(chunks, source_name):
    global library_chunks, library_meta, library_embeddings
    vecs = embed_passages(chunks)
    if library_embeddings is None:
        library_embeddings = vecs
    else:
        library_embeddings = np.vstack([library_embeddings, vecs])
    start = len(library_chunks)
    library_chunks.extend(chunks)
    library_meta.extend({"source": source_name, "idx": i} for i in range(start, start+len(chunks)))
    return len(chunks)

def reset_library():
    global library_chunks, library_meta, library_embeddings
    library_chunks = []
    library_meta = []
    library_embeddings = None

def top_k(query, k=5):
    if library_embeddings is None or len(library_chunks) == 0:
        return []
    qv = embed_query(query)
    scores = library_embeddings @ qv  # cosine, lebo normované
    idx = np.argsort(-scores)[:k]
    return [(float(scores[i]), library_chunks[i], library_meta[i]) for i in idx]

def build_answer(question, hits):
    lines = [f"**Otázka:** {question}", "\n**Odpoveď (zo zdrojov):**"]
    if not hits:
        lines.append("_Nenašli sa relevantné pasáže. Nahraj PDF alebo skús inú formuláciu._")
    for i,(s,txt,meta) in enumerate(hits, start=1):
        snippet = clean_text(txt)[:400]
        src = meta.get("source", "PDF")
        lines.append(f"[{i}] {snippet} … (score {s:.2f}, zdroj: {src})")
    return "\n".join(lines)

# ---------- UI (Gradio 3.x) ----------
with gr.Blocks(title="EDM & Fanuc – PDF RAG (bez Qdrant)") as demo:
    gr.Markdown("# 🔧 EDM & Fanuc – PDF RAG (bez Qdrant)")

    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 📥 Ingest PDF (textové, nie skeny)")
            topic_dd = gr.Dropdown(choices=["EDM","FANUC","general"], value="EDM", label="Topic (len štítok)")
            chunk_chars = gr.Slider(800, 2500, value=1500, step=100, label="Chunk size")
            overlap = gr.Slider(0, 400, value=200, step=20, label="Overlap")
            upload = gr.File(file_count="multiple", file_types=[".pdf"], label="📁 PDF súbory")
            ingest_btn = gr.Button("📥 Pridať do knižnice")
            reset_btn = gr.Button("🗑️ Vyčistiť knižnicu")
            ingest_out = gr.Markdown()

            def _ingest(files, topic, size, ov):
                if not files:
                    return "⚠️ Nahraj aspoň jedno PDF."
                total = 0
                for f in files:
                    name = getattr(f, "name", "document.pdf")
                    text = extract_text(io.BytesIO(f.read()))
                    chunks = chunk_text(text, size, ov)
                    total += add_to_library(chunks, f"{topic}:{name}")
                return f"✅ Pridané chunky: {total}. V knižnici: {len(library_chunks)} pasáží."

            def _reset():
                reset_library()
                return "✅ Knižnica vyčistená."

            ingest_btn.click(_ingest, inputs=[upload, topic_dd, chunk_chars, overlap], outputs=ingest_out)
            reset_btn.click(_reset, outputs=ingest_out)

        with gr.Column(scale=2):
            gr.Markdown("### 💬 Chat z lokálnej knižnice")
            question = gr.Textbox(label="Otázka", placeholder="Napr. G41 vs G42 na Fanuc…", lines=3)
            topk = gr.Slider(1, 10, value=5, step=1, label="Top-K pasáží")
            ask_btn = gr.Button("💬 Spýtať sa")
            answer_md = gr.Markdown()

            def _ask(q, k):
                return build_answer(q, top_k(q, int(k)))

            ask_btn.click(_ask, inputs=[question, topk], outputs=answer_md)

if __name__ == "__main__":
    # V 3.x netreba show_api=False; táto verzia je stabilná a bez "No API found"
    demo.launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)))