import asyncio
import aiohttp
import json
import traceback
from filename_proceess_and_kgquery import get_entity
from config import SHIP_MODEL_NAME

# --- 配置区 ---
BASE_URL = "https://htknow.zkzdht.com/api/v1/knowledge"
ERROR_LOG_FILE = "error_log.txt"

HEADERS = {
    "accept": "application/json",
    "x-role": "admin",
    "x-user-id": "1",
    "x-user-name": "testuser"
}

# 全局 Semaphore，用于限制并发连接数，防止对服务器造成过大压力
# 根据实际情况调整数值
CONCURRENT_LIMIT = 10
semaphore = asyncio.Semaphore(CONCURRENT_LIMIT)

async def log_error(kb_id, file_id, error_msg):
    """
    将错误信息追加写入日志文件
    注意：文件 I/O 在 asyncio 中是阻塞的，如果日志量极大，可以考虑使用线程池执行器
    """
    # 对于少量的日志写入，直接使用 await asyncio.to_thread 或直接同步写入通常足够
    # 这里为了简化，直接使用同步写入（在日志量不大时影响较小）
    # 如果需要完全异步，可以使用 aiofiles 库
    try:
        with open(ERROR_LOG_FILE, "a", encoding="utf-8") as f:
            f.write(f"[ERROR] KnowledgeBaseID: {kb_id}, FileID: {file_id}, Message: {error_msg}\n")
    except Exception as e:
        print(f"写入错误日志时发生异常: {e}")

async def find_ship_info_by_hull(json_file_path, data):
    """从NER结果中提取舷号，再从JSON文件中查找对应的舰船信息"""
    try:
        # 同步文件读取，如果文件很大或数量很多，可以使用 aiofiles
        with open(json_file_path, 'r', encoding='utf-8') as f:
            ship_data = json.load(f)
            
        hulls = [
            entity['properties']['舷号']
            for entity in data.get('entities', [])
            if 'properties' in entity and '舷号' in entity['properties']
        ]
        if not hulls:
            return None

        target_hull = hulls[0]

        for item in ship_data:
            if str(item.get('hull_number')) == str(target_hull):
                return {
                    "model_name": item.get('model_name'),
                    "ship_name": item.get('ship_name')
                }
        return None
    except Exception as e:
        print(f"查找舰船信息时发生错误: {e}")
        return None

def format_entity_text(data):
    """格式化实体文本 (纯计算逻辑，保持同步)"""
    type_mapping = {
        '舰艇': [('舰艇为', '名称'), ('，舷号', '舷号')],
        '系统': [('，系统：', '名称')],
        '子系统': [('，子系统：', '名称')],
        '设备': [('，设备：', '名称')]
    }
    
    text_parts = []
    for entity in data.get('entities', []):
        entity_type = entity.get('type')
        properties = entity.get('properties', {})
        
        if entity_type in type_mapping:
            entity_str = ""
            for prefix, attr_key in type_mapping[entity_type]:
                value = str(properties.get(attr_key, ''))
                entity_str += f"{prefix}{value}"
            text_parts.append(entity_str)
            
    return ''.join(text_parts)

async def get_knowledge_base_files(session, kb_id):
    """
    获取指定知识库的文件列表
    """
    url = f"{BASE_URL}/knowledge_base/{kb_id}"
    try:
        async with session.get(url, headers=HEADERS, timeout=120) as response:
            response.raise_for_status()
            data = await response.json()
            
            if 'files' in data and isinstance(data['files'], list):
                return data['files']
            else:
                print(f"警告: 知识库 {kb_id} 中未找到 'files' 字段。")
                return []
                
    except Exception as e:
        error_msg = f"获取知识库详情失败: {str(e)} | Traceback: {traceback.format_exc()}"
        print(error_msg)
        await log_error(kb_id, "N/A", error_msg)
        return []

async def get_file_slices(session, file_id):
    """获取文件所有切片"""
    url = f"{BASE_URL}/files/{file_id}/slices"
    try:
        async with session.get(url, headers=HEADERS, timeout=120) as response:
            response.raise_for_status()
            return await response.json()
    except Exception as e:
        error_msg = f"获取文件切片失败: {str(e)}"
        print(error_msg)
        # 错误会在调用处记录
        raise Exception(error_msg)

