# -*- coding: utf-8 -*-
from __future__ import annotations
import os, time, uuid, logging, random
from typing import List, Optional, Dict, Any, Tuple

import numpy as np
import requests
from fastapi import FastAPI, BackgroundTasks, Header, HTTPException, Query
from pydantic import BaseModel, Field
from qdrant_client import QdrantClient
from qdrant_client.http.models import VectorParams, Distance, PointStruct

# ---------- logging ----------
logging.basicConfig(level=logging.INFO, format="%(levelname)s:%(name)s:%(message)s")
LOG = logging.getLogger("remote_indexer")

# ---------- ENV (config) ----------
# Ordre des backends d'embeddings à essayer (séparés par des virgules). Ex: "hf,deepinfra"
EMB_BACKEND_ORDER = [s.strip().lower() for s in os.getenv("EMB_BACKEND_ORDER", os.getenv("EMB_BACKEND", "hf")).split(",") if s.strip()]

# HF Inference API
HF_TOKEN   = os.getenv("HF_API_TOKEN", "").strip()
HF_MODEL   = os.getenv("HF_EMBED_MODEL", "sentence-transformers/all-MiniLM-L6-v2").strip()

# URLs configurables
HF_API_URL_USER      = os.getenv("HF_API_URL", "").strip()
HF_API_URL_PIPELINE  = os.getenv("HF_API_URL_PIPELINE", "").strip()
HF_API_URL_MODELS    = os.getenv("HF_API_URL_MODELS", "").strip()

if HF_API_URL_USER:
    if "/pipeline" in HF_API_URL_USER:
        HF_API_URL_PIPELINE = HF_API_URL_USER
    else:
        HF_API_URL_MODELS = HF_API_URL_USER

# Défaults
HF_URL_PIPELINE = (HF_API_URL_PIPELINE or f"https://api-inference.huggingface.co/pipeline/feature-extraction/{HF_MODEL}")
HF_URL_MODELS   = (HF_API_URL_MODELS   or f"https://api-inference.huggingface.co/models/{HF_MODEL}")

HF_TIMEOUT = float(os.getenv("EMB_TIMEOUT_SEC", "120"))
HF_WAIT    = os.getenv("HF_WAIT_FOR_MODEL", "true").lower() in ("1","true","yes","on")
HF_PIPELINE_FIRST = os.getenv("HF_PIPELINE_FIRST", "true").lower() in ("1","true","yes","on")

# DeepInfra Embeddings (OpenAI-like)
DI_TOKEN   = os.getenv("DEEPINFRA_API_KEY", "").strip()
DI_MODEL   = os.getenv("DEEPINFRA_EMBED_MODEL", "thenlper/gte-small").strip()
DI_URL     = os.getenv("DEEPINFRA_EMBED_URL", "https://api.deepinfra.com/v1/embeddings").strip()
DI_TIMEOUT = float(os.getenv("EMB_TIMEOUT_SEC", "120"))

# Retries
RETRY_MAX      = int(os.getenv("EMB_RETRY_MAX", "6"))         # tentatives max par backend
RETRY_BASE_SEC = float(os.getenv("EMB_RETRY_BASE", "1.5"))    # backoff de base (exponentiel)
RETRY_JITTER   = float(os.getenv("EMB_RETRY_JITTER", "0.35")) # jitter fraction (0..1)

# Qdrant
QDRANT_URL = os.getenv("QDRANT_URL", "http://localhost:6333")
QDRANT_API = os.getenv("QDRANT_API_KEY", "").strip()

# Auth d’API du service (simple header)
AUTH_TOKEN = os.getenv("REMOTE_INDEX_TOKEN", "").strip()

LOG.info(f"Embeddings backend order = {EMB_BACKEND_ORDER}")
LOG.info(f"HF pipeline URL = {HF_URL_PIPELINE}")
LOG.info(f"HF models   URL = {HF_URL_MODELS}")
if "hf" in EMB_BACKEND_ORDER and not HF_TOKEN:
    LOG.warning("HF_API_TOKEN manquant — tentatives HF échoueront.")
if "deepinfra" in EMB_BACKEND_ORDER and not DI_TOKEN:
    LOG.warning("DEEPINFRA_API_KEY manquant — tentatives DeepInfra échoueront.")

# ---------- Clients ----------
try:
    qdr = QdrantClient(url=QDRANT_URL, api_key=QDRANT_API if QDRANT_API else None)
except Exception as e:
    LOG.warning(f"Qdrant client init: {e}")

