"""
RAG搜索工具模块
包含知识库搜索、文件处理等功能
"""
import os
import json
import base64
import asyncio
import httpx
from typing import Optional, Dict, Any
from langchain_core.tools import tool

from config import RAG_CONFIG
from utils.function_tracker import track_function_calls


@track_function_calls
async def rag_search(
    query: str,
    top_k: Optional[int] = 5,
    file_id: Optional[int] = None,
    kb_id: Optional[str] = None
) -> Dict[str, Any]:
    """
    知识库搜索工具
    调用知识库搜索API，检索与query相关的文档内容

    Args:
        query: 搜索关键词（必填）
        top_k: 返回结果数量（默认5）
        file_id: 文件ID（可选）
        kb_id: 知识库ID（可选）

    Returns:
        搜索结果，包含sourceCitation格式的数据
    """

    if not query or not query.strip():
        return {"success": False, "error": "查询文本不能为空"}

    SEARCH_URL = RAG_CONFIG["search_endpoint"]

    params = {"query": query.strip()}

    if file_id is None:
        config_file_id = RAG_CONFIG.get("file-id")
        if config_file_id is not None:
            try:
                file_id = int(config_file_id) if isinstance(config_file_id, str) else config_file_id
            except (ValueError, TypeError):
                file_id = None

    if kb_id is None:
        config_kb_id = RAG_CONFIG.get("kb-id")
        if config_kb_id is not None:
            kb_id = str(config_kb_id) if config_kb_id is not None else None

    if file_id is not None and file_id != "":
        params["file_id"] = file_id
    if kb_id is not None and kb_id != "":
        params["kb_id"] = kb_id

    print(111111111111111111)
    print(RAG_CONFIG["x-user-id"])
    print(RAG_CONFIG["x-user-name"])
    print(RAG_CONFIG["x-role"])

    headers = {
        "accept": "application/json",
        "x-user-id": RAG_CONFIG["x-user-id"],
        "x-role": RAG_CONFIG["x-role"]
    }

    print(f"RAG params: {params}")

    max_retries = 1
    retry_delay = 2

    response = None
    last_exception = None

    for attempt in range(max_retries + 1):
        try:
            async with httpx.AsyncClient(timeout=180.0, verify=False) as client:
                response = await client.get(SEARCH_URL, params=params, headers=headers)
            break

        except (httpx.ReadTimeout, httpx.ConnectTimeout, httpx.RequestError) as e:
            last_exception = e
            if attempt < max_retries:
                print(f"RAG请求超时，第 {attempt+1} 次重试，2秒后重新请求...")
                await asyncio.sleep(retry_delay)
            else:
                return {"success": False, "error": f"请求超时: {str(e)}"}

    try:
        if response.status_code != 200:
            return {"success": False, "error": f"HTTP {response.status_code}"}

        result = response.json()
        data = result.get("results", [])

        filtered_data = [item for item in data if item.get("score", 0.0) >= 0.4]

        sorted_data = sorted(
            filtered_data,
            key=lambda x: x.get("score", 0.0),
            reverse=True
        )[:top_k]

        source_citation = {}

        for idx, item in enumerate(sorted_data, start=1):
            if not isinstance(item, dict):
                continue

            file_info = item.get("file")
            kb_info = item.get("kb")

            filename = file_info.get("filename", "") if file_info else ""
            resource = filename or (kb_info.get("name", "未分组文档") if kb_info else "未分组文档")

            positions = item.get("positions", [])
            page_idx = positions[0].get("page_idx", 0) if positions else 0

            citation_item = {
                "index": idx,
                "id": item.get("id", ""),
                "text": item.get("content", ""),
                "resource": resource,
                "filename": filename,
                "score": item.get("score", 0.0),
                "page_idx": page_idx,
                "file_id": item.get("file_id", ""),
                "kb_id": kb_info.get("id", "") if kb_info else "",
                "kb_name": kb_info.get("name", "") if kb_info else "",
                "positions": positions
            }

            if resource not in source_citation:
                source_citation[resource] = []
            source_citation[resource].append(citation_item)

        return {
            "success": True,
            "sourceCitation": source_citation,
            "total": len(sorted_data),
            "total_available": len(data)
        }

    except Exception as e:
        return {"success": False, "error": str(e)}


@tool
async def rag_search_tool(query: str, top_k: Optional[int] = 5) -> str:
    """
    RAG检索工具（供工作流调用）
    在知识库中检索与query相关的文档内容

    Args:
        query: 检索查询
        top_k: 返回结果数量

    Returns:
        检索结果文本
    """
    result = await rag_search(query, top_k=top_k)

    if not result.get("success", False):
        return ""

    source_citation = result.get("sourceCitation", {})
    if not source_citation:
        return ""

    results_parts = []
    for resource, items in source_citation.items():
        for item in items:
            text = item.get("text", "")
            if text:
                results_parts.append(f"【{resource}】{text}")

    results_text = "\n\n".join(results_parts) if results_parts else ""
    if len(results_text) > 2000:
        results_text = results_text[:2000] + "..."
    return results_text


@track_function_calls
async def file_to_text(file_path: str) -> Dict[str, Any]:
    """
    文件转文本工具
    支持多种文件格式：PDF、Word、TXT、图片等

    Args:
        file_path: 文件路径

    Returns:
        转换结果，包含文本内容
    """
    if not file_path or not file_path.strip():
        return {"success": False, "error": "文件路径不能为空"}

    file_path = file_path.strip()

    if not os.path.exists(file_path):
        return {"success": False, "error": f"文件不存在: {file_path}"}

    try:
        file_ext = os.path.splitext(file_path)[1].lower()

        if file_ext == ".txt":
            with open(file_path, "r", encoding="utf-8") as f:
                text = f.read()
            return {"success": True, "text": text}

        elif file_ext == ".pdf":
            import fitz
            doc = fitz.open(file_path)
            text_parts = []
            for page in doc:
                text_parts.append(page.get_text())
            text = "\n".join(text_parts)
            doc.close()
            return {"success": True, "text": text}

        elif file_ext in [".docx", ".doc"]:
            from docx import Document
            doc = Document(file_path)
            text_parts = []
            for para in doc.paragraphs:
                text_parts.append(para.text)
            text = "\n".join(text_parts)
            return {"success": True, "text": text}

        elif file_ext in [".jpg", ".jpeg", ".png", ".bmp", ".gif"]:
            with open(file_path, "rb") as f:
                image_data = f.read()
            base64_image = base64.b64encode(image_data).decode("utf-8")
            base64_url = f"data:image/{file_ext[1:]};base64,{base64_image}"
            
            from modelsAPI.model_api import OpenaiAPI
            text = await OpenaiAPI.vlm_chat(
                image_path=base64_url,
                prompt="请描述这张图片的内容"
            )
            return {"success": True, "text": text}

        else:
            return {"success": False, "error": f"不支持的文件格式: {file_ext}"}

    except Exception as e:
        return {"success": False, "error": f"文件处理失败: {str(e)}"}


@tool
async def file_to_text_tool(file_path: str) -> str:
    """
    文件转文本工具（供工作流调用）
    支持多种文件格式：PDF、Word、TXT、图片等

    Args:
        file_path: 文件路径

    Returns:
        文件文本内容
    """
    result = await file_to_text(file_path)

    if not result.get("success", False):
        return f"文件处理失败: {result.get('error', '未知错误')}"

    return result.get("text", "")
