import json
import logging
from subprocess import run, CalledProcessError

from .api import retry_on_exception

logger = logging.getLogger(__name__)


def _build_akb_query_command(query: str, collection: str) -> str:
    """
    构建 akb 查询命令，包含目录存在性检查

    Args:
        query: 查询字符串
        collection: 知识库名称

    Returns:
        完整的 bash 命令字符串
    """
    collection_dir = f"/share/programs/akb/database/{collection}"

    return f"""if [ ! -d "{collection_dir}" ]; then
    echo "知识库目录不存在: {collection_dir}"
    exit 1
fi
akb simple-query \\
    -i {collection_dir}/{collection}.index \\
    -m /share/programs/BAAI/bge-m3 \\
    -p {collection_dir}/{collection}.parquet \\
    -l 5 \\
    -f json -q "{query}\""""


class RAG:
    """
    搜寻RAG知识

    :param str query: 需要运行的命令
    :param str collection: 搜寻的知识库名称。必须是下面之一："DSDP","MPNN","PySCF","RFdiffusion","gaussian","protenix","GPU4PySCF","MolecularDynamics","RDkit","OpenBabel"
    """

    def call(self, query: str, collection: str):
        cmd = _build_akb_query_command(query, collection)
        logger.info(f"搜索知识库 {collection}: {query}")

        try:
            result = run(cmd, shell='/bin/bash', check=True, text=True, capture_output=True)
        except CalledProcessError as e:
            logger.error(e.stderr.strip())
            return e.stderr.strip()

        try:
            parsed_data = json.loads(result.stdout)
            # 如果返回的是数组，直接使用
            if isinstance(parsed_data, list):
                return "\n".join(parsed_data)
            # 如果返回的是对象且包含 text 字段，使用 text 字段
            elif isinstance(parsed_data, dict) and "text" in parsed_data:
                text_array = parsed_data["text"]
                if isinstance(text_array, list):
                    return "\n".join(text_array)
                else:
                    return str(text_array)
            else:
                # 其他情况返回原始输出
                return result.stdout
        except Exception as e:
            logger.error(e)
            return result.stdout


@retry_on_exception(max_retries=3)
def knowledgeSearch(query_info, messages_prev, project_name, collection_name, max_messages=4):
    """
    知识库检索函数
    """
    if not collection_name:
        raise ValueError("collection_name 参数是必需的")

    try:
        # 使用本地RAG实现
        rag = RAG()
        result = rag.call(query_info, collection_name)

        # 构造返回格式，保持与原有API格式一致
        response_data = {
            "code": 0,
            "data": {
                "collection_name": collection_name,
                "count": 1,
                "result_list": [{
                    "chunk_id": 0,
                    "chunk_source": "document",
                    "chunk_title": result.split("\n")[0] if result else "",
                    "chunk_type": "text",
                    "content": result,
                    "doc_info": {
                        "create_time": 0,
                        "doc_id": "local_doc",
                        "doc_name": f"{collection_name}_knowledge",
                        "doc_type": "text",
                        "source": "local"
                    },
                    "score": 1.0
                }],
                "token_usage": {
                    "embedding_token_usage": {
                        "completion_tokens": 0,
                        "prompt_tokens": 0,
                        "total_tokens": 0
                    },
                    "rerank_token_usage": 0
                }
            },
            "message": "success",
            "request_id": "local_request"
        }

        return json.dumps(response_data, ensure_ascii=False)
    except Exception as e:
        error_response = {
            "code": 1,
            "message": f"知识库搜索失败: {str(e)}"
        }
        return json.dumps(error_response, ensure_ascii=False)
