import llmengine.milvus_connection
from traceback import format_exc
import argparse
from aiohttp import web
from llmengine.base_connection import get_connection_class
from appPublic.registerfunction import RegisterFunction
from appPublic.log import debug, error, info
from ahserver.serverenv import ServerEnv
from ahserver.webapp import webserver
import os
import json

helptext = """Milvus Connection Service API (using pymilvus Collection API):

1. Create Collection Endpoint:
path: /v1/createcollection
method: POST
headers: {"Content-Type": "application/json"}
data: {
    "db_type": "textdb"  // 可选，若不提供则使用默认集合 ragdb
}
response:
- Success: HTTP 200, {"status": "success", "collection_name": "ragdb" or "ragdb_textdb", "message": "集合 ragdb 或 ragdb_textdb 创建成功"}
- Error: HTTP 400, {"status": "error", "collection_name": "ragdb" or "ragdb_textdb", "message": "<error message>"}

2. Delete Collection Endpoint:
path: /v1/deletecollection
method: POST
headers: {"Content-Type": "application/json"}
data: {
    "db_type": "textdb"  // 可选，若不提供则删除默认集合 ragdb
}
response:
- Success: HTTP 200, {"status": "success", "collection_name": "ragdb" or "ragdb_textdb", "message": "集合 ragdb 或 ragdb_textdb 删除成功"}
- Success (collection does not exist): HTTP 200, {"status": "success", "collection_name": "ragdb" or "ragdb_textdb", "message": "集合 ragdb 或 ragdb_textdb 不存在，无需删除"}
- Error: HTTP 400, {"status": "error", "collection_name": "ragdb" or "ragdb_textdb", "message": "<error message>"}

3. Insert File Endpoint:
path: /v1/insertfile
method: POST
headers: {"Content-Type": "application/json"}
data: {
    "file_path": "/path/to/file.txt",  // 必填，文件路径
    "userid": "user123",               // 必填，用户 ID
    "db_type": "textdb",               // 可选，若不提供则使用默认集合 ragdb
    "knowledge_base_id": "kb123"       // 必填，知识库 ID
}
response:
- Success: HTTP 200, {"status": "success", "document_id": "<uuid>", "collection_name": "ragdb" or "ragdb_textdb", "message": "文件 <file_path> 成功嵌入并处理三元组", "status_code": 200}
- Success (triples failed): HTTP 200, {"status": "success", "document_id": "<uuid>", "collection_name": "ragdb" or "ragdb_textdb", "message": "文件 <file_path> 成功嵌入，但三元组处理失败: <error>", "status_code": 200}
- Error: HTTP 400, {"status": "error", "document_id": "", "collection_name": "ragdb" or "ragdb_textdb", "message": "<error message>", "status_code": 400}

4. Delete Document Endpoint:
path: /v1/deletefile
method: POST
headers: {"Content-Type": "application/json"}
data: {
    "userid": "user123",         // 必填，用户 ID
    "filename": "file.txt",      // 必填，文件名
    "db_type": "textdb",         // 可选，若不提供则使用默认集合 ragdb
    "knowledge_base_id": "kb123" // 必填，知识库 ID
}
response:
- Success: HTTP 200, {"status": "success", "document_id": "<uuid1,uuid2>", "collection_name": "ragdb" or "ragdb_textdb", "message": "成功删除 <count> 条 Milvus 记录，<nodes> 个 Neo4j 节点，<rels> 个 Neo4j 关系，userid=<userid>, filename=<filename>", "status_code": 200}
- Success (no records): HTTP 200, {"status": "success", "document_id": "", "collection_name": "ragdb" or "ragdb_textdb", "message": "没有找到 userid=<userid>, filename=<filename>, knowledge_base_id=<knowledge_base_id> 的记录，无需删除", "status_code": 200}
- Success (collection missing): HTTP 200, {"status": "success", "document_id": "", "collection_name": "ragdb" or "ragdb_textdb", "message": "集合 <collection_name> 不存在，无需删除", "status_code": 200}
- Error: HTTP 400, {"status": "error", "document_id": "", "collection_name": "ragdb" or "ragdb_textdb", "message": "<error message>", "status_code": 400}

5. Fused Search Query Endpoint:
path: /v1/fusedsearchquery
method: POST
headers: {"Content-Type": "application/json"}
data: {
    "query": "苹果公司在北京开设新店",
    "userid": "user1",
    "db_type": "textdb",  // 可选，若不提供则使用默认集合 ragdb
    "knowledge_base_ids": ["kb123"],
    "limit": 5,
    "offset": 0,
    "use_rerank": true
}
response:
- Success: HTTP 200, {
    "status": "success",
    "results": [
        {
            "text": "<完整文本内容>",
            "distance": 0.95,
            "source": "fused_query_with_triplets",
            "rerank_score": 0.92,  // 若 use_rerank=true
            "metadata": {
                "userid": "user1",
                "document_id": "<uuid>",
                "filename": "file.txt",
                "file_path": "/path/to/file.txt",
                "upload_time": "<iso_timestamp>",
                "file_type": "txt"
            }
        },
        ...
    ],
    "timing": {
        "collection_load": <float>,  // 集合加载耗时（秒）
        "entity_extraction": <float>,  // 实体提取耗时（秒）
        "triplet_matching": <float>,  // 三元组匹配耗时（秒）
        "triplet_text_combine": <float>,  // 拼接三元组文本耗时（秒）
        "embedding_generation": <float>,  // 嵌入向量生成耗时（秒）
        "vector_search": <float>,  // 向量搜索耗时（秒）
        "deduplication": <float>,  // 去重耗时（秒）
        "reranking": <float>,  // 重排序耗时（秒，若 use_rerank=true）
        "total_time": <float>  // 总耗时（秒）
    },
    "collection_name": "ragdb" or "ragdb_textdb"
}
- Error: HTTP 400, {
    "status": "error",
    "message": "<error message>",
    "collection_name": "ragdb" or "ragdb_textdb"
}
6. Search Query Endpoint:
path: /v1/searchquery
method: POST
headers: {"Content-Type": "application/json"}
data: {
    "query": "知识图谱的知识融合是什么？",
    "userid": "user1",
    "db_type": "textdb",  // 可选，若不提供则使用默认集合 ragdb
    "knowledge_base_ids": ["kb123"],
    "limit": 5,
    "offset": 0,
    "use_rerank": true
}
response:
- Success: HTTP 200, {
    "status": "success",
    "results": [
        {
            "text": "<完整文本内容>",
            "distance": 0.95,
            "source": "vector_query",
            "rerank_score": 0.92,  // 若 use_rerank=true
            "metadata": {
                "userid": "user1",
                "document_id": "<uuid>",
                "filename": "file.txt",
                "file_path": "/path/to/file.txt",
                "upload_time": "<iso_timestamp>",
                "file_type": "txt"
            }
        },
        ...
    ],
    "timing": {
        "collection_load": <float>,  // 集合加载耗时（秒）
        "embedding_generation": <float>,  // 嵌入向量生成耗时（秒）
        "vector_search": <float>,  // 向量搜索耗时（秒）
        "deduplication": <float>,  // 去重耗时（秒）
        "reranking": <float>,  // 重排序耗时（秒，若 use_rerank=true）
        "total_time": <float>  // 总耗时（秒）
    },
    "collection_name": "ragdb" or "ragdb_textdb"
}
- Error: HTTP 400, {
    "status": "error",
    "message": "<error message>",
    "collection_name": "ragdb" or "ragdb_textdb"
}

7. List User Files Endpoint:
path: /v1/listuserfiles
method: POST
headers: {"Content-Type": "application/json"}
data: {
    "userid": "user1",
    "db_type": "textdb"  // 可选，若不提供则使用默认集合 ragdb
}
response:
- Success: HTTP 200, {
    "status": "success",
    "files_by_knowledge_base": {
        "kb123": [
            {
                "document_id": "<uuid>",
                "filename": "file1.txt",
                "file_path": "/path/to/file1.txt",
                "upload_time": "<iso_timestamp>",
                "file_type": "txt",
                "knowledge_base_id": "kb123"
            },
            ...
        ],
        "kb456": [
            {
                "document_id": "<uuid>",
                "filename": "file2.pdf",
                "file_path": "/path/to/file2.pdf",
                "upload_time": "<iso_timestamp>",
                "file_type": "pdf",
                "knowledge_base_id": "kb456"
            },
            ...
        ]
    },
    "collection_name": "ragdb" or "ragdb_textdb"
}
- Error: HTTP 400, {
    "status": "error",
    "message": "<error message>",
    "collection_name": "ragdb" or "ragdb_textdb"
}
8. Connection Endpoint (for compatibility):
path: /v1/connection
method: POST
headers: {"Content-Type": "application/json"}
data: {
    "action": "<initialize|get_params|create_collection|delete_collection|insert_document|delete_document|fused_search|search_query|list_user_files>",
    "params": {...}
}
response:
- Success: HTTP 200, {"status": "success", ...}
- Error: HTTP 400, {"status": "error", "message": "<error message>"}

9. Docs Endpoint:
path: /docs
method: GET
response: This help text

10. Delete Knowledge Base Endpoint:
path: /v1/deleteknowledgebase
method: POST
headers: {"Content-Type": "application/json"}
data: {
    "userid": "user123",         // 必填，用户 ID
    "knowledge_base_id": "kb123",// 必填，知识库 ID
    "db_type": "textdb"          // 可选，若不提供则使用默认集合 ragdb
}
response:
- Success: HTTP 200, {"status": "success", "document_id": "<uuid1,uuid2>", "filename": "<filename1,filename2>", "collection_name": "ragdb" or "ragdb_textdb", "message": "成功删除 <count> 条 Milvus 记录，<nodes> 个 Neo4j 节点，<rels> 个 Neo4j 关系，userid=<userid>, knowledge_base_id=<knowledge_base_id>", "status_code": 200}
- Success (no records): HTTP 200, {"status": "success", "document_id": "", "filename": "", "collection_name": "ragdb" or "ragdb_textdb", "message": "没有找到 userid=<userid>, knowledge_base_id=<knowledge_base_id> 的记录，无需删除", "status_code": 200}
- Success (collection missing): HTTP 200, {"status": "success", "document_id": "", "filename": "", "collection_name": "ragdb" or "ragdb_textdb", "message": "集合 <collection_name> 不存在，无需删除", "status_code": 200}
- Error: HTTP 400, {"status": "error", "document_id": "", "filename": "", "collection_name": "ragdb" or "ragdb_textdb", "message": "<error message>", "status_code": 400}

10. List All Knowledge Bases Endpoint:
path: /v1/listallknowledgebases
method: POST
headers: {"Content-Type": "application/json"}
data: {
    "db_type": "textdb"  // 可选，若不提供则使用默认集合 ragdb
}
response:
- Success: HTTP 200, {
    "status": "success",
    "users_knowledge_bases": {
        "user1": {
            "kb123": [
                {
                    "document_id": "<uuid>",
                    "filename": "file1.txt",
                    "file_path": "/path/to/file1.txt",
                    "upload_time": "<iso_timestamp>",
                    "file_type": "txt",
                    "knowledge_base_id": "kb123"
                },
                ...
            ],
            "kb456": [
                {
                    "document_id": "<uuid>",
                    "filename": "file2.pdf",
                    "file_path": "/path/to/file2.pdf",
                    "upload_time": "<iso_timestamp>",
                    "file_type": "pdf",
                    "knowledge_base_id": "kb456"
                },
                ...
            ]
        },
        "user2": {...}
    },
    "collection_name": "ragdb" or "ragdb_textdb",
    "message": "成功列出 <count> 个用户的知识库和文件",
    "status_code": 200
}
- Error: HTTP 400, {
    "status": "error",
    "users_knowledge_bases": {},
    "collection_name": "ragdb" or "ragdb_textdb",
    "message": "<error message>",
    "status_code": 400
}
"""