# ---------- Pydantic ----------
class FileIn(BaseModel):
    path: str
    text: str

class IndexRequest(BaseModel):
    project_id: str = Field(..., min_length=1)
    files: List[FileIn]
    chunk_size: int = 1200
    overlap: int = 200
    batch_size: int = 8
    store_text: bool = True

class QueryRequest(BaseModel):
    project_id: str
    query: str
    top_k: int = 6

# ---------- Jobs store (mémoire) ----------
JOBS: Dict[str, Dict[str, Any]] = {}  # {job_id: {"status": "...", "logs": [...], "created": ts}}

def _append_log(job_id: str, line: str):
    job = JOBS.get(job_id)
    if job: job["logs"].append(line)

def _set_status(job_id: str, status: str):
    job = JOBS.get(job_id)
    if job: job["status"] = status

def _auth(x_auth: Optional[str]):
    if AUTH_TOKEN and (x_auth or "") != AUTH_TOKEN:
        raise HTTPException(status_code=401, detail="Unauthorized")

# ---------- Embeddings backends avec retry ----------
def _retry_sleep(attempt: int):
    # backoff exponentiel + jitter
    back = (RETRY_BASE_SEC ** attempt)
    jitter = 1.0 + random.uniform(-RETRY_JITTER, RETRY_JITTER)
    return max(0.25, back * jitter)

def _with_task_param(url: str, task: str = "feature-extraction") -> str:
    # Ajoute ?task=feature-extraction (ou &task=...) si absent
    return url + ("&" if "?" in url else "?") + f"task={task}"

def _hf_http(url: str, payload: Dict[str, Any], headers_extra: Optional[Dict[str, str]] = None) -> Tuple[np.ndarray, int]:
    if not HF_TOKEN:
        raise RuntimeError("HF_API_TOKEN manquant (backend=hf).")

    headers = {
        "Authorization": f"Bearer {HF_TOKEN}",
        "Content-Type": "application/json",
        "Accept": "application/json",
    }
    # options.wait_for_model dans le JSON + X-Wait-For-Model côté header -> compat maximale
    if HF_WAIT:
        payload.setdefault("options", {})["wait_for_model"] = True
        headers["X-Wait-For-Model"] = "true"
        headers["X-Use-Cache"] = "true"

    if headers_extra:
        headers.update(headers_extra)

    r = requests.post(url, headers=headers, json=payload, timeout=HF_TIMEOUT)
    size = int(r.headers.get("Content-Length", "0"))
    if r.status_code >= 400:
        LOG.error(f"HF error {r.status_code}: {r.text[:1000]}")
        r.raise_for_status()

    data = r.json()
    arr = np.array(data, dtype=np.float32)
    # data peut être: [tokens, dim], [batch, tokens, dim], [batch, dim], [dim]
    if arr.ndim == 3:   # [batch, tokens, dim]
        arr = arr.mean(axis=1)
    elif arr.ndim == 2:
        pass
    elif arr.ndim == 1: # [dim] -> [1, dim]
        arr = arr.reshape(1, -1)
    else:
        raise RuntimeError(f"HF: unexpected embeddings shape: {arr.shape}")

    # normalisation L2
    norms = np.linalg.norm(arr, axis=1, keepdims=True) + 1e-12
    arr = arr / norms
    return arr.astype(np.float32), size

def _hf_post_embeddings_once(batch: List[str]) -> Tuple[np.ndarray, int]:
    """
    1) Essaie PIPELINE feature-extraction (si dispo)
    2) Fallback MODELS + header X-Task: feature-extraction
    3) Si encore 400 à cause de SentenceSimilarityPipeline, force aussi ?task=feature-extraction sur l'URL MODELS
    """
    payload: Dict[str, Any] = {"inputs": (batch if len(batch) > 1 else batch[0])}

    urls = [HF_URL_PIPELINE, HF_URL_MODELS] if HF_PIPELINE_FIRST else [HF_URL_MODELS, HF_URL_PIPELINE]
    last_exc: Optional[Exception] = None

    for idx, url in enumerate(urls, 1):
        try:
            if "/models/" in url:
                # 2) MODELS avec header X-Task
                return _hf_http(url, payload, headers_extra={"X-Task": "feature-extraction"})
            else:
                # 1) PIPELINE
                return _hf_http(url, payload, headers_extra=None)
        except requests.HTTPError as he:
            code = he.response.status_code if he.response is not None else 0
            body = he.response.text if he.response is not None else ""
            last_exc = he
            if code in (404, 405, 501) and idx < len(urls):
                LOG.warning(f"HF endpoint {url} non dispo ({code}), fallback vers alternative ...")
                continue
            # Si on a tapé MODELS et reçu SentenceSimilarityPipeline -> réessaie avec ?task=feature-extraction
            if "/models/" in url and "SentenceSimilarityPipeline" in (body or ""):
                try:
                    forced_url = _with_task_param(url, "feature-extraction")
                    LOG.warning("HF MODELS a choisi Similarity -> retry avec %s + X-Task", forced_url)
                    return _hf_http(forced_url, payload, headers_extra={"X-Task": "feature-extraction"})
                except Exception as he2:
                    last_exc = he2
            raise
        except Exception as e:
            last_exc = e
            raise

    # ne devrait pas arriver
    raise RuntimeError(f"HF: aucun endpoint utilisable ({last_exc})")

