"""Few-shot 样本的 Qdrant 向量存储。

独立于 LightRAG 的 Qdrant 客户端，使用专用 collection ``few_shot_samples``，
与知识库 RAG 的 collection 隔离。所有操作失败都不抛异常（记日志返回空），
保证主链路不阻塞。

向量来自 :class:`EmbeddingProvider`，payload 带业务过滤字段（scene/label/domain/risk_type），
检索时按这些字段过滤 + 向量相似度排序。
"""

from __future__ import annotations

import os
import uuid
from typing import Any

from app.core.logging import get_logger
from app.services.knowledge_rag import _resolve_default_qdrant_url

logger = get_logger("app.services.few_shot_store")

FEW_SHOT_COLLECTION = "few_shot_samples"
FEW_SHOT_VECTOR_NAMESPACE = uuid.UUID("0ecb5868-cf75-47c1-b3bf-81e4a6ce1df3")


def stable_vector_id(*, tenant_id: str, sample_id: str) -> str:
    """同一租户样本始终映射到同一个 Qdrant point。"""

    tenant = str(tenant_id or "").strip()
    sample = str(sample_id or "").strip()
    if not tenant or not sample:
        raise ValueError("tenant_id and sample_id are required")
    return str(uuid.uuid5(FEW_SHOT_VECTOR_NAMESPACE, f"{tenant}:{sample}"))


def _resolve_qdrant_config() -> tuple[str, str]:
    """复用 knowledge_rag 的 Qdrant URL/key 解析逻辑。"""

    url = os.environ.get("QDRANT_URL", "").strip() or _resolve_default_qdrant_url()
    api_key = os.environ.get("QDRANT_API_KEY", "").strip()
    return url, api_key


