# -*- coding: utf-8 -*-
"""
HF Space - main.py de substitution pour tests Qdrant / indexation minimale

Fonctions clés :
- POST /wipe?project_id=XXX         : supprime la collection Qdrant
- POST /index                       : lance un job d'indexation (JSON files=[{path,text},...])
- GET  /status/{job_id}             : état du job + logs
- GET  /collections/{proj}/count    : retourne le nombre de points dans Qdrant
- POST /query                       : recherche sémantique (top_k, text, project_id)

Une UI Gradio minimale est montée sur "/" pour déclencher les tests sans console.

ENV attendues :
- QDRANT_URL            : ex. https://xxxxx.eu-central-1-0.aws.cloud.qdrant.io:6333
- QDRANT_API_KEY        : clé Qdrant Cloud
- COLLECTION_PREFIX     : défaut "proj_"
- EMB_PROVIDER          : "hf" (défaut) ou "dummy"
- HF_EMBED_MODEL        : défaut "BAAI/bge-m3"
- HUGGINGFACEHUB_API_TOKEN : token HF Inference (si EMB_PROVIDER=hf)
- LOG_LEVEL             : DEBUG (défaut), INFO...

Dépendances (requirements) suggérées :
fastapi>=0.111
uvicorn>=0.30
httpx>=0.27
pydantic>=2.7
gradio>=4.43
numpy>=2.0
"""

from __future__ import annotations
import os
import time
import uuid
import math
import json
import hashlib
import logging
import asyncio
from typing import List, Dict, Any, Optional, Tuple

import numpy as np
import httpx
from pydantic import BaseModel, Field, ValidationError
from fastapi import FastAPI, HTTPException, Query
from fastapi.middleware.cors import CORSMiddleware

import gradio as gr

# ------------------------------------------------------------------------------
# Configuration & logs
# ------------------------------------------------------------------------------
LOG_LEVEL = os.getenv("LOG_LEVEL", "DEBUG").upper()
logging.basicConfig(
    level=getattr(logging, LOG_LEVEL, logging.DEBUG),
    format="%(asctime)s - %(levelname)s - %(message)s",
)
LOG = logging.getLogger("remote_indexer_min")

QDRANT_URL = os.getenv("QDRANT_URL", "").rstrip("/")
QDRANT_API_KEY = os.getenv("QDRANT_API_KEY", "")
COLLECTION_PREFIX = os.getenv("COLLECTION_PREFIX", "proj_").strip() or "proj_"

EMB_PROVIDER = os.getenv("EMB_PROVIDER", "hf").lower()  # "hf" | "dummy"
HF_EMBED_MODEL = os.getenv("HF_EMBED_MODEL", "BAAI/bge-m3")
HF_TOKEN = os.getenv("HUGGINGFACEHUB_API_TOKEN", "")

if not QDRANT_URL or not QDRANT_API_KEY:
    LOG.warning("QDRANT_URL / QDRANT_API_KEY non fournis : l'upsert échouera. Fournis-les dans les Secrets du Space.")

if EMB_PROVIDER == "hf" and not HF_TOKEN:
    LOG.warning("EMB_PROVIDER=hf mais HUGGINGFACEHUB_API_TOKEN absent. Tu peux basculer EMB_PROVIDER=dummy pour tester sans token.")

# ------------------------------------------------------------------------------
# Schémas Pydantic
# ------------------------------------------------------------------------------
class FileItem(BaseModel):
    path: str
    text: str

class IndexRequest(BaseModel):
    project_id: str = Field(..., min_length=1)
    files: List[FileItem] = Field(default_factory=list)
    chunk_size: int = Field(200, ge=64, le=4096)
    overlap: int = Field(20, ge=0, le=512)
    batch_size: int = Field(32, ge=1, le=1024)
    store_text: bool = True

class QueryRequest(BaseModel):
    project_id: str
    text: str
    top_k: int = Field(5, ge=1, le=100)