def init():
    rf = RegisterFunction()
    rf.register('createcollection', create_collection)
    rf.register('deletecollection', delete_collection)
    rf.register('insertfile', insert_file)
    rf.register('deletefile', delete_file)
    rf.register('deleteknowledgebase', delete_knowledge_base)
    rf.register('fusedsearchquery', fused_search_query)
    rf.register('searchquery', search_query)
    rf.register('listuserfiles', list_user_files)
    rf.register('listallknowledgebases', list_all_knowledge_bases)
    rf.register('connection', handle_connection)
    rf.register('docs', docs)

async def docs(request, params_kw, *params, **kw):
    return web.Response(text=helptext, content_type='text/plain')

async def not_implemented(request, params_kw, *params, **kw):
    return web.json_response({
        "status": "error",
        "message": "功能尚未实现"
    }, dumps=lambda obj: json.dumps(obj, ensure_ascii=False), status=501)

async def create_collection(request, params_kw, *params, **kw):
    debug(f'{params_kw=}')
    se = ServerEnv()
    engine = se.engine
    db_type = params_kw.get('db_type', '')
    collection_name = "ragdb" if not db_type else f"ragdb_{db_type}"
    try:
        result = await engine.handle_connection("create_collection", {"db_type": db_type})
        debug(f'{result=}')
        return web.json_response(result, dumps=lambda obj: json.dumps(obj, ensure_ascii=False))
    except Exception as e:
        error(f'创建集合失败: {str(e)}')
        return web.json_response({
            "status": "error",
            "collection_name": collection_name,
            "message": str(e)
        }, dumps=lambda obj: json.dumps(obj, ensure_ascii=False), status=400)

