"""
Dify 知识库客户端
================
封装与 Dify 知识库(v1 API)的交互：连通性测试、检索(RAG)、文本入库。
设计用于替换原 FastGPT 外部知识库对接。

已验证的接入点（2026-07-14 实测）:
- Base:  https://dify.qihuyun.com/v1
         （直接 HTTP 访问会被 nginx 301 跳转到 HTTPS，必须 HTTPS 直连，
          且内网自签证书需 verify=False）
- Auth:  Bearer {api_key}
- GET  /v1/datasets                               列出知识库(可据此解析 dataset id)
- POST /v1/datasets/{id}/document/create-by-text  文本入库
- POST /v1/datasets/{id}/document/create-by-file  文件入库(multipart)
- POST /v1/datasets/{id}/retrieve                 检索(RAG)
  retrieve 必需完整的 retrieval_model 字典，否则返回 400。
"""
import logging
import threading
import time
import httpx

logger = logging.getLogger("aiflow.dify_kb")

# dataset id 缓存: key -> (dataset_id, ts)
_dataset_cache: dict = {}
_dataset_cache_lock = threading.Lock()
_DATASET_CACHE_TTL = 3600


def _normalize_base(api_url: str) -> str:
    """统一 base url -> https 形式，且不带 /v1 后缀。"""
    u = (api_url or "").strip()
    if not u:
        u = "https://dify.qihuyun.com/v1"
    u = u.rstrip("/")
    if u.endswith("/v1"):
        u = u[: -len("/v1")]
    if u.startswith("http://"):
        u = "https://" + u[len("http://"):]
    return u


def dify_test(api_url: str, api_key: str) -> dict:
    """验证 Dify 知识库「可通讯 + 可用」。

    返回:
        {
          "connected": bool,
          "usable": bool,
          "status": "connected" | "auth_error" | "timeout" | "error",
          "detail": str,
          "datasets": [{"id":..., "name":...}, ...]
        }
    """
    base = _normalize_base(api_url)
    headers = {"Authorization": f"Bearer {api_key}"}
    try:
        with httpx.Client(base_url=base, timeout=15, verify=False) as client:
            r = client.get("/v1/datasets", headers=headers)
            if r.status_code == 200:
                data = r.json()
                datasets = data.get("data", []) or []
                names = [d.get("name") for d in datasets]
                # 能列出知识库即代表鉴权通过、可读取 —— 视为可用。
                return {
                    "connected": True,
                    "usable": True,
                    "status": "connected",
                    "detail": f"Dify 连接成功，发现 {len(datasets)} 个知识库: {', '.join([n for n in names if n]) or '无'}",
                    "datasets": [{"id": d.get("id"), "name": d.get("name")} for d in datasets],
                }
            elif r.status_code == 401:
                return {
                    "connected": True,
                    "usable": False,
                    "status": "auth_error",
                    "detail": "已连通 Dify，但 API Key 鉴权失败(401)，请检查 API Key。",
                    "datasets": [],
                }
            else:
                return {
                    "connected": False,
                    "usable": False,
                    "status": "error",
                    "detail": f"Dify 返回状态码 {r.status_code}",
                    "datasets": [],
                }
    except httpx.TimeoutException:
        return {"connected": False, "usable": False, "status": "timeout", "detail": "连接超时", "datasets": []}
    except Exception as e:
        return {"connected": False, "usable": False, "status": "error", "detail": str(e), "datasets": []}


def resolve_dataset_id(api_url: str, api_key: str, dataset_name: str = None) -> str:
    """解析 dataset id（优先按名称匹配，否则取列表第一个），带缓存。"""
    base = _normalize_base(api_url)
    headers = {"Authorization": f"Bearer {api_key}"}
    cache_key = f"{base}|{api_key}|{dataset_name}"
    now = time.time()
    with _dataset_cache_lock:
        c = _dataset_cache.get(cache_key)
        if c and now - c[1] < _DATASET_CACHE_TTL:
            return c[0]
    ds_id = None
    try:
        with httpx.Client(base_url=base, timeout=15, verify=False) as client:
            r = client.get("/v1/datasets", headers=headers)
            if r.status_code == 200:
                datasets = r.json().get("data", []) or []
                if dataset_name:
                    match = [d for d in datasets if d.get("name") == dataset_name]
                    if match:
                        ds_id = match[0].get("id")
                if not ds_id and datasets:
                    ds_id = datasets[0].get("id")
    except Exception as e:
        logger.error(f"resolve_dataset_id failed: {e}")
    if ds_id:
        with _dataset_cache_lock:
            _dataset_cache[cache_key] = (ds_id, now)
    return ds_id