# ------------------------------------------------------------------------------
# Job store (en mémoire)
# ------------------------------------------------------------------------------
class JobState(BaseModel):
    job_id: str
    project_id: str
    stage: str = "pending"      # pending -> embedding -> upserting -> done/failed
    total_files: int = 0
    total_chunks: int = 0
    embedded: int = 0
    upserted: int = 0
    errors: List[str] = Field(default_factory=list)
    messages: List[str] = Field(default_factory=list)
    started_at: float = Field(default_factory=time.time)
    finished_at: Optional[float] = None

    def log(self, msg: str) -> None:
        stamp = time.strftime("%H:%M:%S")
        line = f"[{stamp}] {msg}"
        self.messages.append(line)
        LOG.debug(f"[{self.job_id}] {msg}")

JOBS: Dict[str, JobState] = {}

# ------------------------------------------------------------------------------
# Utilitaires
# ------------------------------------------------------------------------------
def hash8(s: str) -> str:
    return hashlib.sha256(s.encode("utf-8")).hexdigest()[:16]

def l2_normalize(vec: List[float]) -> List[float]:
    arr = np.array(vec, dtype=np.float32)
    n = float(np.linalg.norm(arr))
    if n > 0:
        arr = arr / n
    return arr.astype(np.float32).tolist()

def flatten_any(x: Any) -> List[float]:
    """
    Certaines APIs renvoient [[...]] ou [[[...]]]; on aplanit en 1D.
    """
    if isinstance(x, (list, tuple)):
        if len(x) > 0 and isinstance(x[0], (list, tuple)):
            # Aplanit récursif
            return flatten_any(x[0])
        return list(map(float, x))
    raise ValueError("Embedding vector mal formé")

def chunk_text(text: str, chunk_size: int, overlap: int) -> List[Tuple[int, int, str]]:
    """
    Retourne une liste de (start, end, chunk_text)
    Ignore les petits fragments (< 30 chars) pour éviter le bruit.
    """
    text = text or ""
    if not text.strip():
        return []
    res = []
    n = len(text)
    i = 0
    while i < n:
        j = min(i + chunk_size, n)
        chunk = text[i:j]
        if len(chunk.strip()) >= 30:
            res.append((i, j, chunk))
        i = j - overlap
        if i <= 0:
            i = j
    return res

async def ensure_collection(client: httpx.AsyncClient, coll: str, vector_size: int) -> None:
    """
    Crée ou ajuste la collection Qdrant (distance = Cosine).
    """
    url = f"{QDRANT_URL}/collections/{coll}"
    # Vérifie l'existence
    r = await client.get(url, headers={"api-key": QDRANT_API_KEY}, timeout=20)
    if r.status_code == 200:
        # Optionnel: vérifier la taille du vecteur ; si mismatch, on peut supprimer/recréer
        data = r.json()
        existing_size = data.get("result", {}).get("vectors", {}).get("size")
        if existing_size and int(existing_size) != int(vector_size):
            LOG.warning(f"Collection {coll} dim={existing_size} ≠ attendu {vector_size} → recréation")
            await client.delete(url, headers={"api-key": QDRANT_API_KEY}, timeout=20)
        else:
            LOG.debug(f"Collection {coll} déjà prête (dim={existing_size})")
    # (Re)création
    body = {
        "vectors": {"size": vector_size, "distance": "Cosine"}
    }
    r2 = await client.put(url, headers={"api-key": QDRANT_API_KEY}, json=body, timeout=30)
    if r2.status_code not in (200, 201):
        raise HTTPException(status_code=500, detail=f"Qdrant PUT collection a échoué: {r2.text}")

