import logging
from typing import Dict, List, Optional, Union
from reportlab.lib import colors
from reportlab.lib.pagesizes import A4
from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, Table, TableStyle
from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
from reportlab.pdfbase import pdfmetrics
from reportlab.pdfbase.ttfonts import TTFont
from reportlab.platypus.paragraph import Paragraph
import markdown2
from bs4 import BeautifulSoup
import os
from datetime import datetime
import shutil

import matplotlib
matplotlib.use('Agg')

# 修改导入路径，使用相对导入
try:
    # 尝试相对导入
    from .chat_bot import ChatBot
    from .chat_bot_with_offline import ChatBot as OfflineChatBot
except ImportError:
    # 如果相对导入失败，尝试绝对导入
    try:
        from src.fundamentals_llm.chat_bot import ChatBot
        from src.fundamentals_llm.chat_bot_with_offline import ChatBot as OfflineChatBot
    except ImportError:
        # 最后尝试直接导入
        from chat_bot import ChatBot
        from chat_bot_with_offline import ChatBot as OfflineChatBot

# 设置日志记录
logger = logging.getLogger(__name__)

class PDFGenerator:
    """PDF生成器类，用于将内容转换为PDF格式"""
    
    def __init__(self):
        """初始化PDF生成器"""
        # 尝试注册中文字体
        self.font_name = self._register_chinese_font()
        self.chat_bot = OfflineChatBot(platform="volc", model_type="offline_model")  # 使用离线模式
        self.styles = getSampleStyleSheet()
        
        # 创建自定义样式
        self.styles.add(ParagraphStyle(
            name='ReportTitle',
            parent=self.styles['Heading1'],
            fontSize=24,
            spaceAfter=30,
            alignment=1,  # 居中对齐
            fontName=self.font_name
        ))
        
        self.styles.add(ParagraphStyle(
            name='ReportHeading',
            parent=self.styles['Heading2'],
            fontSize=16,
            spaceAfter=12,
            textColor=colors.HexColor('#2C3E50'),
            fontName=self.font_name
        ))
        
        self.styles.add(ParagraphStyle(
            name='ReportBody',
            parent=self.styles['Normal'],
            fontSize=12,
            leading=18,
            spaceAfter=12,
            fontName=self.font_name
        ))
        
        # 添加Markdown样式
        self.styles.add(ParagraphStyle(
            name='MarkdownHeading1',
            parent=self.styles['Heading1'],
            fontSize=20,
            spaceAfter=20,
            fontName=self.font_name
        ))
        
        self.styles.add(ParagraphStyle(
            name='MarkdownHeading2',
            parent=self.styles['Heading2'],
            fontSize=16,
            spaceAfter=16,
            fontName=self.font_name
        ))
        
        self.styles.add(ParagraphStyle(
            name='MarkdownHeading3',
            parent=self.styles['Heading3'],
            fontSize=14,
            spaceAfter=14,
            fontName=self.font_name
        ))
        
        self.styles.add(ParagraphStyle(
            name='MarkdownList',
            parent=self.styles['Normal'],
            fontSize=12,
            leading=18,
            spaceAfter=12,
            leftIndent=20,
            fontName=self.font_name
        ))
        
        self.styles.add(ParagraphStyle(
            name='MarkdownCode',
            parent=self.styles['Normal'],
            fontSize=10,
            leading=16,
            spaceAfter=12,
            leftIndent=20,
            fontName='Courier',
            textColor=colors.HexColor('#333333')
        ))
    
    def _register_chinese_font(self):
        """查找并注册中文字体
        
        Returns:
            str: 注册的字体名称，如果失败则使用默认字体
        """
        try:
            # 可能的字体文件位置列表
            font_locations = [
                # 当前目录/子目录字体位置
                os.path.join(os.path.dirname(__file__), "fonts", "simhei.ttf"),
                os.path.join(os.path.dirname(__file__), "fonts", "wqy-microhei.ttc"),
                os.path.join(os.path.dirname(__file__), "fonts", "wqy-zenhei.ttc"),
                # Docker中可能的位置
                "/app/src/fundamentals_llm/fonts/simhei.ttf",
                # Windows系统字体位置
                "C:/Windows/Fonts/simhei.ttf",
                "C:/Windows/Fonts/simfang.ttf",
                "C:/Windows/Fonts/simsun.ttc",
                # Linux/Mac系统字体位置
                "/usr/share/fonts/truetype/wqy/wqy-microhei.ttc",
                "/usr/share/fonts/truetype/wqy/wqy-zenhei.ttc",
                "/usr/share/fonts/wqy-microhei/wqy-microhei.ttc",
                "/usr/share/fonts/wqy-zenhei/wqy-zenhei.ttc",
                "/usr/share/fonts/truetype/droid/DroidSansFallback.ttf",
                "/usr/share/fonts/opentype/noto/NotoSansCJK-Regular.ttc",
                "/System/Library/Fonts/PingFang.ttc"  # macOS
            ]
            
            # 尝试各个位置
            for font_path in font_locations:
                if os.path.exists(font_path):
                    logger.info(f"使用字体文件: {font_path}")
                    # 为防止字体文件名称不同，统一拷贝到字体目录并重命名
                    font_dir = os.path.join(os.path.dirname(__file__), "fonts")
                    os.makedirs(font_dir, exist_ok=True)
                    target_path = os.path.join(font_dir, "simhei.ttf")
                    
                    # 如果字体文件不在目标位置，拷贝过去
                    if font_path != target_path and not os.path.exists(target_path):
                        try:
                            shutil.copy2(font_path, target_path)
                            logger.info(f"已将字体文件复制到: {target_path}")
                        except Exception as e:
                            logger.warning(f"复制字体文件失败: {str(e)}")
                    
                    # 注册字体
                    pdfmetrics.registerFont(TTFont('SimHei', target_path))
                    return 'SimHei'
            
            # 如果所有位置都找不到，使用默认字体
            logger.warning("找不到中文字体文件，将使用默认字体")
            return 'Helvetica'
            
        except Exception as e:
            logger.error(f"注册中文字体失败: {str(e)}")
            return 'Helvetica'  # 使用默认字体
    
    def _convert_markdown_to_flowables(self, markdown_text: str) -> List:
        """将Markdown文本转换为PDF流对象
        
        Args:
            markdown_text: Markdown格式的文本
            
        Returns:
            List: PDF流对象列表
        """
        # 将Markdown转换为HTML
        html = markdown2.markdown(markdown_text)
        soup = BeautifulSoup(html, 'html.parser')
        
        flowables = []
        for element in soup.children:
            if element.name is None:  # 纯文本
                if element.strip():
                    flowables.append(Paragraph(element.strip(), self.styles['ReportBody']))
            elif element.name in ['h1', 'h2', 'h3']:  # 标题
                level = int(element.name[1])
                style_name = f'MarkdownHeading{level}'
                flowables.append(Paragraph(element.get_text(), self.styles[style_name]))
            elif element.name == 'p':  # 段落
                flowables.append(Paragraph(element.get_text(), self.styles['ReportBody']))
            elif element.name == 'ul':  # 无序列表
                for li in element.find_all('li'):
                    flowables.append(Paragraph(f"• {li.get_text()}", self.styles['MarkdownList']))
            elif element.name == 'ol':  # 有序列表
                for i, li in enumerate(element.find_all('li'), 1):
                    flowables.append(Paragraph(f"{i}. {li.get_text()}", self.styles['MarkdownList']))
            elif element.name == 'code':  # 代码块
                flowables.append(Paragraph(element.get_text(), self.styles['MarkdownCode']))
            elif element.name == 'blockquote':  # 引用
                flowables.append(Paragraph(f"「{element.get_text()}」", self.styles['ReportBody']))
            elif element.name == 'hr':  # 分隔线
                flowables.append(Spacer(1, 20))
        
        return flowables
    
    def _optimize_content(self, content: str) -> str:
        # 移除markdown代码块标记
        content = content.replace('```markdown', '').replace('```', '')
        """优化内容格式并转换为Markdown
        
        Args:
            content: 原始内容
            
        Returns:
            str: 优化后的Markdown格式内容
        """
        try:
            prompt = f"""请对以下内容进行优化和格式化。要求：
                    1. 保持原文的专业性和准确性
                    2. 将零散的内容整合成连贯的段落，并对不重要的内容精简
                    3. 使用适当的标点符号和换行
                    4. 使用Markdown格式进行排版
                    5. 移除所有引用内容（包括"参考资料："等）
                    6. 不要返回其他多余的内容
                    7. 确保内容结构清晰，层次分明
                    8. 将零散的内容整合成完整的段落，避免过于零散的表述
                    9. 使用自然流畅的语言，避免过于机械的结构化表达
                    
                    原始内容：
                    {content}"""
            
            result = self.chat_bot.chat(prompt)
            return result
        except Exception as e:
            logger.error(f"优化内容失败: {str(e)}")
            return content
    
    def generate_pdf(self, 
                    title: str,
                    content_dict: Dict[str, str],
                    output_dir: str = "reports",
                    filename: Optional[str] = None) -> Optional[str]:
        """生成PDF报告
        
        Args:
            title: 报告标题
            content_dict: 包含各维度内容的字典，key为维度名称，value为内容
            output_dir: 输出目录
            filename: 文件名（可选），如果不提供则自动生成
            
        Returns:
            Optional[str]: 生成的PDF文件路径，如果失败则返回None
        """
        try:
            # 创建输出目录
            os.makedirs(output_dir, exist_ok=True)
            
            # 生成文件名
            if not filename:
                timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
                filename = f"report_{timestamp}.pdf"
            filepath = os.path.join(output_dir, filename)
            
            # 创建PDF文档
            doc = SimpleDocTemplate(
                filepath,
                pagesize=A4,
                rightMargin=72,
                leftMargin=72,
                topMargin=72,
                bottomMargin=72
            )
            
            # 准备内容
            story = []
            
            # 添加标题
            title_para = Paragraph(title, self.styles['ReportTitle'])
            story.append(title_para)
            story.append(Spacer(1, 20))
            
            # 添加生成时间
            time_str = datetime.now().strftime("%Y年%m月%d日 %H:%M:%S")
            time_para = Paragraph(f"生成时间：{time_str}", self.styles['ReportBody'])
            story.append(time_para)
            story.append(Spacer(1, 20))
            
            # 添加各维度内容
            for section_title, content in content_dict.items():
                if content:  # 只添加非空内容
                    # 优化内容并转换为Markdown
                    optimized_content = self._optimize_content(content)
                    # 转换为PDF流对象
                    # story.append(Paragraph(section_title, self.styles['ReportHeading']))
                    story.extend(self._convert_markdown_to_flowables(optimized_content))
                    story.append(Spacer(1, 20))
            
            # 生成PDF
            doc.build(story)
            logger.info(f"PDF报告已生成: {filepath}")
            return filepath
            
        except Exception as e:
            logger.error(f"生成PDF报告失败: {str(e)}")
            return None

# 使用示例
if __name__ == "__main__":
    # 设置日志级别
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
    )
    
    # 创建PDF生成器实例
    generator = PDFGenerator()
    
    # 示例内容
    test_content = {
        "公司简介": """# 公司简介
这是一段公司简介的内容...

## 主营业务
- 业务1
- 业务2
- 业务3

## 发展历程
1. 2020年成立
2. 2021年上市
3. 2022年扩张""",
        "财务分析": """# 财务分析
## 收入情况
2023年收入达到100亿...

## 利润分析
净利润同比增长20%...""",
        "行业分析": """# 行业分析
## 市场格局
目前市场呈现...""",
        "投资建议": """# 投资建议
## 优势
- 优势1
- 优势2

## 风险
- 风险1
- 风险2"""
    }
    
    # 生成PDF
    filepath = generator.generate_pdf(
        title="测试报告",
        content_dict=test_content,
        filename="test_report.pdf"
    )
    
    if filepath:
        print(f"报告已生成: {filepath}")
    else:
        print("报告生成失败") 