async def delete_collection(request, params_kw, *params, **kw):
    debug(f'{params_kw=}')
    se = ServerEnv()
    engine = se.engine
    db_type = params_kw.get('db_type', '')
    collection_name = "ragdb" if not db_type else f"ragdb_{db_type}"
    try:
        result = await engine.handle_connection("delete_collection", {"db_type": db_type})
        debug(f'{result=}')
        return web.json_response(result, dumps=lambda obj: json.dumps(obj, ensure_ascii=False))
    except Exception as e:
        error(f'删除集合失败: {str(e)}')
        return web.json_response({
            "status": "error",
            "collection_name": collection_name,
            "message": str(e)
        }, dumps=lambda obj: json.dumps(obj, ensure_ascii=False), status=400)

async def insert_file(request, params_kw, *params, **kw):
    debug(f'Received params: {params_kw=}')
    se = ServerEnv()
    engine = se.engine
    file_path = params_kw.get('file_path', '')
    userid = params_kw.get('userid', '')
    db_type = params_kw.get('db_type', '')
    knowledge_base_id = params_kw.get('knowledge_base_id', '')
    document_id = params_kw.get('document_id', '')
    collection_name = "ragdb" if not db_type else f"ragdb_{db_type}"
    try:
        required_fields = ['file_path', 'userid', 'knowledge_base_id', 'document_id']
        missing_fields = [field for field in required_fields if field not in params_kw or not params_kw[field]]
        if missing_fields:
            raise ValueError(f"缺少必填字段: {', '.join(missing_fields)}")

        debug(
            f'Calling insert_document with: file_path={file_path}, userid={userid}, db_type={db_type}, knowledge_base_id={knowledge_base_id}, document_id={document_id}')
        result = await engine.handle_connection("insert_document", {
            "file_path": file_path,
            "userid": userid,
            "db_type": db_type,
            "knowledge_base_id": knowledge_base_id,
            "document_id": document_id
        })
        debug(f'Insert result: {result=}')
        status = 200 if result.get("status") == "success" else 400
        return web.json_response(result, dumps=lambda obj: json.dumps(obj, ensure_ascii=False), status=status)
    except Exception as e:
        error(f'插入文件失败: {str(e)}')
        return web.json_response({
            "status": "error",
            "collection_name": collection_name,
            "document_id": document_id,
            "message": str(e)
        }, dumps=lambda obj: json.dumps(obj, ensure_ascii=False), status=400)