async def qdrant_upsert(
    client: httpx.AsyncClient,
    coll: str,
    points: List[Dict[str, Any]],
) -> int:
    if not points:
        return 0
    url = f"{QDRANT_URL}/collections/{coll}/points?wait=true"
    body = {"points": points}
    r = await client.put(url, headers={"api-key": QDRANT_API_KEY}, json=body, timeout=60)
    if r.status_code not in (200, 202):
        raise HTTPException(status_code=500, detail=f"Qdrant upsert échoué: {r.text}")
    return len(points)

async def qdrant_count(client: httpx.AsyncClient, coll: str) -> int:
    url = f"{QDRANT_URL}/collections/{coll}/points/count"
    r = await client.post(
        url,
        headers={"api-key": QDRANT_API_KEY},
        json={"exact": True},
        timeout=20,
    )
    if r.status_code != 200:
        raise HTTPException(status_code=500, detail=f"Qdrant count échoué: {r.text}")
    return int(r.json().get("result", {}).get("count", 0))

async def qdrant_search(
    client: httpx.AsyncClient,
    coll: str,
    vector: List[float],
    limit: int = 5,
) -> Dict[str, Any]:
    url = f"{QDRANT_URL}/collections/{coll}/points/search"
    r = await client.post(
        url,
        headers={"api-key": QDRANT_API_KEY},
        json={"vector": vector, "limit": limit, "with_payload": True},
        timeout=30,
    )
    if r.status_code != 200:
        raise HTTPException(status_code=500, detail=f"Qdrant search échoué: {r.text}")
    return r.json()

# ------------------------------------------------------------------------------
# Embeddings (HF Inference ou dummy)
# ------------------------------------------------------------------------------
async def embed_hf(
    client: httpx.AsyncClient,
    texts: List[str],
    model: str = HF_EMBED_MODEL,
    token: str = HF_TOKEN,
) -> List[List[float]]:
    """
    Appel HuggingFace Inference (feature extraction) - batch.
    Normalise L2 les vecteurs.
    """
    if not token:
        raise HTTPException(status_code=400, detail="HUGGINGFACEHUB_API_TOKEN manquant pour EMB_PROVIDER=hf")
    url = f"https://api-inference.huggingface.co/models/{model}"
    headers = {"Authorization": f"Bearer {token}"}
    # HF accepte une liste de textes directement
    payload = {"inputs": texts, "options": {"wait_for_model": True}}
    r = await client.post(url, headers=headers, json=payload, timeout=120)
    if r.status_code != 200:
        raise HTTPException(status_code=502, detail=f"HF Inference error: {r.text}")
    data = r.json()
    # data peut être une liste de listes (ou de listes de listes...)
    embeddings: List[List[float]] = []
    if isinstance(data, list):
        for row in data:
            vec = flatten_any(row)
            embeddings.append(l2_normalize(vec))
    else:
        vec = flatten_any(data)
        embeddings.append(l2_normalize(vec))
    return embeddings

