#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
用户端知识库搜索 API（增强版）
支持：
1. 法条定位搜索 —— 《xxx法》第xx条 第xx款
2. 多集合向量搜索 —— 同时搜索 laws / discretion / cases
3. 全文语义搜索 —— 使用 bge-m3，最大 12000 字符输入
"""

import re
import logging
import httpx
from fastapi import APIRouter, Depends, Query, Request
from typing import List, Optional, Dict, Any
from qdrant_client import QdrantClient

from ...core.database import get_db
from ...core.permission import require_login
from ...core.config import settings
from ...core.llm_client import llm_client

router = APIRouter()

# ==================== 常量配置 ====================
EMBEDDING_MAX_INPUT_CHARS = 12000
EMBEDDING_PAYLOAD_MAX_CHARS = 12000
QDRANT_HOST = getattr(settings, "QDRANT_HOST", "127.0.0.1")
QDRANT_PORT = getattr(settings, "QDRANT_PORT", "6333")


# source → collections 映射
SOURCE_COLLECTION_MAP = {
    "laws": ["market_regulation_laws"],
    "discretion": ["discretion_standards"],
    "cases": ["case_precedents"],
    "exemptions": ["exemption_lists"],
    "all": ["market_regulation_laws", "discretion_standards", "case_precedents"],
}

# 中文数字转阿拉伯数字
CN_NUMBERS = {
    '零': 0, '一': 1, '二': 2, '三': 3, '四': 4,
    '五': 5, '六': 6, '七': 7, '八': 8, '九': 9,
    '十': 10, '百': 100, '千': 1000, '万': 10000
}


def chinese_to_arabic(cn: str) -> int:
    """将中文数字转为阿拉伯数字，如 '五十' -> 50, '一百零五' -> 105"""
    if not cn:
        return 0
    if cn.isdigit():
        return int(cn)
    total = 0
    current = 0
    for ch in cn:
        if ch in CN_NUMBERS:
            num = CN_NUMBERS[ch]
            if num >= 10:
                if current == 0:
                    current = 1
                total += current * num
                current = 0
            else:
                current = current * 10 + num if current > 0 else num
    return total + current


def parse_article_query(query: str) -> Dict[str, Any]:
    """
    解析法条定位查询
    支持格式：
      - 《食品安全法》第五十条
      - 食品安全法 第50条
      - 《广告法》第5条第1款（一）
      - 消费者权益保护法 第三十四条 第二款
    """
    result = {
        "law_name": None,
        "article_num": None,
        "article_arabic": None,
        "clause": None,
        "is_article_query": False,
    }

    # 1. 提取法律名称（《xxx》或 xxx法/条例/规定/办法）
    law_match = re.search(r'《([^》]+)》', query)
    if law_match:
        result["law_name"] = law_match.group(1).strip()

    # 2. 提取法条编号：第X条
    article_match = re.search(r'第\s*([一二三四五六七八九十百千万\d]+)\s*条', query)
    if article_match:
        num_str = article_match.group(1).strip()
        result["article_num"] = f"第{num_str}条"
        result["article_arabic"] = chinese_to_arabic(num_str)
        result["is_article_query"] = True

    # 3. 提取款项：第X款 或 （X）
    clause_match = re.search(r'第\s*([一二三四五六七八九十\d]+)\s*款', query)
    if not clause_match:
        clause_match = re.search(r'[（(]([一二三四五六七八九十\d]+)[)）]', query)
    if clause_match:
        result["clause"] = clause_match.group(1).strip()

    # 4. 没有《》但提到了法条，尝试提取法律名称
    if not result["law_name"] and result["is_article_query"]:
        # 取"第X条"之前的文本作为法律名称
        parts = re.split(r'第\s*[一二三四五六七八九十百千万\d]+\s*条', query)
        if parts and parts[0].strip():
            candidate = parts[0].strip().rstrip(' \n,，;；')
            # 只保留有意义的名称（至少包含"法/条例/规定/办法/规则"或长度>3）
            if len(candidate) > 3 or re.search(r'[法规条例规定办法规则]', candidate):
                result["law_name"] = candidate

    return result


def _get_collections(source: str) -> List[str]:
    """根据 source 参数获取要搜索的集合列表"""
    if not source or source == "all":
        return SOURCE_COLLECTION_MAP["all"]
    return SOURCE_COLLECTION_MAP.get(source, SOURCE_COLLECTION_MAP["all"])


def _build_article_filter(parsed: Dict[str, Any], law_name: Optional[str] = None) -> Optional[Dict]:
    """构建法条定位的 Qdrant filter"""
    if not parsed["is_article_query"]:
        return None

    must_conditions = []

    # 法律名称过滤（h1 字段存储的是法律名称）
    # 支持全称/简称匹配：如 "食品安全法" → 同时匹配 "中华人民共和国食品安全法"
    if parsed["law_name"] or law_name:
        name = (parsed["law_name"] or law_name).strip()
        h1_conditions = [{"key": "h1", "match": {"text": name}}]
        # 尝试全称变体（自动补全"中华人民共和国"前缀）
        if not name.startswith("中华人民共和国") and not name.startswith("中华"):
            h1_conditions.append({"key": "h1", "match": {"text": f"中华人民共和国{name}"}})
        if len(h1_conditions) > 1:
            must_conditions.append({"should": h1_conditions})
        else:
            must_conditions.append(h1_conditions[0])

    # 法条编号过滤（支持 article_num 和 article_number 两个字段）
    if parsed["article_num"]:
        must_conditions.append({
            "should": [
                {"key": "article_num", "match": {"text": parsed["article_num"]}},
                {"key": "article_number", "match": {"text": parsed["article_num"]}},
            ]
        })

    if must_conditions:
        return {"must": must_conditions}
    return None


async def _get_embedding(text: str) -> List[float]:
    """调用 Embedding API 生成向量"""
    result = await llm_client.embed(
        texts=[text[:EMBEDDING_MAX_INPUT_CHARS]],
    )
    if not result or not result[0]:
        raise RuntimeError("Embedding 服务返回空结果")
    return result[0]


# 共享的 Qdrant 客户端（使用官方 SDK）
_qdrant_sdk: Optional[QdrantClient] = None

def _get_qdrant_sdk() -> QdrantClient:
    """获取或创建 Qdrant SDK 客户端"""
    global _qdrant_sdk
    if _qdrant_sdk is None:
        # 尝试 gRPC 端口（6334）而不是 HTTP 端口
        _qdrant_sdk = QdrantClient(host=QDRANT_HOST, port=6334, prefer_grpc=True)
    return _qdrant_sdk


async def _search_qdrant_single(
    collection: str,
    vector: List[float],
    limit: int,
    filter_body: Optional[Dict] = None,
    region_filter: Optional[str] = None,
) -> List[Dict]:
    """在单个 Qdrant 集合中搜索（使用官方 SDK）"""
    search_body = {
        "vector": vector,
        "limit": limit,
        "with_payload": True,
    }

    # 构建 filter
    must_conditions = []
    if filter_body and filter_body.get("must"):
        must_conditions.extend(filter_body["must"])
    if region_filter:
        must_conditions.append({
            "key": "region",
            "match": {"text": region_filter}
        })

    try:
        sdk = _get_qdrant_sdk()
        search_result = sdk.query_points(
            collection_name=collection,
            query=vector,
            limit=limit,
            with_payload=True,
        )
        results = list(search_result.points)
    except Exception as e:
        logging.getLogger(__name__).warning(f"Qdrant 搜索失败 [{collection}]: {e}")
        raise

    formatted = []
    for r in results:
        payload = dict(r.payload) if r.payload else {}
        payload["_collection"] = collection
        formatted.append({
            "id": str(r.id),
            "score": r.score,
            "payload": payload,
        })
    return formatted


async def _search_qdrant_scroll(
    collection: str,
    limit: int,
    filter_body: Dict,
) -> List[Dict]:
    """使用 scroll API 进行精确 filter 查询（无法量时）"""
    scroll_body = {
        "limit": limit,
        "with_payload": True,
        "filter": filter_body,
    }

    async with httpx.AsyncClient() as client:
        qdrant_res = await client.post(
            f"http://{QDRANT_HOST}:{QDRANT_PORT}/collections/{collection}/points/scroll",
            json=scroll_body,
            timeout=30.0,
        )
        if qdrant_res.status_code != 200:
            return []

        search_data = qdrant_res.json()
        results = []
        for r in search_data.get("result", {}).get("points", []):
            payload = r.get("payload", {})
            payload["_collection"] = collection
            results.append({
                "id": r.get("id"),
                "score": 1.0,  # filter 查询没有 score，设为 1.0
                "payload": payload,
            })
        return results


async def _do_search(
    query: str,
    limit: int,
    offset: int,
    source: str,
    mode: str,
    user: dict,
    db,
    rerank: bool = False,
) -> Dict[str, Any]:
    """核心搜索逻辑"""
    if not query:
        return {"success": True, "results": [], "total": 0, "mode": "empty"}

    # 1. 解析查询意图
    parsed = parse_article_query(query)

    # 2. 确定搜索模式
    if mode == "auto":
        if parsed["is_article_query"]:
            mode = "article"
        else:
            mode = "vector"

    # 3. 获取用户区域过滤
    region_filter = None
    try:
        conn, cur = db
        await cur.execute(
            "SELECT enforcement_region FROM users WHERE id = %s",
            (user["userId"],),
        )
        user_row = await cur.fetchone()
        if user_row and dict(user_row).get("enforcement_region"):
            region_filter = user_row["enforcement_region"]
    except Exception:
        pass

    # 4. 确定要搜索的集合
    collections = _get_collections(source)
    all_results: List[Dict] = []

    # ==================== 法条定位搜索 ====================
    if mode == "article" and parsed["is_article_query"]:
        # 策略 1: 精确 filter（法律名称 + 法条编号）
        article_filter = _build_article_filter(parsed)
        if article_filter:
            for collection in collections:
                try:
                    results = await _search_qdrant_scroll(collection, limit, article_filter)
                    all_results.extend(results)
                except Exception:
                    continue

        # 策略 2: 如果精确 filter 无结果，尝试只用法条编号 filter
        if not all_results and parsed["article_num"]:
            num_only_filter = {
                "must": [{
                    "should": [
                        {"key": "article_num", "match": {"text": parsed["article_num"]}},
                        {"key": "article_number", "match": {"text": parsed["article_num"]}},
                    ]
                }]
            }
            for collection in collections:
                try:
                    results = await _search_qdrant_scroll(collection, limit, num_only_filter)
                    all_results.extend(results)
                except Exception:
                    continue

        # 如果 filter 仍然没有结果，fallback 到向量搜索
        if not all_results:
            mode = "vector"

    # ==================== 向量语义搜索 ====================
    if mode == "vector" or not all_results:
        try:
            vector = await _get_embedding(query)
        except Exception as e:
            return {"error": "Embedding 服务暂时不可用", "detail": str(e)}

        # 法条定位模式下，如果已有 filter 结果不够，也做向量搜索补充
        article_filter = _build_article_filter(parsed) if parsed["is_article_query"] else None

        for collection in collections:
            try:
                # 如果是法条定位查询，优先用 filter + vector 组合搜索
                if parsed["is_article_query"] and article_filter:
                    results = await _search_qdrant_single(
                        collection, vector, limit,
                        filter_body=article_filter,
                        region_filter=region_filter,
                    )
                else:
                    results = await _search_qdrant_single(
                        collection, vector, limit,
                        filter_body=None,
                        region_filter=region_filter,
                    )
                all_results.extend(results)
            except Exception as ex:
                logging.getLogger(__name__).warning(f"搜索集合 [{collection}] 失败: {ex}")
                continue

    # 5. 合并结果并排序
    all_results.sort(key=lambda x: x["score"], reverse=True)

    # 5.5 Rerank 重排（仅向量搜索结果）
    rerank_model_used = None
    if rerank and len(all_results) > 1:
        try:
            documents = []
            for r in all_results:
                p = r.get("payload", {})
                doc_text = f"{p.get('title', '')} {p.get('article_num', '')} {p.get('article_number', '')} {p.get('content', '')[:300]}"
                documents.append(doc_text.strip())
            rerank_results = await llm_client.rerank(
                query=query,
                documents=documents,
                top_n=limit,
            )
            if rerank_results:
                reranked = []
                for rr in rerank_results:
                    idx = rr.get("index", -1)
                    if 0 <= idx < len(all_results):
                        all_results[idx]["score"] = round(rr.get("relevance_score", all_results[idx]["score"]), 4)
                        reranked.append(all_results[idx])
                all_results = sorted(reranked, key=lambda x: x["score"], reverse=True)
                rerank_model_used = "rerank"
        except Exception as rerank_err:
            logging.getLogger(__name__).warning(f"Rerank failed (user search): {rerank_err}")

    # 6. 分页
    total = len(all_results)
    paginated = all_results[offset:offset + limit]

    # 7. 格式化返回
    formatted = []
    for r in paginated:
        payload = r["payload"]
        collection_name = payload.pop("_collection", "market_regulation_laws")

        # source_type 映射
        source_type_map = {
            "market_regulation_laws": "law",
            "discretion_standards": "discretion",
            "case_precedents": "case",
            "exemption_lists": "exemption",
        }
        st = payload.get("source_type", "")
        if not st:
            st = source_type_map.get(collection_name, "law")

        formatted.append({
            "id": r["id"],
            "score": r["score"],
            "collection": collection_name,
            "source_type": st,
            "title": payload.get("title", ""),
            "content": payload.get("content", "")[:EMBEDDING_PAYLOAD_MAX_CHARS],
            "article_num": payload.get("article_num") or payload.get("article_number", ""),
            "category": payload.get("category", ""),
            "region": payload.get("region", ""),
            "status": payload.get("status", ""),
            "effective_date": payload.get("effective_date", ""),
            "source_file": payload.get("source_file", ""),
            "payload": payload,
        })

    return {
        "success": True,
        "results": formatted,
        "total": total,
        "mode": mode,
        "parsed": parsed if parsed["is_article_query"] else None,
        "reranked": bool(rerank_model_used),
        "rerank_model": rerank_model_used,
    }


# ==================== API 端点 ====================

@router.post("/search")
async def search_knowledge_post(
    request: Request,
    user: dict = Depends(require_login),
    db=Depends(get_db),
):
    """
    知识库搜索（增强版）

    请求体:
    {
        "query": "搜索关键词",
        "limit": 10,           // 每页数量
        "offset": 0,           // 偏移量
        "source": "all",       // all | laws | discretion | cases | exemptions
        "mode": "auto",        // auto | vector | article
    }
    """
    body = await request.json()
    return await _do_search(
        query=body.get("query", "").strip(),
        limit=min(max(body.get("limit", 10), 1), 50),
        offset=max(body.get("offset", 0), 0),
        source=body.get("source", "all") or "all",
        mode=body.get("mode", "auto"),
        user=user,
        db=db,
        rerank=body.get("rerank", False),
    )


@router.get("/search")
async def search_knowledge_get(
    request: Request,
    q: str = Query(""),
    query: str = Query(""),  # 兼容旧接口参数名
    limit: int = Query(10, ge=1, le=50),
    pageSize: int = Query(0, ge=0),  # 兼容旧接口分页参数
    page: int = Query(1, ge=1),      # 兼容旧接口分页参数
    offset: int = Query(0, ge=0),
    source: str = Query("all"),
    mode: str = Query("auto"),
    user: dict = Depends(require_login),
    db=Depends(get_db),
):
    """GET 方式搜索（兼容旧接口，支持 q/query/page/pageSize 等参数）"""
    search_query = (q or query).strip()
    # 兼容 page/pageSize 分页参数
    if pageSize > 0 and offset == 0:
        limit = min(pageSize, 50)
        offset = (page - 1) * pageSize
    return await _do_search(
        query=search_query,
        limit=limit,
        offset=offset,
        source=source,
        mode=mode,
        user=user,
        db=db,
        rerank=False,
    )