async def delete_file(request, params_kw, *params, **kw):
    debug(f'Received delete_file params: {params_kw=}')
    se = ServerEnv()
    engine = se.engine
    userid = params_kw.get('userid', '')
    file_path = params_kw.get('file_path', '')
    db_type = params_kw.get('db_type', '')
    knowledge_base_id = params_kw.get('knowledge_base_id', '')
    document_id = params_kw.get('document_id', '')
    collection_name = "ragdb" if not db_type else f"ragdb_{db_type}"
    try:
        required_fields = ['userid', 'file_path', 'knowledge_base_id', 'document_id']
        missing_fields = [field for field in required_fields if field not in params_kw or not params_kw[field]]
        if missing_fields:
            raise ValueError(f"缺少必填字段: {', '.join(missing_fields)}")

        debug(f'Calling delete_document with: userid={userid}, file_path={file_path}, db_type={db_type}, knowledge_base_id={knowledge_base_id}, document_id={document_id}')
        result = await engine.handle_connection("delete_document", {
            "userid": userid,
            "file_path": file_path,
            "knowledge_base_id": knowledge_base_id,
            "document_id": document_id,
            "db_type": db_type
        })
        debug(f'Delete result: {result=}')
        status = 200 if result.get("status") == "success" else 400
        return web.json_response(result, dumps=lambda obj: json.dumps(obj, ensure_ascii=False), status=status)
    except Exception as e:
        error(f'删除文件失败: {str(e)}')
        return web.json_response({
            "status": "error",
            "collection_name": collection_name,
            "document_id": document_id,
            "message": str(e),
            "status_code": 400
        }, dumps=lambda obj: json.dumps(obj, ensure_ascii=False), status=400)

