"""
研究综述生成模块
"""

import json
import os
import re
from typing import List, Optional, Dict, Any
from pathlib import Path
from datetime import datetime

from openai import AsyncOpenAI
from docx import Document
from docx.shared import Pt, Inches
from docx.enum.text import WD_ALIGN_PARAGRAPH

from .models import ProjectInfo, SearchResults


class ReviewGenerator:
    """研究综述生成器"""
    
    def __init__(self, config=None):
        """
        Args:
            config: LLMConfig 或 None（从环境变量读取）
        """
        if config:
            self.client = AsyncOpenAI(
                api_key=config.api_key,
                base_url=config.base_url,
            )
            self.model = config.model
        else:
            self.client = AsyncOpenAI(
                api_key=os.getenv("OPENAI_API_KEY"),
                base_url=os.getenv("NSFC_LLM_BASE_URL"),
            )
            self.model = os.getenv("NSFC_LLM_MODEL", "gpt-4o")
    
    # 配置参数
    BATCH_SIZE = 10  # 每批处理项目数
    
    async def generate_review(
        self, 
        projects: List[ProjectInfo],
        topic: str,
        output_path: Optional[str] = None,
        stream: bool = True,
        custom_requirements: Optional[str] = None
    ) -> str:
        """
        生成研究综述（分批处理 + 合并）
        
        Args:
            projects: 项目列表
            topic: 研究主题
            output_path: 输出文件路径
            stream: 是否流式输出
            custom_requirements: 用户自定义要求
        
        Returns:
            综述内容
        """
        if not projects:
            return "无项目数据，无法生成综述"
        
        total = len(projects)
        
        # 项目少于等于 BATCH_SIZE，直接生成
        if total <= self.BATCH_SIZE:
            print(f"  📝 整理 {total} 个项目信息...")
            project_info = self._prepare_project_info(projects)
            print("  🤖 生成研究综述...")
            review = await self._generate_final_review(project_info, topic, stream, custom_requirements)
        else:
            # 分批处理
            print(f"  📊 共 {total} 个项目，分批处理...")
            batch_summaries = await self._process_batches(projects, topic)
            
            # 合并生成最终综述
            print("  🔄 合并生成最终综述...")
            review = await self._merge_summaries(batch_summaries, topic, stream, custom_requirements)
        
        # 保存结果
        if output_path:
            self._save_review(review, projects, topic, output_path)
            print(f"\n  ✓ 综述已保存到: {output_path}")
        
        return review
    
    def _prepare_project_info(self, projects: List[ProjectInfo], start_idx: int = 1) -> str:
        """准备项目信息摘要（使用完整摘要）
        
        Args:
            projects: 项目列表
            start_idx: 起始编号（用于分批时保持全局编号）
        """
        lines = []
        for i, p in enumerate(projects, start_idx):
            info = f"[{i}] 《{p.project_name}》"
            info += f"\n    负责人: {p.project_leader} | 单位: {p.institution}"
            if p.project_type:
                info += f" | 类别: {p.project_type}"
            if p.approval_year:
                info += f" | 年度: {p.approval_year}"
            if p.funding_amount:
                info += f" | 经费: {p.funding_amount}万"
            if p.keywords:
                info += f"\n    关键词: {p.keywords}"
            
            # 优先使用中文摘要，其次结题摘要，不截断
            abstract = p.abstract_cn or p.conclusion_abstract
            if abstract:
                info += f"\n    摘要: {abstract}"
            
            lines.append(info)
        
        return "\n\n".join(lines)
    
    async def _process_batches(self, projects: List[ProjectInfo], topic: str) -> List[Dict[str, Any]]:
        """分批处理项目，生成每批的摘要总结
        
        Returns:
            每批的总结列表，包含 batch_num, project_indices, summary
        """
        batch_summaries = []
        total = len(projects)
        total_batches = (total + self.BATCH_SIZE - 1) // self.BATCH_SIZE
        
        for batch_num in range(total_batches):
            start_idx = batch_num * self.BATCH_SIZE
            end_idx = min(start_idx + self.BATCH_SIZE, total)
            batch = projects[start_idx:end_idx]
            
            print(f"    [{batch_num + 1}/{total_batches}] 处理项目 {start_idx + 1}-{end_idx}...")
            
            # 准备该批次的项目信息
            project_info = self._prepare_project_info(batch, start_idx + 1)
            
            # 生成该批次的总结
            summary = await self._summarize_batch(project_info, topic, start_idx + 1, end_idx)
            
            batch_summaries.append({
                'batch_num': batch_num + 1,
                'project_indices': list(range(start_idx + 1, end_idx + 1)),
                'summary': summary
            })
        
        return batch_summaries
    
    async def _summarize_batch(
        self, 
        project_info: str, 
        topic: str, 
        start_idx: int, 
        end_idx: int
    ) -> str:
        """生成单批项目的摘要总结"""
        prompt = f"""请分析以下国家自然科学基金项目（项目编号 {start_idx}-{end_idx}），提取与"{topic}"相关的关键信息。

## 项目信息
{project_info}

## 分析要求
请提取并总结：
1. **主要研究方向**：这些项目涉及哪些研究方向？
2. **关键技术方法**：采用了哪些技术方法？
3. **研究机构分布**：主要有哪些研究机构？
4. **创新点**：有哪些值得关注的创新点？

## 输出要求
- 保留项目编号引用，格式如 [1]、[2-3] 等
- 简洁但信息完整
- 500-800字"""

        response = await self.client.chat.completions.create(
            model=self.model,
            messages=[{"role": "user", "content": prompt}],
            temperature=0.5,
            max_tokens=1500,
        )
        return response.choices[0].message.content or ""
    
    async def _merge_summaries(
        self, 
        batch_summaries: List[Dict[str, Any]], 
        topic: str,
        stream: bool = True,
        custom_requirements: Optional[str] = None
    ) -> str:
        """合并各批次总结，生成最终综述"""
        # 组合所有批次总结
        summaries_text = ""
        for bs in batch_summaries:
            summaries_text += f"\n### 批次 {bs['batch_num']} (项目 {bs['project_indices'][0]}-{bs['project_indices'][-1]})\n"
            summaries_text += bs['summary'] + "\n"
        
        prompt = self._build_merge_prompt(summaries_text, topic, custom_requirements)
        
        if stream:
            return await self._stream_generate(prompt)
        else:
            response = await self.client.chat.completions.create(
                model=self.model,
                messages=[{"role": "user", "content": prompt}],
                temperature=0.7,
                max_tokens=4000,
            )
            return response.choices[0].message.content or ""
    
    async def _generate_final_review(
        self, 
        project_info: str, 
        topic: str,
        stream: bool = True,
        custom_requirements: Optional[str] = None
    ) -> str:
        """直接生成最终综述（项目数少时使用）"""
        prompt = self._build_review_prompt(project_info, topic, custom_requirements)
        
        if stream:
            return await self._stream_generate(prompt)
        else:
            response = await self.client.chat.completions.create(
                model=self.model,
                messages=[{"role": "user", "content": prompt}],
                temperature=0.7,
                max_tokens=4000,
            )
            return response.choices[0].message.content or ""
    
    async def _stream_generate(self, prompt: str) -> str:
        """流式生成并实时输出"""
        print()  # 换行
        full_content = ""
        
        stream = await self.client.chat.completions.create(
            model=self.model,
            messages=[{"role": "user", "content": prompt}],
            temperature=0.7,
            max_tokens=4000,
            stream=True,
        )
        
        async for chunk in stream:
            # 安全检查：确保 choices 存在且非空
            if chunk.choices and len(chunk.choices) > 0:
                delta = chunk.choices[0].delta
                if delta and delta.content:
                    print(delta.content, end="", flush=True)
                    full_content += delta.content
        
        print()  # 结束换行
        return full_content
    
    def _build_review_prompt(self, project_info: str, topic: str, custom_requirements: Optional[str] = None) -> str:
        """构建综述生成 prompt"""
        custom_section = ""
        if custom_requirements:
            custom_section = f"\n## 用户特别要求\n{custom_requirements}\n"
        
        return f"""你是一位资深科研工作者，请根据以下国家自然科学基金项目信息，撰写一份关于"{topic}"的研究综述。

## 项目信息
{project_info}
{custom_section}
## 综述要求
请按以下结构撰写综述：

### 1. 研究背景与意义
- 简述该领域的研究背景
- 说明研究的重要性和意义

### 2. 研究现状分析
- 归纳这些项目涉及的主要研究方向
- 分析各研究方向的特点和进展
- 指出主要研究机构和团队

### 3. 关键技术与方法
- 总结这些项目采用的主要技术方法
- 分析技术发展趋势

### 4. 研究热点与趋势
- 识别当前的研究热点
- 预测未来发展趋势

### 5. 研究空白与机遇
- 指出现有研究的不足
- 提出可能的创新方向
- 给出具体的研究建议

### 6. 总结
- 概括主要发现
- 展望未来研究方向

## 输出要求
- 使用中文撰写
- 内容要有深度，结合具体项目分析
- 引用项目时使用编号，如 [1]、[2-3]
- 每个部分都要有实质性内容
- 总字数 2000-3000 字
"""
    
    def _build_merge_prompt(self, summaries_text: str, topic: str, custom_requirements: Optional[str] = None) -> str:
        """构建合并综述 prompt"""
        custom_section = ""
        if custom_requirements:
            custom_section = f"\n## 用户特别要求\n{custom_requirements}\n"
        
        return f"""你是一位资深科研工作者，请根据以下各批次项目分析结果，撰写一份关于"{topic}"的完整研究综述。

## 各批次分析结果
{summaries_text}
{custom_section}
## 综述要求
请整合以上分析，按以下结构撰写完整综述：

### 1. 研究背景与意义
### 2. 研究现状分析
### 3. 关键技术与方法
### 4. 研究热点与趋势
### 5. 研究空白与机遇
### 6. 总结

## 输出要求
- 使用中文撰写
- 整合各批次信息，避免重复
- 保留项目编号引用，如 [1]、[5-8]
- 内容要有深度和逻辑性
- 总字数 2000-3000 字
"""

    async def _generate_review_content(self, project_info: str, topic: str) -> str:
        """调用 LLM 生成综述"""
        prompt = f"""你是一位资深科研工作者，请根据以下国家自然科学基金项目信息，撰写一份关于"{topic}"的研究综述。

## 项目信息
{project_info}

## 综述要求
请按以下结构撰写综述：

### 1. 研究背景与意义
- 简述该领域的研究背景
- 说明研究的重要性和意义

### 2. 研究现状分析
- 归纳这些项目涉及的主要研究方向
- 分析各研究方向的特点和进展
- 指出主要研究机构和团队

### 3. 关键技术与方法
- 总结这些项目采用的主要技术方法
- 分析技术发展趋势

### 4. 研究热点与趋势
- 识别当前的研究热点
- 预测未来发展趋势

### 5. 研究空白与机遇
- 指出现有研究的不足
- 提出可能的创新方向
- 给出具体的研究建议

### 6. 总结
- 概括主要发现
- 展望未来研究方向

## 输出要求
- 使用中文撰写
- 内容要有深度，结合具体项目分析
- 每个部分都要有实质性内容
- 总字数 2000-3000 字
"""
        
        response = await self.client.chat.completions.create(
            model=self.model,
            messages=[{"role": "user", "content": prompt}],
            temperature=0.7,
            max_tokens=4000,
        )
        return response.choices[0].message.content or ""
    
    def _save_review(
        self, 
        review: str, 
        projects: List[ProjectInfo],
        topic: str,
        output_path: str
    ):
        """保存综述到 docx 文件"""
        path = Path(output_path)
        # 确保扩展名为 .docx
        if path.suffix.lower() != '.docx':
            path = path.with_suffix('.docx')
        path.parent.mkdir(parents=True, exist_ok=True)
        
        doc = Document()
        
        # 标题
        title = doc.add_heading(f'{topic} - 研究综述报告', level=0)
        title.alignment = WD_ALIGN_PARAGRAPH.CENTER
        
        # 元信息
        meta = doc.add_paragraph()
        meta.add_run(f'生成时间: {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}\n')
        meta.add_run(f'项目数量: {len(projects)}')
        
        doc.add_paragraph()  # 空行
        
        # 解析 Markdown 内容并添加到文档
        self._add_markdown_to_doc(doc, review)
        
        # 添加分隔线
        doc.add_paragraph('─' * 50)
        
        # 附录：项目列表
        doc.add_heading('附录：项目列表', level=1)
        
        # 创建表格
        table = doc.add_table(rows=1, cols=5)
        table.style = 'Table Grid'
        
        # 表头
        header_cells = table.rows[0].cells
        headers = ['序号', '项目名称', '负责人', '依托单位', '年度']
        for i, header in enumerate(headers):
            header_cells[i].text = header
            header_cells[i].paragraphs[0].runs[0].bold = True
        
        # 表格内容
        for i, p in enumerate(projects, 1):
            row = table.add_row().cells
            row[0].text = str(i)
            row[1].text = p.project_name[:50] + '...' if len(p.project_name) > 50 else p.project_name
            row[2].text = p.project_leader
            row[3].text = p.institution
            row[4].text = p.approval_year or '-'
        
        doc.save(str(path))
    
    def _add_markdown_to_doc(self, doc, markdown_text: str):
        """将 Markdown 文本解析并添加到 Word 文档"""
        lines = markdown_text.split('\n')
        
        for line in lines:
            line = line.rstrip()
            
            # 跳过空行
            if not line:
                doc.add_paragraph()
                continue
            
            # 标题处理
            if line.startswith('### '):
                doc.add_heading(line[4:], level=2)
            elif line.startswith('## '):
                doc.add_heading(line[3:], level=1)
            elif line.startswith('# '):
                doc.add_heading(line[2:], level=0)
            # 列表项
            elif line.startswith('- '):
                p = doc.add_paragraph(line[2:], style='List Bullet')
            elif re.match(r'^\d+\.\s', line):
                text = re.sub(r'^\d+\.\s', '', line)
                p = doc.add_paragraph(text, style='List Number')
            # 普通段落
            else:
                p = doc.add_paragraph()
                self._add_formatted_text(p, line)
    
    def _add_formatted_text(self, paragraph, text: str):
        """处理文本中的粗体和斜体"""
        # 简单处理：移除 ** 和 * 标记，直接添加文本
        # 更复杂的格式化可以后续扩展
        text = re.sub(r'\*\*([^*]+)\*\*', r'\1', text)  # 移除粗体标记
        text = re.sub(r'\*([^*]+)\*', r'\1', text)      # 移除斜体标记
        paragraph.add_run(text)