async def update_slices(session, slices_data, file_id):
    """更新切片"""
    url = f"{BASE_URL}/files/{file_id}/slices"
    headers = {
        **HEADERS,
        "content-type": "application/json"
    }
    payload = {"slices": slices_data}
    try:
        async with session.put(url, headers=headers, json=payload, timeout=120) as response:
            response.raise_for_status()
            return await response.json()
    except Exception as e:
        error_msg = f"更新切片失败: {str(e)}"
        print(error_msg)
        raise Exception(error_msg)

async def process_single_file(session, kb_id, file):
    """
    处理单个文件的逻辑，用于并发执行
    """
    file_id = file.get('id')
    filename = file.get('filename', 'Unknown')
    
    print(f"  处理文件: {filename} (ID: {file_id})")
    
    try:
        # --- 步骤 A: 提取实体信息 ---
        final_result = get_entity(filename)
        
        # 处理舰船型号
        xinghao = await find_ship_info_by_hull(SHIP_MODEL_NAME, final_result)
        model_name = xinghao.get('model_name', '') if xinghao else ''
        
        # 格式化其他实体文本
        other_data = format_entity_text(final_result)
        
        # --- 步骤 B: 获取切片 ---
        slices = await get_file_slices(session, file_id)
        
        # --- 步骤 C: 处理切片内容 ---
        slices_v2 = []
        info_suffix = f"({other_data}，型号为{model_name})" if model_name and other_data else ""            
        for ins in slices:
            ins_id = ins.get('id')
            content = ins.get('content', '')
            
            # 查找第一行的结尾（换行符位置）
            newline_idx = content.find('\n')
            
            # 拼接新内容
            if newline_idx == -1:
                # 内容中没有换行符，直接在末尾添加
                new_content = content + info_suffix
            else:
                # 在第一行末尾插入
                new_content = content[:newline_idx] + info_suffix + content[newline_idx:]
            
            slices_v2.append({"id": ins_id, "content": new_content})
        
        # --- 步骤 D: 更新切片 ---
        await update_slices(session, slices_v2, file_id)
        print(f"   成功更新文件: {filename}")
        
    except Exception as e:
        # 捕获当前文件处理过程中的所有异常
        error_msg = f"处理文件失败: {str(e)}"
        print(error_msg)
        # 记录错误日志
        await log_error(kb_id, file_id, error_msg)
        # 跳过当前文件
        return False
    return True

async def process_files_in_knowledge_base(session, kb_id):
    """
    处理单个知识库中的所有文件
    """
    print(f"\n=== 开始处理知识库 ID: {kb_id} ===")
    
    # 1. 获取文件列表
    files = await get_knowledge_base_files(session, kb_id)
    if not files:
        print(f"知识库 {kb_id} 中没有文件或获取文件列表失败，跳过。")
        return

    # 2. 并发处理每个文件
    # 使用 asyncio.gather 来并发处理所有文件
    # 如果文件数量巨大，可以分批处理以避免内存溢出
    file_tasks = [
        process_single_file(session, kb_id, file) 
        for file in files
    ]
    
    # 并发执行所有文件任务
    # return_exceptions=True 会捕获任务内部的异常并作为结果返回，而不是中断整个 gather
    results = await asyncio.gather(*file_tasks, return_exceptions=True)
    
    # 统计结果
    success_count = sum(1 for r in results if r is True)
    failed_count = len(results) - success_count
    print(f"知识库 {kb_id} 处理完成。成功: {success_count}, 失败: {failed_count}")

async def main():
    KNOWLEDGE_BASE_IDS = ['8']
    print(f"开始批量处理 {len(KNOWLEDGE_BASE_IDS)} 个知识库...")
    
    # 创建一个全局的 TCP 连接器，复用连接
    connector = aiohttp.TCPConnector(limit=CONCURRENT_LIMIT, limit_per_host=5)
    
    async with aiohttp.ClientSession(connector=connector) as session:
        # 遍历所有知识库ID，并发处理
        # 如果知识库数量很多，也可以在这里使用 asyncio.gather
        for kb_id in KNOWLEDGE_BASE_IDS:
            try:
                await process_files_in_knowledge_base(session, kb_id)
            except Exception as e:
                error_msg = f"处理知识库顶层异常: {str(e)}"
                print(error_msg)
                await log_error(kb_id, "N/A", error_msg)
                continue
    
    print("所有知识库处理完成。")

if __name__ == "__main__":
    # 运行异步主函数
    asyncio.run(main())