async def delete_knowledge_base(request, params_kw, *params, **kw):
    debug(f'Received delete_knowledge_base params: {params_kw=}')
    se = ServerEnv()
    engine = se.engine
    userid = params_kw.get('userid', '')
    knowledge_base_id = params_kw.get('knowledge_base_id', '')
    db_type = params_kw.get('db_type', '')
    collection_name = "ragdb" if not db_type else f"ragdb_{db_type}"
    try:
        required_fields = ['userid', 'knowledge_base_id']
        missing_fields = [field for field in required_fields if field not in params_kw or not params_kw[field]]
        if missing_fields:
            raise ValueError(f"缺少必填字段: {', '.join(missing_fields)}")

        debug(
            f'Calling delete_knowledge_base with: userid={userid}, knowledge_base_id={knowledge_base_id}, db_type={db_type}')
        result = await engine.handle_connection("delete_knowledge_base", {
            "userid": userid,
            "knowledge_base_id": knowledge_base_id,
            "db_type": db_type
        })
        debug(f'Delete knowledge base result: {result=}')
        status = 200 if result.get("status") == "success" else 400
        return web.json_response(result, dumps=lambda obj: json.dumps(obj, ensure_ascii=False), status=status)
    except Exception as e:
        error(f'删除知识库失败: {str(e)}')
        return web.json_response({
            "status": "error",
            "collection_name": collection_name,
            "document_id": "",
            "filename": "",
            "message": str(e),
            "status_code": 400
        }, dumps=lambda obj: json.dumps(obj, ensure_ascii=False), status=400)

async def search_query(request, params_kw, *params, **kw):
    debug(f'Received search_query params: {params_kw=}')
    se = ServerEnv()
    engine = se.engine
    query = params_kw.get('query')
    userid = params_kw.get('userid')
    db_type = params_kw.get('db_type', '')
    knowledge_base_ids = params_kw.get('knowledge_base_ids')
    limit = params_kw.get('limit', 5)
    offset = params_kw.get('offset', 0)
    use_rerank = params_kw.get('use_rerank', True)
    collection_name = "ragdb" if not db_type else f"ragdb_{db_type}"
    try:
        if not all([query, userid, knowledge_base_ids]):
            debug(f'query, userid 或 knowledge_base_ids 未提供')
            return web.json_response({
                "status": "error",
                "message": "query, userid 或 knowledge_base_ids 未提供",
                "collection_name": collection_name
            }, dumps=lambda obj: json.dumps(obj, ensure_ascii=False), status=400)
        result = await engine.handle_connection("search_query", {
            "query": query,
            "userid": userid,
            "knowledge_base_ids": knowledge_base_ids,
            "limit": limit,
            "offset": offset,
            "use_rerank": use_rerank,
            "db_type": db_type
        })
        debug(f'Search result: {result=}')
        response = {
            "status": "success",
            "results": result.get("results", []),
            "timing": result.get("timing", {}),
            "collection_name": collection_name
        }
        return web.json_response(response, dumps=lambda obj: json.dumps(obj, ensure_ascii=False))
    except Exception as e:
        error(f'纯向量搜索失败: {str(e)}')
        return web.json_response({
            "status": "error",
            "message": str(e),
            "collection_name": collection_name
        }, dumps=lambda obj: json.dumps(obj, ensure_ascii=False), status=400)

async def fused_search_query(request, params_kw, *params, **kw):
    debug(f'Received fused_search_query params: {params_kw=}')
    se = ServerEnv()
    engine = se.engine
    query = params_kw.get('query')
    userid = params_kw.get('userid')
    db_type = params_kw.get('db_type', '')
    knowledge_base_ids = params_kw.get('knowledge_base_ids')
    limit = params_kw.get('limit', 5)
    offset = params_kw.get('offset', 0)
    use_rerank = params_kw.get('use_rerank', True)
    collection_name = "ragdb" if not db_type else f"ragdb_{db_type}"
    try:
        if not all([query, userid, knowledge_base_ids]):
            debug(f'query, userid 或 knowledge_base_ids 未提供')
            return web.json_response({
                "status": "error",
                "message": "query, userid 或 knowledge_base_ids 未提供",
                "collection_name": collection_name
            }, dumps=lambda obj: json.dumps(obj, ensure_ascii=False), status=400)
        result = await engine.handle_connection("fused_search", {
            "query": query,
            "userid": userid,
            "knowledge_base_ids": knowledge_base_ids,
            "limit": limit,
            "offset": offset,
            "use_rerank": use_rerank,
            "db_type": db_type
        })
        debug(f'Fused search result: {result=}')
        response = {
            "status": "success",
            "results": result.get("results", []),
            "timing": result.get("timing", {}),
            "collection_name": collection_name
        }
        return web.json_response(response, dumps=lambda obj: json.dumps(obj, ensure_ascii=False))
    except Exception as e:
        error(f'融合搜索失败: {str(e)}')
        return web.json_response({
            "status": "error",
            "message": str(e),
            "collection_name": collection_name
        }, dumps=lambda obj: json.dumps(obj, ensure_ascii=False), status=400)

