import sys
import os
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import json
import re
import asyncio
import aiofiles
from concurrent.futures import ThreadPoolExecutor
from preprocessing.agent_system.extractor_agent.agent import DepartmentAgent
from typing import Dict, List, Tuple
import gc
import time
import psutil
from tqdm.asyncio import tqdm as async_tqdm
from tqdm import tqdm

json_path = "dataset/merged_pretrain_extra.jsonl"
output_path = "dataset/processed_triples.jsonl"

# 优化后的配置参数 - 降低资源消耗
BATCH_SIZE = 5000  # 减少批次大小：每批1万条数据
MAX_CONCURRENT = 200  # 减少并发数：最多50条并发处理
AGENT_POOL_SIZE = 20 # 大幅减少agent池大小：只创建5个agent实例

def get_memory_usage():
    """获取当前内存使用情况"""
    process = psutil.Process(os.getpid())
    memory_info = process.memory_info()
    memory_mb = memory_info.rss / 1024 / 1024
    return memory_mb

def print_memory_info(stage=""):
    """打印内存使用信息"""
    memory_mb = get_memory_usage()
    print(f"🔧 {stage} - 内存使用: {memory_mb:.1f} MB")

# 创建extractor_agent池，避免并发冲突
def create_extractor_pool(pool_size: int = 5):
    """创建extractor_agent池"""
    print(f"正在创建 {pool_size} 个agent实例...")
    agents = []
    for i in range(pool_size):
        try:
            agent = DepartmentAgent(model_type="deepseek")
            agents.append(agent)
            print(f"  ✓ Agent {i+1}/{pool_size} 创建成功")
        except Exception as e:
            print(f"  ✗ Agent {i+1} 创建失败: {e}")
    print(f"Agent池创建完成，实际创建了 {len(agents)} 个实例")
    return agents

# 延迟初始化agent池
AGENT_POOL = None
agent_pool_index = 0

def get_agent_pool():
    """获取agent池，延迟初始化"""
    global AGENT_POOL
    if AGENT_POOL is None:
        print_memory_info("创建Agent池前")
        AGENT_POOL = create_extractor_pool(pool_size=AGENT_POOL_SIZE)
        print_memory_info("创建Agent池后")
    return AGENT_POOL

def get_next_agent():
    """轮询获取下一个可用的agent"""
    global agent_pool_index
    pool = get_agent_pool()
    agent = pool[agent_pool_index % len(pool)]
    agent_pool_index += 1
    return agent

def clean_and_split_text(text):
    """
    去除文本开头结尾的标记，并按句子分割
    """
    # 去除开头的<|im_start|>和结尾的<|im_end|>
    text = text.strip()
    if text.startswith('<|im_start|>'):
        text = text[len('<|im_start|>'):]
    if text.endswith('<|im_end|>'):
        text = text[:-len('<|im_end|>')]
    
    # 清理文本，去除多余的空白字符
    text = text.strip()
    
    # 按句子分割（根据句号、问号、感叹号等标点符号）
    # 使用正则表达式匹配句子结束标志
    sentence_endings = r'[.!?。！？]'
    sentences = re.split(sentence_endings, text)
    
    # 清理每个句子，去除空白和空句子
    cleaned_sentences = []
    for sentence in sentences:
        sentence = sentence.strip()
        if sentence and len(sentence) > 5:  # 只保留非空且有意义的句子
            cleaned_sentences.append(sentence)
    
    return cleaned_sentences

async def extract_triple_from_sentence_async(sentence: str, context: str = None) -> Dict:
    """
    异步使用extractor_agent从句子中提取三元组
    """
    try:
        # 获取一个agent实例
        agent = get_next_agent()
        result = await agent.async_run(sentence=sentence, context=context)
        return {
            "sentence": sentence,
            "triple": {
                "subject": result.triple.subject,
                "predicate": result.triple.predicate,
                "object": result.triple.object
            },
            "confidence": result.confidence
        }
    except Exception as e:
        return {
            "sentence": sentence,
            "triple": {
                "subject": "",
                "predicate": "",
                "object": ""
            },
            "confidence": 0.0,
            "error": str(e)
        }

async def process_paragraph_async(line_num: int, original_text: str, semaphore: asyncio.Semaphore) -> Dict:
    """
    异步处理单个段落
    """
    async with semaphore:  # 控制并发数量
        try:
            # 清理并分割文本
            sentences = clean_and_split_text(original_text)
            
            if not sentences:
                return None
                
            # 构建当前段落的结果
            paragraph_result = {
                "source_line": line_num,
                "original_paragraph": original_text,
                "sentences": [],
                "triples": []
            }
            
            # 异步处理所有句子
            tasks = []
            for sentence in sentences:
                task = extract_triple_from_sentence_async(sentence, context=original_text)
                tasks.append(task)
            
            # 等待所有句子处理完成
            triple_results = await asyncio.gather(*tasks)
            
            # 整理结果
            for i, sentence in enumerate(sentences):
                paragraph_result["sentences"].append(sentence)
                paragraph_result["triples"].append(triple_results[i])
            
            return paragraph_result
            
        except Exception as e:
            print(f"处理第 {line_num} 行时出错: {e}")
            return None