def _di_post_embeddings_once(batch: List[str]) -> Tuple[np.ndarray, int]:
    if not DI_TOKEN:
        raise RuntimeError("DEEPINFRA_API_KEY manquant (backend=deepinfra).")
    headers = {"Authorization": f"Bearer {DI_TOKEN}", "Content-Type": "application/json", "Accept": "application/json"}
    payload = {"model": DI_MODEL, "input": batch}
    r = requests.post(DI_URL, headers=headers, json=payload, timeout=DI_TIMEOUT)
    size = int(r.headers.get("Content-Length", "0"))
    if r.status_code >= 400:
        LOG.error(f"DeepInfra error {r.status_code}: {r.text[:1000]}")
        r.raise_for_status()
    js = r.json()
    data = js.get("data")
    if not isinstance(data, list) or not data:
        raise RuntimeError(f"DeepInfra embeddings: réponse invalide {js}")
    embs = [d.get("embedding") for d in data]
    arr = np.asarray(embs, dtype=np.float32)
    if arr.ndim != 2:
        raise RuntimeError(f"DeepInfra: unexpected embeddings shape: {arr.shape}")
    norms = np.linalg.norm(arr, axis=1, keepdims=True) + 1e-12
    arr = arr / norms
    return arr.astype(np.float32), size

def _call_with_retries(func, batch: List[str], label: str, job_id: Optional[str] = None) -> Tuple[np.ndarray, int]:
    last_exc = None
    for attempt in range(RETRY_MAX):
        try:
            if job_id:
                _append_log(job_id, f"{label}: try {attempt+1}/{RETRY_MAX} (batch={len(batch)})")
            return func(batch)
        except requests.HTTPError as he:
            code = he.response.status_code if he.response is not None else "HTTP"
            retriable = code in (429, 500, 502, 503, 504)
            if not retriable:
                raise
            sleep_s = _retry_sleep(attempt)
            msg = f"{label}: HTTP {code}, retry in {sleep_s:.1f}s"
            LOG.warning(msg); _append_log(job_id, msg)
            time.sleep(sleep_s)
            last_exc = he
        except Exception as e:
            sleep_s = _retry_sleep(attempt)
            msg = f"{label}: error {type(e).__name__}: {e}, retry in {sleep_s:.1f}s"
            LOG.warning(msg); _append_log(job_id, msg)
            time.sleep(sleep_s)
            last_exc = e
    raise RuntimeError(f"{label}: retries exhausted: {last_exc}")

def _post_embeddings(batch: List[str], job_id: Optional[str] = None) -> Tuple[np.ndarray, int]:
    """
    Essaie les backends dans EMB_BACKEND_ORDER avec retries.
    Ex: EMB_BACKEND_ORDER=hf,deepinfra
    """
    last_err = None
    for b in EMB_BACKEND_ORDER:
        if b == "hf":
            try:
                return _call_with_retries(_hf_post_embeddings_once, batch, "HF", job_id)
            except Exception as e:
                last_err = e
                _append_log(job_id, f"HF failed: {e}.")
                LOG.error(f"HF failed: {e}")
        elif b == "deepinfra":
            try:
                return _call_with_retries(_di_post_embeddings_once, batch, "DeepInfra", job_id)
            except Exception as e:
                last_err = e
                _append_log(job_id, f"DeepInfra failed: {e}.")
                LOG.error(f"DeepInfra failed: {e}")
        else:
            _append_log(job_id, f"Backend inconnu ignoré: {b}")
    raise RuntimeError(f"Tous les backends ont échoué: {last_err}")