async def list_user_files(request, params_kw, *params, **kw):
    debug(f'{params_kw=}')
    se = ServerEnv()
    engine = se.engine
    userid = params_kw.get('userid')
    db_type = params_kw.get('db_type', '')
    collection_name = "ragdb" if not db_type else f"ragdb_{db_type}"
    try:
        if not userid:
            debug(f'userid 未提供')
            return web.json_response({
                "status": "error",
                "message": "userid 未提供",
                "collection_name": collection_name
            }, dumps=lambda obj: json.dumps(obj, ensure_ascii=False), status=400)
        result = await engine.handle_connection("list_user_files", {
            "userid": userid,
            "db_type": db_type
        })
        debug(f'{result=}')
        response = {
            "status": "success",
            "files_by_knowledge_base": result,
            "collection_name": collection_name
        }
        return web.json_response(response, dumps=lambda obj: json.dumps(obj, ensure_ascii=False))
    except Exception as e:
        error(f'列出用户文件失败: {str(e)}')
        return web.json_response({
            "status": "error",
            "message": str(e),
            "collection_name": collection_name
        }, dumps=lambda obj: json.dumps(obj, ensure_ascii=False), status=400)

async def list_all_knowledge_bases(request, params_kw, *params, **kw):
    debug(f'{params_kw=}')
    se = ServerEnv()
    engine = se.engine
    db_type = params_kw.get('db_type', '')
    collection_name = "ragdb" if not db_type else f"ragdb_{db_type}"
    try:
        result = await engine.handle_connection("list_all_knowledge_bases", {
            "db_type": db_type
        })
        debug(f'{result=}')
        response = {
            "status": result.get("status", "success"),
            "users_knowledge_bases": result.get("users_knowledge_bases", {}),
            "collection_name": collection_name,
            "message": result.get("message", ""),
            "status_code": result.get("status_code", 200)
        }
        return web.json_response(response, dumps=lambda obj: json.dumps(obj, ensure_ascii=False), status=response["status_code"])
    except Exception as e:
        error(f'列出所有用户知识库失败: {str(e)}')
        return web.json_response({
            "status": "error",
            "users_knowledge_bases": {},
            "collection_name": collection_name,
            "message": str(e),
            "status_code": 400
        }, dumps=lambda obj: json.dumps(obj, ensure_ascii=False), status=400)

async def handle_connection(request, params_kw, *params, **kw):
    debug(f'{params_kw=}')
    se = ServerEnv()
    engine = se.engine
    try:
        data = await request.json()
        action = data.get('action')
        if not action:
            debug(f'action 未提供')
            return web.json_response({
                "status": "error",
                "message": "action 参数未提供"
            }, dumps=lambda obj: json.dumps(obj, ensure_ascii=False), status=400)
        result = await engine.handle_connection(action, data.get('params', {}))
        debug(f'{result=}')
        return web.json_response(result, dumps=lambda obj: json.dumps(obj, ensure_ascii=False))
    except Exception as e:
        error(f'处理连接操作失败: {str(e)}')
        return web.json_response({
            "status": "error",
            "message": str(e)
        }, dumps=lambda obj: json.dumps(obj, ensure_ascii=False), status=400)

def main():
    parser = argparse.ArgumentParser(prog="Milvus Connection Service")
    parser.add_argument('-w', '--workdir')
    parser.add_argument('-p', '--port', default='8888')
    parser.add_argument('connection_path')
    args = parser.parse_args()
    debug(f"Arguments: {args}")
    Klass = get_connection_class(args.connection_path)
    se = ServerEnv()
    se.engine = Klass()
    workdir = args.workdir or os.getcwd()
    port = args.port
    debug(f'{args=}')
    webserver(init, workdir, port)

if __name__ == '__main__':
    main()