def embed_dummy(texts: List[str], dim: int = 128) -> List[List[float]]:
    """
    Embedding déterministe basé sur un hash -> vecteur pseudo-aléatoire stable.
    Suffisant pour tester le pipeline Qdrant (dimensions cohérentes, upsert, count, search).
    """
    out: List[List[float]] = []
    for t in texts:
        h = hashlib.sha256(t.encode("utf-8")).digest()
        # Étale sur dim floats
        arr = np.frombuffer((h * ((dim // len(h)) + 1))[:dim], dtype=np.uint8).astype(np.float32)
        # Centrage et normalisation
        arr = (arr - 127.5) / 127.5
        arr = arr / (np.linalg.norm(arr) + 1e-9)
        out.append(arr.astype(np.float32).tolist())
    return out

async def embed_texts(client: httpx.AsyncClient, texts: List[str]) -> List[List[float]]:
    if EMB_PROVIDER == "hf":
        return await embed_hf(client, texts)
    return embed_dummy(texts, dim=128)

# ------------------------------------------------------------------------------
# Pipeline d'indexation
# ------------------------------------------------------------------------------
async def run_index_job(job: JobState, req: IndexRequest) -> None:
    job.stage = "embedding"
    job.total_files = len(req.files)
    job.log(f"Index start project={req.project_id} files={len(req.files)} chunk_size={req.chunk_size} overlap={req.overlap} batch_size={req.batch_size} store_text={req.store_text}")

    # Dédup global par hash du texte brut des fichiers
    file_hashes = [hash8(f.text) for f in req.files]
    uniq = len(set(file_hashes))
    if uniq != len(file_hashes):
        job.log(f"Attention: {len(file_hashes)-uniq} fichiers ont un texte identique (hash dupliqué).")

    # Chunking
    records: List[Dict[str, Any]] = []
    for f in req.files:
        chunks = chunk_text(f.text, req.chunk_size, req.overlap)
        if not chunks:
            job.log(f"{f.path}: 0 chunk (trop court ou vide)")
        for idx, (start, end, ch) in enumerate(chunks):
            payload = {
                "path": f.path,
                "chunk": idx,
                "start": start,
                "end": end,
            }
            if req.store_text:
                payload["text"] = ch
            records.append({"payload": payload, "raw": ch})
    job.total_chunks = len(records)
    job.log(f"Total chunks = {job.total_chunks}")

    if job.total_chunks == 0:
        job.stage = "failed"
        job.errors.append("Aucun chunk à indexer.")
        job.finished_at = time.time()
        return

    # Embedding + Upsert (en batches)
    async with httpx.AsyncClient(timeout=120) as client:
        # Dimension à partir du 1er embedding (warmup)
        warmup_vec = (await embed_texts(client, [records[0]["raw"]]))[0]
        vec_dim = len(warmup_vec)
        job.log(f"Warmup embeddings dim={vec_dim} provider={EMB_PROVIDER}")

        # Qdrant collection
        coll = f"{COLLECTION_PREFIX}{req.project_id}"
        await ensure_collection(client, coll, vector_size=vec_dim)

        job.stage = "upserting"
        batch_vectors: List[List[float]] = []
        batch_points: List[Dict[str, Any]] = []

        async def flush_batch():
            nonlocal batch_vectors, batch_points
            if not batch_points:
                return 0
            added = await qdrant_upsert(client, coll, batch_points)
            job.upserted += added
            job.log(f"+{added} points upsert (total={job.upserted})")
            batch_vectors = []
            batch_points = []
            return added

        # Traite par lot d'embeddings (embedding_batch_size indépendant de l'upsert batch_size)
        EMB_BATCH = max(8, min(64, req.batch_size * 2))
        i = 0
        while i < len(records):
            sub = records[i : i + EMB_BATCH]
            texts = [r["raw"] for r in sub]
            vecs = await embed_texts(client, texts)
            if len(vecs) != len(sub):
                raise HTTPException(status_code=500, detail="Embedding batch size mismatch")
            job.embedded += len(vecs)

            for r, v in zip(sub, vecs):
                payload = r["payload"]
                point = {
                    "id": str(uuid.uuid4()),
                    "vector": v,
                    "payload": payload,
                }
                batch_points.append(point)
                if len(batch_points) >= req.batch_size:
                    await flush_batch()
            i += EMB_BATCH

        # Flush final
        await flush_batch()

    job.stage = "done"
    job.finished_at = time.time()
    job.log("Index job terminé.")

# ------------------------------------------------------------------------------
# FastAPI app + endpoints
# ------------------------------------------------------------------------------
fastapi_app = FastAPI(title="Remote Indexer - Minimal Test Space")
fastapi_app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

@fastapi_app.get("/")
async def root():
    return {"ok": True, "service": "remote-indexer-min", "qdrant": bool(QDRANT_URL), "emb_provider": EMB_PROVIDER}

@fastapi_app.post("/wipe")
async def wipe(project_id: str = Query(..., min_length=1)):
    if not QDRANT_URL or not QDRANT_API_KEY:
        raise HTTPException(status_code=400, detail="QDRANT_URL / QDRANT_API_KEY requis")
    coll = f"{COLLECTION_PREFIX}{project_id}"
    async with httpx.AsyncClient() as client:
        r = await client.delete(f"{QDRANT_URL}/collections/{coll}", headers={"api-key": QDRANT_API_KEY}, timeout=30)
        if r.status_code not in (200, 202, 404):
            raise HTTPException(status_code=500, detail=f"Echec wipe: {r.text}")
    return {"ok": True, "collection": coll, "wiped": True}

@fastapi_app.post("/index")
async def index(req: IndexRequest):
    if not QDRANT_URL or not QDRANT_API_KEY:
        raise HTTPException(status_code=400, detail="QDRANT_URL / QDRANT_API_KEY requis")
    job_id = uuid.uuid4().hex[:12]
    job = JobState(job_id=job_id, project_id=req.project_id)
    JOBS[job_id] = job
    # Lance en tâche de fond
    asyncio.create_task(run_index_job(job, req))
    job.log(f"Job {job_id} créé pour project {req.project_id}")
    return {"job_id": job_id, "project_id": req.project_id}

@fastapi_app.get("/status/{job_id}")
async def status(job_id: str):
    job = JOBS.get(job_id)
    if not job:
        raise HTTPException(status_code=404, detail="job_id inconnu")
    return job.model_dump()

@fastapi_app.get("/collections/{project_id}/count")
async def coll_count(project_id: str):
    if not QDRANT_URL or not QDRANT_API_KEY:
        raise HTTPException(status_code=400, detail="QDRANT_URL / QDRANT_API_KEY requis")
    coll = f"{COLLECTION_PREFIX}{project_id}"
    async with httpx.AsyncClient() as client:
        cnt = await qdrant_count(client, coll)
    return {"project_id": project_id, "collection": coll, "count": cnt}

@fastapi_app.post("/query")
async def query(req: QueryRequest):
    if not QDRANT_URL or not QDRANT_API_KEY:
        raise HTTPException(status_code=400, detail="QDRANT_URL / QDRANT_API_KEY requis")
    coll = f"{COLLECTION_PREFIX}{req.project_id}"
    async with httpx.AsyncClient() as client:
        vec = (await embed_texts(client, [req.text]))[0]
        data = await qdrant_search(client, coll, vec, limit=req.top_k)
    return data

# ------------------------------------------------------------------------------
# Gradio UI (montée sur "/")
# ------------------------------------------------------------------------------
def _default_two_docs() -> List[Dict[str, str]]:
    a = "Alpha bravo charlie delta echo foxtrot golf hotel india. " * 3
    b = "Lorem ipsum dolor sit amet, consectetuer adipiscing elit, sed diam nonummy." * 3
    return [
        {"path": "a.txt", "text": a},
        {"path": "b.txt", "text": b},
    ]

async def ui_wipe(project: str):
    try:
        resp = await wipe(project)  # appelle la route interne
        return f"✅ Wipe ok — collection {resp['collection']} supprimée."
    except Exception as e:
        LOG.exception("wipe UI error")
        return f"❌ Wipe erreur: {e}"

async def ui_index_sample(project: str, chunk_size: int, overlap: int, batch_size: int, store_text: bool):
    files = _default_two_docs()
    req = IndexRequest(
        project_id=project,
        files=[FileItem(**f) for f in files],
        chunk_size=chunk_size,
        overlap=overlap,
        batch_size=batch_size,
        store_text=store_text,
    )
    try:
        data = await index(req)
        job_id = data["job_id"]
        return f"🚀 Job lancé: {job_id}"
    except ValidationError as ve:
        return f"❌ Payload invalide: {ve}"
    except Exception as e:
        LOG.exception("index UI error")
        return f"❌ Index erreur: {e}"

async def ui_status(job_id: str):
    if not job_id.strip():
        return "⚠️ Renseigne un job_id"
    try:
        st = await status(job_id)
        # Formatage
        lines = [f"Job {st['job_id']} — stage={st['stage']} files={st['total_files']} chunks={st['total_chunks']} embedded={st['embedded']} upserted={st['upserted']}"]
        lines += st.get("messages", [])[-50:]  # dernières lignes
        if st.get("errors"):
            lines.append("Erreurs:")
            lines += [f" - {e}" for e in st["errors"]]
        return "\n".join(lines)
    except Exception as e:
        return f"❌ Status erreur: {e}"

async def ui_count(project: str):
    try:
        resp = await coll_count(project)
        return f"📊 Count — collection={resp['collection']} → {resp['count']} points"
    except Exception as e:
        LOG.exception("count UI error")
        return f"❌ Count erreur: {e}"

async def ui_query(project: str, text: str, topk: int):
    try:
        data = await query(QueryRequest(project_id=project, text=text, top_k=topk))
        hits = data.get("result", [])
        if not hits:
            return "Aucun résultat."
        out = []
        for h in hits:
            score = h.get("score")
            payload = h.get("payload", {})
            path = payload.get("path")
            chunk = payload.get("chunk")
            preview = (payload.get("text") or "")[:120].replace("\n", " ")
            out.append(f"{score:.4f} — {path} [chunk {chunk}] — {preview}…")
        return "\n".join(out)
    except Exception as e:
        LOG.exception("query UI error")
        return f"❌ Query erreur: {e}"

with gr.Blocks(title="Remote Indexer - Minimal Test", analytics_enabled=False) as ui:
    gr.Markdown("## 🔬 Remote Indexer — Tests sans console\n"
                "Wipe → Index 2 docs → Status → Count → Query\n"
                f"- **Embeddings**: `{EMB_PROVIDER}` (model: `{HF_EMBED_MODEL}`)\n"
                f"- **Qdrant**: `{'OK' if QDRANT_URL else 'ABSENT'}`\n"
                "Conseil: si tu n'as pas de token HF, mets `EMB_PROVIDER=dummy` dans les Secrets du Space.")
    with gr.Row():
        project_tb = gr.Textbox(label="Project ID", value="DEEPWEB")
        jobid_tb = gr.Textbox(label="Job ID (pour Status)", value="", interactive=True)
    with gr.Row():
        wipe_btn = gr.Button("🧨 Wipe collection", variant="stop")
        index_btn = gr.Button("🚀 Indexer 2 documents", variant="primary")
        count_btn = gr.Button("📊 Count points", variant="secondary")
    with gr.Row():
        chunk_size = gr.Slider(64, 1024, value=200, step=8, label="chunk_size")
        overlap = gr.Slider(0, 256, value=20, step=2, label="overlap")
        batch_size = gr.Slider(1, 128, value=32, step=1, label="batch_size")
        store_text = gr.Checkbox(value=True, label="store_text (payload)")
    out_log = gr.Textbox(lines=18, label="Logs / Résultats", interactive=False)
    with gr.Row():
        query_tb = gr.Textbox(label="Query text", value="alpha bravo")
        topk = gr.Slider(1, 20, value=5, step=1, label="top_k")
        query_btn = gr.Button("🔎 Query")
    query_out = gr.Textbox(lines=10, label="Résultats Query", interactive=False)

    wipe_btn.click(ui_wipe, inputs=[project_tb], outputs=[out_log])
    index_btn.click(ui_index_sample, inputs=[project_tb, chunk_size, overlap, batch_size, store_text], outputs=[out_log])
    # Petit auto-poll status: on relance ui_status à la main en collant le job_id
    count_btn.click(ui_count, inputs=[project_tb], outputs=[out_log])
    query_btn.click(ui_query, inputs=[project_tb, query_tb, topk], outputs=[query_out])

# Monte l'UI Gradio sur la FastAPI
app = gr.mount_gradio_app(fastapi_app, ui, path="/")