# ---------- Qdrant helpers ----------
def _ensure_collection(name: str, dim: int):
    try:
        qdr.get_collection(name); return
    except Exception:
        pass
    qdr.create_collection(
        collection_name=name,
        vectors_config=VectorParams(size=dim, distance=Distance.COSINE),
    )

def _chunk_with_spans(text: str, size: int, overlap: int):
    n = len(text or "")
    if size <= 0:
        yield (0, n, text); return
    i = 0
    while i < n:
        j = min(n, i + size)
        yield (i, j, text[i:j])
        i = max(0, j - overlap)
        if i >= n: break

# ---------- Background task ----------
def run_index_job(job_id: str, req: IndexRequest):
    try:
        _set_status(job_id, "running")
        total_chunks = 0
        _append_log(job_id, f"Start project={req.project_id} files={len(req.files)} | backends={EMB_BACKEND_ORDER}")
        LOG.info(f"[{job_id}] Index start project={req.project_id} files={len(req.files)}")

        # Warmup -> dimension (1er morceau non vide si possible)
        warm = "warmup"
        if req.files:
            for _, _, chunk_txt in _chunk_with_spans(req.files[0].text or "", req.chunk_size, req.overlap):
                if (chunk_txt or "").strip():
                    warm = chunk_txt; break
        embs, sz = _post_embeddings([warm], job_id=job_id)
        dim = embs.shape[1]
        col = f"proj_{req.project_id}"
        _ensure_collection(col, dim)
        _append_log(job_id, f"Collection ready: {col} (dim={dim})")

        point_id = 0
        # Boucle sur les fichiers
        for fi, f in enumerate(req.files, 1):
            if not (f.text or "").strip():
                _append_log(job_id, f"file {fi}: vide — ignoré")
                continue
            chunks, metas = [], []
            for ci, (start, end, chunk_txt) in enumerate(_chunk_with_spans(f.text, req.chunk_size, req.overlap)):
                if not (chunk_txt or "").strip():
                    continue
                chunks.append(chunk_txt)
                meta = {"path": f.path, "chunk": ci, "start": start, "end": end}
                if req.store_text:
                    meta["text"] = chunk_txt
                metas.append(meta)
                # flush par lots
                if len(chunks) >= req.batch_size:
                    vecs, sz = _post_embeddings(chunks, job_id=job_id)
                    batch_points = [
                        PointStruct(id=point_id + k, vector=vec.tolist(), payload=metas[k])
                        for k, vec in enumerate(vecs)
                    ]
                    qdr.upsert(collection_name=col, points=batch_points)
                    point_id += len(batch_points)
                    total_chunks += len(chunks)
                    _append_log(job_id, f"file {fi}/{len(req.files)}: +{len(chunks)} chunks (total={total_chunks}) ~{sz/1024:.1f}KiB")
                    chunks, metas = [], []

            # flush fin de fichier
            if chunks:
                vecs, sz = _post_embeddings(chunks, job_id=job_id)
                batch_points = [
                    PointStruct(id=point_id + k, vector=vec.tolist(), payload=metas[k])
                    for k, vec in enumerate(vecs)
                ]
                qdr.upsert(collection_name=col, points=batch_points)
                point_id += len(batch_points)
                total_chunks += len(chunks)
                _append_log(job_id, f"file {fi}/{len(req.files)}: +{len(chunks)} chunks (total={total_chunks}) ~{sz/1024:.1f}KiB")

        _append_log(job_id, f"Done. chunks={total_chunks}")
        _set_status(job_id, "done")
        LOG.info(f"[{job_id}] Index finished. chunks={total_chunks}")
    except Exception as e:
        LOG.exception("Index job failed")
        _append_log(job_id, f"ERROR: {e}")
        _set_status(job_id, "error")

# ---------- API ----------
app = FastAPI()

@app.get("/")
def root():
    return {
        "ok": True,
        "service": "remote-indexer",
        "backends": EMB_BACKEND_ORDER,
        "hf_url_pipeline": HF_URL_PIPELINE if "hf" in EMB_BACKEND_ORDER else None,
        "hf_url_models": HF_URL_MODELS if "hf" in EMB_BACKEND_ORDER else None,
        "di_model": DI_MODEL if "deepinfra" in EMB_BACKEND_ORDER else None,
        "docs": "/health, /index, /status/{job_id}, /query, /wipe"
    }

@app.get("/health")
def health():
    return {"ok": True}

