# -*- 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) ----------
# Par défaut on met DeepInfra d'abord pour être opérationnel.
DEFAULT_BACKENDS = "deepinfra,hf"
EMB_BACKEND_ORDER = [s.strip().lower() for s in os.getenv("EMB_BACKEND_ORDER", os.getenv("EMB_BACKEND", DEFAULT_BACKENDS)).split(",") if s.strip()]

# Auto-fallback vers DeepInfra si HF renvoie la fameuse erreur Similarity
ALLOW_DI_AUTOFALLBACK = os.getenv("ALLOW_DI_AUTOFALLBACK", "true").lower() in ("1","true","yes","on")

# 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()

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

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-compatible)
DI_TOKEN   = os.getenv("DEEPINFRA_API_KEY", "").strip()
# 👇 IMPORTANT : modèle existant chez DeepInfra (multilingue)
DI_MODEL   = os.getenv("DEEPINFRA_EMBED_MODEL", "BAAI/bge-m3").strip()
# 👇 IMPORTANT : endpoint OpenAI-compatible
DI_URL     = os.getenv("DEEPINFRA_EMBED_URL", "https://api.deepinfra.com/v1/openai/embeddings").strip()
DI_TIMEOUT = float(os.getenv("EMB_TIMEOUT_SEC", "120"))

# Retries
RETRY_MAX      = int(os.getenv("EMB_RETRY_MAX", "6"))
RETRY_BASE_SEC = float(os.getenv("EMB_RETRY_BASE", "1.5"))
RETRY_JITTER   = float(os.getenv("EMB_RETRY_JITTER", "0.35"))

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

# Auth
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 ----------
JOBS: Dict[str, Dict[str, Any]] = {}

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")

# ---------- Helpers retry ----------
def _retry_sleep(attempt: int):
    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:
    return url + ("&" if "?" in url else "?") + f"task={task}"

# ---------- HF embeddings ----------
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",
    }
    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)
    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}")

    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]:
    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:
                return _hf_http(url, payload, headers_extra={"X-Task": "feature-extraction"})
            else:
                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
            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
    raise RuntimeError(f"HF: aucun endpoint utilisable ({last_exc})")

# ---------- DeepInfra embeddings ----------
def _di_post_embeddings_once(batch: List[str]) -> Tuple[np.ndarray, int]:
    if not DI_TOKEN:
        raise RuntimeError("DEEPINFRA_API_KEY manquant (backend=deepinfra).")
    # OpenAI-compatible embeddings endpoint
    headers = {"Authorization": f"Bearer {DI_TOKEN}", "Content-Type": "application/json", "Accept": "application/json"}
    payload = {"model": DI_MODEL, "input": batch}
    # NB: on peut aussi ajouter "encoding_format":"float" si nécessaire
    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

# ---------- Retry orchestrator ----------
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.
    Auto-fallback optionnel vers DeepInfra si HF renvoie la Similarity.
    """
    last_err = None
    similarity_misroute = False

    for b in EMB_BACKEND_ORDER:
        if b == "hf":
            try:
                return _call_with_retries(_hf_post_embeddings_once, batch, "HF", job_id)
            except requests.HTTPError as he:
                body = he.response.text if getattr(he, "response", None) is not None else ""
                if "SentenceSimilarityPipeline.__call__()" in (body or ""):
                    similarity_misroute = True
                last_err = he
                _append_log(job_id, f"HF failed: {he}.")
                LOG.error(f"HF failed: {he}")
        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}")

    if ALLOW_DI_AUTOFALLBACK and similarity_misroute and DI_TOKEN:
        LOG.warning("HF a routé sur SentenceSimilarity => auto-fallback DeepInfra (override ordre).")
        _append_log(job_id, "Auto-fallback DeepInfra (HF => SentenceSimilarity).")
        return _call_with_retries(_di_post_embeddings_once, batch, "DeepInfra", job_id)

    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)}")

        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
        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():  # pas d'embeddings sur des blancs
                    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)
                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 = [], []

            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_url": DI_URL if "deepinfra" 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()
    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 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()
    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:
        qdrant.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)