def dify_retrieve(api_url, api_key, query, top_k=5, dataset_name=None) -> list:
    """检索，返回 [{content, score, source, metadata}]。"""
    ds_id = resolve_dataset_id(api_url, api_key, dataset_name)
    if not ds_id:
        return []
    base = _normalize_base(api_url)
    headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}
    payload = {
        "query": query,
        "retrieval_model": {
            "search_method": "semantic_search",
            "reranking_enable": False,
            "score_threshold_enabled": False,
            "top_k": top_k,
        },
    }
    try:
        with httpx.Client(base_url=base, timeout=30, verify=False) as client:
            r = client.post(f"/v1/datasets/{ds_id}/retrieve", headers=headers, json=payload)
            if r.status_code == 200:
                data = r.json()
                out = []
                for item in data.get("records", []):
                    seg = item.get("segment", {})
                    out.append({
                        "content": seg.get("content", ""),
                        "score": float(item.get("score", 0.5)),
                        "source": f"Dify: {seg.get('document_id', '')}",
                        "metadata": item,
                    })
                return out
            else:
                logger.warning(f"Dify retrieve returned {r.status_code}: {r.text[:200]}")
                return []
    except Exception as e:
        logger.error(f"Dify retrieve error: {e}")
        return []


def dify_ingest_text(api_url, api_key, text, name=None, dataset_name=None) -> bool:
    """文本入库（create-by-text）。返回是否成功。"""
    if not text or not text.strip():
        return False
    ds_id = resolve_dataset_id(api_url, api_key, dataset_name)
    if not ds_id:
        return False
    base = _normalize_base(api_url)
    headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}
    payload = {
        "name": name or f"doc-{int(time.time())}",
        "text": text,
        "indexing_technique": "high_quality",
    }
    try:
        with httpx.Client(base_url=base, timeout=60, verify=False) as client:
            r = client.post(f"/v1/datasets/{ds_id}/document/create-by-text", headers=headers, json=payload)
            if r.status_code in (200, 201):
                return True
            else:
                logger.warning(f"Dify ingest text returned {r.status_code}: {r.text[:200]}")
                return False
    except Exception as e:
        logger.error(f"Dify ingest text error: {e}")
        return False


async def aingest_text(api_url, api_key, text, name=None, dataset_name=None) -> bool:
    """异步文本入库（供 chat 实时落库使用）。"""
    if not text or not text.strip():
        return False
    ds_id = resolve_dataset_id(api_url, api_key, dataset_name)
    if not ds_id:
        return False
    base = _normalize_base(api_url)
    headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}
    payload = {
        "name": name or f"doc-{int(time.time())}",
        "text": text,
        "indexing_technique": "high_quality",
    }
    try:
        async with httpx.AsyncClient(base_url=base, timeout=60, verify=False) as client:
            r = await client.post(f"/v1/datasets/{ds_id}/document/create-by-text", headers=headers, json=payload)
            if r.status_code in (200, 201):
                return True
            else:
                logger.warning(f"Dify aingest returned {r.status_code}: {r.text[:200]}")
                return False
    except Exception as e:
        logger.error(f"Dify aingest error: {e}")
        return False


def get_dify_kb_context(db, query: str, top_k: int = 5) -> str:
    """读取已启用的 Dify 外部服务，检索并返回拼接后的上下文文本（供 chat 注入）。"""
    from sqlalchemy import text as _text
    try:
        result = db.execute(_text(
            "SELECT api_url, api_key, config FROM external_kb_services "
            "WHERE service_type='dify' AND is_enabled=1 AND status='connected' LIMIT 1"
        ))
        row = result.fetchone()
        if not row:
            return ""
        api_url, api_key, config_raw = row
        dataset_name = None
        if config_raw:
            try:
                cfg = config_raw if isinstance(config_raw, dict) else __import__("json").loads(config_raw)
                dataset_name = cfg.get("dataset_name")
            except Exception:
                pass
        hits = dify_retrieve(api_url, api_key, query, top_k=top_k, dataset_name=dataset_name)
        if not hits:
            return ""
        parts = []
        for i, h in enumerate(hits, 1):
            content = (h.get("content") or "").strip()
            if content:
                parts.append(f"[{i}] {content}")
        return "\n\n".join(parts)
    except Exception as e:
        logger.error(f"get_dify_kb_context error: {e}")
        return ""