def _check_backend_ready():
    if "hf" in EMB_BACKEND_ORDER and not HF_TOKEN:
        raise HTTPException(400, "HF_API_TOKEN manquant côté serveur (backend=hf).")
    if "deepinfra" in EMB_BACKEND_ORDER and not DI_TOKEN and EMB_BACKEND_ORDER == ["deepinfra"]:
        raise HTTPException(400, "DEEPINFRA_API_KEY manquant côté serveur (backend=deepinfra).")

@app.post("/index")
def start_index(req: IndexRequest, background_tasks: BackgroundTasks, x_auth_token: Optional[str] = Header(default=None)):
    if AUTH_TOKEN and (x_auth_token or "") != AUTH_TOKEN:
        raise HTTPException(401, "Unauthorized")
    _check_backend_ready()

    # Filtrage défensif des fichiers vides pour éviter 422
    non_empty = [f for f in req.files if (f.text or "").strip()]
    if not non_empty:
        raise HTTPException(422, "Aucun fichier non vide à indexer.")
    req.files = non_empty

    job_id = uuid.uuid4().hex[:12]
    JOBS[job_id] = {"status": "queued", "logs": [], "created": time.time()}
    background_tasks.add_task(run_index_job, job_id, req)
    return {"job_id": job_id}

@app.get("/status/{job_id}")
def status(job_id: str, x_auth_token: Optional[str] = Header(default=None)):
    if AUTH_TOKEN and (x_auth_token or "") != AUTH_TOKEN:
        raise HTTPException(401, "Unauthorized")
    j = JOBS.get(job_id)
    if not j:
        raise HTTPException(404, "job inconnu")
    return {"status": j["status"], "logs": j["logs"][-800:]}

# --- Compat endpoints (clients legacy) ---
@app.get("/status")
def status_qp(job_id: str = Query(None), x_auth_token: Optional[str] = Header(default=None)):
    if AUTH_TOKEN and (x_auth_token or "") != AUTH_TOKEN:
        raise HTTPException(401, "Unauthorized")
    if not job_id:
        raise HTTPException(404, "job inconnu")
    j = JOBS.get(job_id)
    if not j:
        raise HTTPException(404, "job inconnu")
    return {"status": j["status"], "logs": j["logs"][-800:]}

class _StatusBody(BaseModel):
    job_id: str

@app.post("/status")
def status_post(body: _StatusBody, x_auth_token: Optional[str] = Header(default=None)):
    if AUTH_TOKEN and (x_auth_token or "") != AUTH_TOKEN:
        raise HTTPException(401, "Unauthorized")
    j = JOBS.get(body.job_id)
    if not j:
        raise HTTPException(404, "job inconnu")
    return {"status": j["status"], "logs": j["logs"][-800:]}

@app.post("/query")
def query(req: QueryRequest, x_auth_token: Optional[str] = Header(default=None)):
    if AUTH_TOKEN and (x_auth_token or "") != AUTH_TOKEN:
        raise HTTPException(401, "Unauthorized")
    _check_backend_ready()

    # bornes du top_k
    k = int(max(1, min(50, req.top_k or 6)))

    vecs, _ = _post_embeddings([req.query])
    col = f"proj_{req.project_id}"
    try:
        res = qdr.search(collection_name=col, query_vector=vecs[0].tolist(), limit=k)
    except Exception as e:
        raise HTTPException(400, f"Search failed: {e}")
    out = []
    for p in res:
        pl = p.payload or {}
        txt = pl.get("text")
        if txt and len(txt) > 800:
            txt = txt[:800] + "..."
        out.append({
            "path": pl.get("path"),
            "chunk": pl.get("chunk"),
            "start": pl.get("start"),
            "end": pl.get("end"),
            "text": txt,
            "score": float(p.score) if hasattr(p, "score") else None,
        })
    return {"results": out}

@app.post("/wipe")
def wipe_collection(project_id: str, x_auth_token: Optional[str] = Header(default=None)):
    if AUTH_TOKEN and (x_auth_token or "") != AUTH_TOKEN:
        raise HTTPException(401, "Unauthorized")
    col = f"proj_{project_id}"
    try:
        qdr.delete_collection(col); return {"ok": True}
    except Exception as e:
        raise HTTPException(400, f"wipe failed: {e}")

# ---------- Entrypoint ----------
if __name__ == "__main__":
    import uvicorn
    port = int(os.getenv("PORT", "7860"))
    LOG.info(f"===== Application Startup on PORT {port} =====")
    uvicorn.run(app, host="0.0.0.0", port=port)