async def process_batch_async(batch_data: List[Tuple[int, str]], batch_num: int) -> List[Dict]:
    """
    异步处理一个批次的数据，带进度条和内存监控
    """
    print(f"\n=== 异步处理批次 {batch_num} ===")
    print(f"批次大小: {len(batch_data)} 条记录")
    print_memory_info(f"批次 {batch_num} 开始前")
    
    start_time = time.time()
    
    # 创建信号量控制并发数量
    semaphore = asyncio.Semaphore(MAX_CONCURRENT)
    
    # 分块处理任务，避免一次性创建太多任务
    chunk_size = 1000  # 每次处理1000个任务
    all_results = []
    
    for chunk_start in range(0, len(batch_data), chunk_size):
        chunk_end = min(chunk_start + chunk_size, len(batch_data))
        chunk_data = batch_data[chunk_start:chunk_end]
        
        print(f"处理子块 {chunk_start//chunk_size + 1}/{(len(batch_data)-1)//chunk_size + 1} ({len(chunk_data)} 条记录)")
        
        # 创建当前块的异步任务
        tasks = []
        for line_num, original_text in chunk_data:
            task = process_paragraph_async(line_num, original_text, semaphore)
            tasks.append(task)
        
        # 使用进度条处理当前块
        progress_bar = tqdm(total=len(tasks), desc=f"批次{batch_num}-块{chunk_start//chunk_size + 1}", unit="段落", ncols=100)
        
        chunk_results = []
        completed_tasks = 0
        
        # 使用as_completed来获取完成的任务，并更新进度条
        for coro in asyncio.as_completed(tasks):
            try:
                result = await coro
                chunk_results.append(result)
                completed_tasks += 1
                
                # 更新进度条
                progress_bar.update(1)
                
                # 每完成50个任务更新一次描述
                if completed_tasks % 50 == 0:
                    valid_results = [r for r in chunk_results if r is not None]
                    progress_bar.set_postfix({
                        '有效': len(valid_results),
                        '完成': completed_tasks,
                        '成功率': f"{len(valid_results)/completed_tasks*100:.1f}%"
                    })
            except Exception as e:
                print(f"任务执行失败: {e}")
                completed_tasks += 1
                progress_bar.update(1)
        
        progress_bar.close()
        all_results.extend(chunk_results)
        
        # 每个块完成后清理内存
        del tasks, chunk_results
        gc.collect()
        
        print_memory_info(f"批次 {batch_num} 块 {chunk_start//chunk_size + 1} 完成后")
    
    # 过滤None结果
    valid_results = [result for result in all_results if result is not None]
    
    # 统计信息
    batch_sentences = sum(len(result["sentences"]) for result in valid_results)
    batch_triples = sum(
        sum(1 for triple in result["triples"] if triple["confidence"] > 0.0)
        for result in valid_results
    )
    
    end_time = time.time()
    processing_time = end_time - start_time
    
    print(f"批次 {batch_num} 异步处理完成:")
    print(f"  - 有效段落: {len(valid_results)}/{len(batch_data)} ({len(valid_results)/len(batch_data)*100:.1f}%)")
    print(f"  - 总句子数: {batch_sentences}")
    print(f"  - 成功三元组: {batch_triples}")
    print(f"  - 三元组成功率: {batch_triples/batch_sentences*100:.1f}%" if batch_sentences > 0 else "无句子")
    print(f"  - 处理时间: {processing_time:.2f}秒")
    print(f"  - 处理速度: {len(batch_data)/processing_time:.2f}段落/秒")
    
    print_memory_info(f"批次 {batch_num} 完成后")
    
    return valid_results

async def write_results_batch(results: List[Dict], output_path: str):
    """
    异步批量写入结果，带进度提示
    """
    try:
        print(f"开始批量写入 {len(results)} 条结果...")
        
        # 准备写入内容
        content_lines = []
        for result in results:
            content_lines.append(json.dumps(result, ensure_ascii=False))
        
        # 异步批量写入
        async with aiofiles.open(output_path, "a", encoding="utf-8") as f:
            await f.write("\n".join(content_lines) + "\n")
        
        print(f"✓ 成功批量写入 {len(results)} 条结果到 {output_path}")
        
    except Exception as e:
        print(f"✗ 批量写入失败: {e}")
        print("尝试逐条写入...")
        
        # 如果批量写入失败，回退到逐条写入（带进度条）
        async with aiofiles.open(output_path, "a", encoding="utf-8") as f:
            for result in tqdm(results, desc="逐条写入", unit="条"):
                await f.write(json.dumps(result, ensure_ascii=False) + "\n")
        print(f"✓ 逐条写入完成")