class FewShotStore:
    """对 Qdrant 的轻量封装，专供 few-shot 样本检索使用。

    设计要点：
    - 惰性创建 client 和 collection，首次操作时初始化。
    - 所有公共方法吞异常（返回空/False），主链路永远不被拖崩。
    - 向量写入和检索都依赖外部传入的 :class:`EmbeddingProvider`，
      由调用方保证与配置一致。
    """

    def __init__(self, embedding_provider: Any) -> None:
        self._embedding_provider = embedding_provider
        self._client: Any = None
        self._ensured = False

    def _client_or_none(self) -> Any:
        """惰性初始化 QdrantClient，失败返回 None。"""

        if self._client is not None:
            return self._client
        try:
            from qdrant_client import QdrantClient

            url, api_key = _resolve_qdrant_config()
            self._client = QdrantClient(url=url, api_key=api_key or None)
        except Exception:
            logger.warning("few-shot QdrantClient 初始化失败，本轮操作跳过", exc_info=True)
            self._client = None
        return self._client

    def _ensure_collection(self) -> bool:
        """确保 collection 存在，成功返回 True。"""

        if self._ensured:
            return True
        client = self._client_or_none()
        if client is None:
            return False
        try:
            from qdrant_client.http.exceptions import UnexpectedResponse

            try:
                client.get_collection(FEW_SHOT_COLLECTION)
            except UnexpectedResponse as exc:
                if exc.status_code != 404:
                    raise
                dim = self._embedding_provider.dimension()
                from qdrant_client.http.models import Distance, VectorParams

                client.create_collection(
                    collection_name=FEW_SHOT_COLLECTION,
                    vectors_config=VectorParams(size=dim, distance=Distance.COSINE),
                )
                logger.info("few-shot collection 创建成功 dim=%s", dim)

            # 老 collection 也要补齐过滤索引，不能只在首次建表时创建。
            from qdrant_client.http.models import (
                PayloadSchemaType,
            )

            for field, field_type in [
                ("sample_id", PayloadSchemaType.KEYWORD),
                ("tenant_id", PayloadSchemaType.KEYWORD),
                ("scene", PayloadSchemaType.KEYWORD),
                ("policy_ref", PayloadSchemaType.KEYWORD),
                ("rule_version", PayloadSchemaType.KEYWORD),
                ("label", PayloadSchemaType.KEYWORD),
                ("domain", PayloadSchemaType.KEYWORD),
                ("risk_type", PayloadSchemaType.KEYWORD),
                ("status", PayloadSchemaType.KEYWORD),
            ]:
                try:
                    client.create_payload_index(
                        collection_name=FEW_SHOT_COLLECTION,
                        field_name=field,
                        field_schema=field_type,
                    )
                except Exception:
                    logger.debug("payload index 创建跳过 field=%s", field, exc_info=True)
            self._ensured = True
            return True
        except Exception:
            logger.warning("few-shot collection 初始化失败，本轮操作跳过", exc_info=True)
            return False

    def upsert(self, sample: Any) -> str | None:
        """把一条样本向量化并写入 Qdrant，返回 vector_id，失败返回 None。"""

        tenant_id = str(getattr(sample, "tenant_id", "") or "").strip()
        sample_id = str(getattr(sample, "id", "") or "").strip()
        if not tenant_id or not sample_id:
            logger.warning("few-shot upsert 缺少 tenant_id/sample_id，已拒绝")
            return None
        if not self._ensure_collection():
            return None
        client = self._client
        try:
            vector = self._embedding_provider.embed([sample.case_text])[0]
        except Exception:
            logger.warning(
                "few-shot embedding 失败 sample_key=%s",
                getattr(sample, "sample_key", ""),
                exc_info=True,
            )
            return None
        vector_id = stable_vector_id(tenant_id=tenant_id, sample_id=sample_id)
        previous_vector_id = str(getattr(sample, "vector_id", "") or "").strip()
        payload = {
            "sample_id": sample_id,
            "tenant_id": tenant_id,
            "scene": sample.scene,
            "policy_ref": getattr(sample, "policy_ref", "") or "",
            "rule_version": getattr(sample, "rule_version", "") or "",
            "label": sample.label,
            "domain": sample.domain,
            "risk_type": sample.risk_type,
            "risk_level": sample.risk_level,
            "status": getattr(sample, "status", "active"),
            "conclusion_text": sample.conclusion_text,
            "payload_json": sample.payload_json,
        }
        try:
            points = [{"id": vector_id, "vector": vector, "payload": payload}]
            if previous_vector_id and previous_vector_id != vector_id:
                # 先用最新判定覆盖旧 point，避免后续删除短暂失败时暴露陈旧标签。
                points.append({"id": previous_vector_id, "vector": vector, "payload": payload})
            client.upsert(
                collection_name=FEW_SHOT_COLLECTION,
                points=points,
            )
            if previous_vector_id and previous_vector_id != vector_id:
                try:
                    client.delete(
                        collection_name=FEW_SHOT_COLLECTION,
                        points_selector=[previous_vector_id],
                    )
                except Exception:
                    logger.warning(
                        "few-shot 旧向量清理失败，已保留同内容副本 vector_id=%s",
                        previous_vector_id,
                        exc_info=True,
                    )
            return vector_id
        except Exception:
            logger.warning(
                "few-shot upsert 失败 sample_key=%s",
                getattr(sample, "sample_key", ""),
                exc_info=True,
            )
            return None

    def search(
        self,
        case_text: str,
        *,
        tenant_id: str,
        scene: str | None = None,
        policy_ref: str | None = None,
        rule_version: str | None = None,
        labels: list[str] | None = None,
        top_k: int = 3,
    ) -> list[dict[str, Any]]:
        """按 case_text 检索相似样本，可按 scene/label 过滤。失败返回空列表。"""

        tenant = str(tenant_id or "").strip()
        if not tenant:
            raise ValueError("tenant_id is required for few-shot search")
        if not case_text or not self._ensure_collection():
            return []
        client = self._client
        try:
            vector = self._embedding_provider.embed([case_text])[0]
        except Exception:
            logger.warning("few-shot 检索 embedding 失败", exc_info=True)
            return []
        must: list[dict[str, Any]] = [
            {"key": "tenant_id", "match": {"value": tenant}},
            {"key": "status", "match": {"value": "active"}},
        ]
        if scene:
            must.append({"key": "scene", "match": {"value": scene}})
        if policy_ref:
            must.append({"key": "policy_ref", "match": {"value": policy_ref}})
        if rule_version:
            must.append({"key": "rule_version", "match": {"value": rule_version}})
        if labels:
            must.append({"key": "label", "match": {"any": labels}})
        try:
            from qdrant_client.http.models import Filter

            results = client.query_points(
                collection_name=FEW_SHOT_COLLECTION,
                query=vector,
                query_filter=Filter(must=must),
                limit=top_k,
                with_payload=True,
            ).points
        except Exception:
            logger.warning("few-shot 检索失败", exc_info=True)
            return []
        hits: list[dict[str, Any]] = []
        for point in results:
            payload = getattr(point, "payload", None) or {}
            hits.append(
                {
                    "sample_id": payload.get("sample_id"),
                    "tenant_id": payload.get("tenant_id"),
                    "scene": payload.get("scene"),
                    "policy_ref": payload.get("policy_ref") or "",
                    "rule_version": payload.get("rule_version") or "",
                    "score": float(getattr(point, "score", 0.0)),
                    "label": payload.get("label"),
                    "domain": payload.get("domain"),
                    "risk_type": payload.get("risk_type"),
                    "conclusion_text": payload.get("conclusion_text") or "",
                    "payload_json": payload.get("payload_json") or {},
                }
            )
        return hits

    def delete_by_vector_id(self, vector_id: str) -> bool:
        """按 vector_id 删除向量，失败返回 False。"""

        if not vector_id or not self._ensure_collection():
            return False
        try:
            self._client.delete(
                collection_name=FEW_SHOT_COLLECTION,
                points_selector=[vector_id],
            )
            return True
        except Exception:
            logger.warning("few-shot 删除失败 vector_id=%s", vector_id, exc_info=True)
            return False