# 主处理流程
async def main_async():
    total_processed = 0
    total_sentences = 0
    total_triples = 0
    batch_num = 0
    
    print("=== 开始异步批次处理JSONL文件 ===")
    print(f"优化后的配置信息:")
    print(f"  - 批次大小: {BATCH_SIZE:,} 条记录")
    print(f"  - 最大并发数: {MAX_CONCURRENT}")
    print(f"  - Agent池大小: {AGENT_POOL_SIZE}")
    print(f"  - 输入文件: {json_path}")
    print(f"  - 输出文件: {output_path}")
    print()
    
    print_memory_info("程序开始")
    
    # 清空输出文件
    async with aiofiles.open(output_path, "w", encoding="utf-8") as f:
        pass
    
    # 读取并处理数据
    with open(json_path, "r", encoding="utf-8") as f_in:
        batch_data = []
        
        for line_num, line in enumerate(f_in):
            if line.strip():  # 跳过空行
                try:
                    item = json.loads(line)
                    original_text = item.get("text", "")
                    
                    if original_text:
                        batch_data.append((line_num + 1, original_text))
                        
                        # 当批次达到指定大小时，异步处理这个批次
                        if len(batch_data) >= BATCH_SIZE:
                            batch_num += 1
                            
                            # 异步处理批次
                            batch_results = await process_batch_async(batch_data, batch_num)
                            
                            # 批量写入结果
                            if batch_results:
                                await write_results_batch(batch_results, output_path)
                            
                            # 统计信息
                            batch_sentences = sum(len(result["sentences"]) for result in batch_results)
                            batch_triples = sum(
                                sum(1 for triple in result["triples"] if triple["confidence"] > 0.0)
                                for result in batch_results
                            )
                            
                            total_processed += len(batch_data)
                            total_sentences += batch_sentences
                            total_triples += batch_triples
                            
                            print(f"\n📊 批次 {batch_num} 累计统计:")
                            print(f"  - 累计处理段落: {total_processed:,}")
                            print(f"  - 累计句子数: {total_sentences:,}")
                            print(f"  - 累计三元组: {total_triples:,}")
                            print(f"  - 整体成功率: {total_triples/total_sentences*100:.1f}%")
                            print("-" * 80)
                            
                            # 清理批次数据，释放内存
                            batch_data.clear()
                            batch_results.clear()
                            gc.collect()  # 强制垃圾回收
                            
                            print_memory_info(f"批次 {batch_num} 清理后")
                            
                except json.JSONDecodeError as e:
                    print(f"第 {line_num + 1} 行JSON解析错误: {e}")
                except Exception as e:
                    print(f"处理第 {line_num + 1} 行时出错: {e}")
        
        # 处理最后一个不完整的批次
        if batch_data:
            batch_num += 1
            batch_results = await process_batch_async(batch_data, batch_num)
            
            if batch_results:
                await write_results_batch(batch_results, output_path)
            
            batch_sentences = sum(len(result["sentences"]) for result in batch_results)
            batch_triples = sum(
                sum(1 for triple in result["triples"] if triple["confidence"] > 0.0)
                for result in batch_results
            )
            
            total_processed += len(batch_data)
            total_sentences += batch_sentences
            total_triples += batch_triples
    
    # 最终统计
    print(f"\n{'='*80}")
    print(f"🎉 所有批次异步处理完成!")
    print(f"{'='*80}")
    print(f"最终统计:")
    print(f"  - 总批次数: {batch_num}")
    print(f"  - 总段落数: {total_processed:,}")
    print(f"  - 总句子数: {total_sentences:,}")
    print(f"  - 总三元组: {total_triples:,}")
    print(f"  - 整体成功率: {total_triples/total_sentences*100:.1f}%" if total_sentences > 0 else "无有效句子")
    print(f"  - 输出文件: {output_path}")
    print(f"{'='*80}")
    
    print_memory_info("程序结束前")
    
    # 显示示例结果
    await show_sample_results()

async def show_sample_results():
    """显示前几个处理结果作为示例"""
    print("\n📋 前3个处理结果示例:")
    try:
        async with aiofiles.open(output_path, "r", encoding="utf-8") as f:
            i = 0
            async for line in f:
                if i >= 3:
                    break
                item = json.loads(line)
                print(f"\n--- 段落 {i+1} (来源行: {item['source_line']}) ---")
                print(f"原始段落: {item['original_paragraph'][:100]}...")
                print(f"句子数量: {len(item['sentences'])}")
                if item['triples']:
                    for j, triple in enumerate(item['triples'][:2]):  # 只显示前2个三元组
                        print(f"  句子 {j+1}: {triple['sentence'][:50]}...")
                        if triple['confidence'] > 0:
                            print(f"  三元组: {triple['triple']['subject']} -> {triple['triple']['predicate']} -> {triple['triple']['object']}")
                            print(f"  置信度: {triple['confidence']:.2f}")
                        else:
                            print(f"  提取失败: {triple.get('error', '未知错误')}")
                i += 1
    except Exception as e:
        print(f"读取示例结果时出错: {e}")

def main():
    """主入口函数"""
    try:
        # 运行异步主函数
        asyncio.run(main_async())
    except KeyboardInterrupt:
        print("\n⚠️ 用户中断处理")
    except Exception as e:
        print(f"❌ 处理过程中出现错误: {e}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    main()