skills/ppt-station-skill/ppt_station/chart_builder/parser.py

"""
图表解析器 - ChartBuilder 的反向操作

这是 ChartBuilder 的逆向工程：
- ChartBuilder: Config -> XML
- ChartParser: XML -> Config

核心职责：
1. 解析坐标轴映射（主轴/次轴）
2. 解析绘图区中的所有系列
3. 提取嵌入的 Excel 数据
4. 生成可用于 ChartBuilder 的配置
"""

from typing import Dict, List, Tuple, Optional
from lxml import etree
from pptx import Presentation
from pptx.chart.chart import Chart
import pandas as pd
import io
from openpyxl import load_workbook

from .oxml_ns import NAMESPACES


class ChartParser:
    """
    图表解析器 - 从现有图表提取配置
    
    使用场景：
    1. 从"黄金标准"PPTX 提取图表配置
    2. 批量复现图表
    3. 数据更新和刷新
    """
    
    def __init__(self, chart: Chart):
        """
        初始化解析器
        
        Args:
            chart: python-pptx Chart 对象
        """
        self.chart = chart
        self.chart_element = chart._element
        self.plotArea = self.chart_element.find('.//c:plotArea', namespaces=NAMESPACES)
        
        # 解析坐标轴映射
        self.axis_map = self._parse_axes()
        
        print("\n📖 ChartParser 初始化完成")
        print(f"  - 发现坐标轴: {len(self.axis_map)} 个")
        for ax_id, ax_info in self.axis_map.items():
            print(f"    • {ax_id}: {ax_info['type']} ({ax_info.get('position', 'N/A')})")
    
    def _parse_axes(self) -> Dict[str, Dict]:
        """
        解析所有坐标轴，建立 ID 映射表
        
        Returns:
            {
                "123456": {"type": "category", "position": None},
                "789001": {"type": "primary", "position": "l"},
                "789002": {"type": "secondary", "position": "r"}
            }
        """
        axis_map = {}
        
        # 1. 解析所有值轴 (Y轴)
        val_axes = self.plotArea.findall('.//c:valAx', namespaces=NAMESPACES)
        for val_ax in val_axes:
            ax_id_elem = val_ax.find('.//c:axId', namespaces=NAMESPACES)
            if ax_id_elem is None:
                continue
            
            ax_id = ax_id_elem.get('val')
            
            # 查找轴的位置
            ax_pos_elem = val_ax.find('.//c:axPos', namespaces=NAMESPACES)
            position = ax_pos_elem.get('val') if ax_pos_elem is not None else 'l'
            
            # 判断主轴还是次轴
            # 'r' (right) = 次轴, 'l' (left) = 主轴
            if position == 'r':
                axis_map[ax_id] = {"type": "secondary", "position": position}
            else:
                axis_map[ax_id] = {"type": "primary", "position": position}
        
        # 2. 解析所有分类轴 (X轴)
        cat_axes = self.plotArea.findall('.//c:catAx', namespaces=NAMESPACES)
        for cat_ax in cat_axes:
            ax_id_elem = cat_ax.find('.//c:axId', namespaces=NAMESPACES)
            if ax_id_elem is None:
                continue
            
            ax_id = ax_id_elem.get('val')
            axis_map[ax_id] = {"type": "category", "position": None}
        
        return axis_map
    
    def parse(self) -> Tuple[List[Dict], pd.DataFrame, str, Optional[Dict]]:
        """
        解析图表配置和数据
        
        Returns:
            (series_config, df, categories_col, layout_info)
            
            series_config: 系列配置列表
            df: 数据 DataFrame
            categories_col: 分类列名
            layout_info: 布局信息（图例、轴配置等）
        """
        print("\n" + "=" * 80)
        print("🔍 开始解析图表")
        print("=" * 80)
        
        # 1. 提取嵌入的 Excel 数据
        df, categories_col = self._extract_embedded_data()
        
        print(f"\n📊 数据提取完成:")
        print(f"  - 行数: {len(df)}")
        print(f"  - 列数: {len(df.columns)}")
        print(f"  - 分类列: {categories_col}")
        print(f"  - 数据列: {list(df.columns[1:])}")
        
        # 2. 解析系列配置
        series_config = self._parse_series_config()
        
        print(f"\n📋 系列配置解析完成:")
        for i, cfg in enumerate(series_config):
            print(f"  {i+1}. {cfg['name']}: type={cfg['type']}, axis={cfg['axis']}, key={cfg['key']}")
        
        # 3. ⭐ 新增：解析布局配置
        layout_info = self._parse_layout_info()
        
        if layout_info:
            print(f"\n⚙️ 布局配置:")
            if 'legend' in layout_info:
                leg = layout_info['legend']
                print(f"  - 图例: 位置={leg.get('position')}, 字体={leg.get('font_size_pt')}pt")
            if 'category_axis' in layout_info:
                cat = layout_info['category_axis']
                print(f"  - 横轴: 类型={cat.get('type')}, 格式={cat.get('number_format')}")
        
        print("\n" + "=" * 80)
        print("✅ 图表解析完成！")
        print("=" * 80)
        
        return series_config, df, categories_col, layout_info
    
    def _parse_layout_info(self) -> Dict:
        """
        解析布局信息（图例、轴配置、标题、尺寸等）
        
        Returns:
            布局信息字典
        """
        layout_info = {}
        
        try:
            # 1. 解析图表标题
            try:
                if self.chart.has_title:
                    title_frame = self.chart.chart_title.text_frame
                    if title_frame and title_frame.text:
                        layout_info["title"] = {
                            "text": title_frame.text,
                            "has_title": True,
                        }
            except Exception as e:
                print(f"  ⚠️ 标题解析失败: {e}")
            
            # 2. 解析图例
            if self.chart.has_legend:
                legend = self.chart.legend
                layout_info["legend"] = {
                    "position": legend.position,
                    "font_size_pt": legend.font.size.pt if legend.font.size else None,
                    "include_in_layout": legend.include_in_layout,
                }
            
            # 3. 解析横轴（分类轴）
            try:
                category_axis = self.chart.category_axis
                layout_info["category_axis"] = {
                    "type": category_axis.category_type,
                    "major_unit": category_axis.major_unit if hasattr(category_axis, 'major_unit') else None,
                    "number_format": category_axis.tick_labels.number_format if hasattr(category_axis.tick_labels, 'number_format') else None,
                    "font_size_pt": category_axis.tick_labels.font.size.pt if category_axis.tick_labels.font.size else None,
                }
            except Exception as e:
                print(f"  ⚠️ 横轴解析失败: {e}")
            
            # 4. 解析纵轴（值轴）
            try:
                value_axis = self.chart.value_axis
                layout_info["value_axis"] = {
                    "number_format": value_axis.tick_labels.number_format if hasattr(value_axis.tick_labels, 'number_format') else None,
                    "font_size_pt": value_axis.tick_labels.font.size.pt if value_axis.tick_labels.font.size else None,
                    "has_major_gridlines": value_axis.has_major_gridlines if hasattr(value_axis, 'has_major_gridlines') else None,
                }
            except Exception as e:
                print(f"  ⚠️ 纵轴解析失败: {e}")
                
        except Exception as e:
            print(f"  ⚠️ 布局信息解析失败: {e}")
        
        return layout_info
    
    def _extract_embedded_data(self) -> Tuple[pd.DataFrame, str]:
        """
        提取嵌入在图表中的 Excel 数据
        
        Returns:
            (df, categories_col)
        """
        try:
            # 方法1: 尝试从 chart 的内部部分获取 Excel 数据
            # python-pptx 的图表对象包含嵌入的 Excel workbook
            chart_part = self.chart.part
            
            # 查找 Excel workbook 关系
            for rel in chart_part.rels.values():
                if 'package' in rel.reltype or 'oleObject' in rel.reltype or 'embeddings' in rel.reltype:
                    try:
                        xlsx_blob = rel.target_part.blob
                        wb = load_workbook(io.BytesIO(xlsx_blob))
                        ws = wb.active
                        
                        # 转换为 DataFrame
                        data = []
                        for row in ws.iter_rows(values_only=True):
                            data.append(row)
                        
                        if data:
                            # 第一行是表头
                            df = pd.DataFrame(data[1:], columns=data[0])
                            categories_col = df.columns[0]
                            print(f"  ✅ 成功从嵌入的 Excel 提取数据")
                            return df, categories_col
                    except Exception as e2:
                        continue
            
            # 如果上面都失败，尝试从 XML 中提取
            print(f"  ⚠️ 未找到嵌入的 Excel 数据，从 XML 提取")
            return self._extract_data_from_xml()
            
        except Exception as e:
            print(f"  ⚠️ 提取嵌入数据失败: {e}")
            # 如果失败，尝试从 XML 中提取数据点
            return self._extract_data_from_xml()
    
    def _extract_data_from_xml(self) -> Tuple[pd.DataFrame, str]:
        """
        从 XML 中提取数据点（备用方法）
        
        Returns:
            (df, categories_col)
        """
        print("  → 尝试从 XML 提取数据...")
        from datetime import datetime, timedelta
        
        # 查找所有绘图元素
        plot_types = ['barChart', 'lineChart', 'areaChart', 'scatterChart']
        
        data_dict = {}
        categories = []
        categories_col = "分类"
        max_data_points = 0
        
        # 第一遍：找到所有系列和最大数据点数量
        for plot_tag in plot_types:
            plot_elements = self.plotArea.findall(f'.//c:{plot_tag}', namespaces=NAMESPACES)
            
            for plot_element in plot_elements:
                # 提取分类数据（只需提取一次）
                if not categories:
                    cat_elem = plot_element.find('.//c:cat', namespaces=NAMESPACES)
                    if cat_elem is not None:
                        # ⭐ 优先尝试 strCache（文本）
                        cat_cache = cat_elem.find('.//c:strCache', namespaces=NAMESPACES)
                        if cat_cache is not None:
                            # ⭐ 修复 A: 强制将字符串转为 datetime
                            print("  → 发现 strCache，尝试转换为 datetime...")
                            for pt in cat_cache.findall('.//c:pt', namespaces=NAMESPACES):
                                v = pt.find('.//c:v', namespaces=NAMESPACES)
                                if v is not None:
                                    try:
                                        # 尝试用 pandas 智能解析多种日期格式
                                        dt = pd.to_datetime(v.text).to_pydatetime()
                                        categories.append(dt)
                                        print(f"    • 转换: {v.text} → {dt}")
                                    except (ValueError, TypeError):
                                        categories.append(v.text)  # 回退到字符串
                            if categories:
                                print(f"  → 成功将 {len(categories)} 个字符串日期转换为 datetime")
                        else:
                            # ⭐ 如果没有 strCache，尝试 numCache（日期轴）
                            num_cache = cat_elem.find('.//c:numCache', namespaces=NAMESPACES)
                            if num_cache is not None:
                                # 提取格式代码（如 "yyyy/mm"）
                                format_code = num_cache.find('.//c:formatCode', namespaces=NAMESPACES)
                                date_format = format_code.text if format_code is not None else None
                                
                                # ⭐ 修复 B: 强制将 Excel 序列号转为 datetime
                                print(f"  → 发现 numCache（格式: {date_format}），转换为 datetime...")
                                for pt in num_cache.findall('.//c:pt', namespaces=NAMESPACES):
                                    v = pt.find('.//c:v', namespaces=NAMESPACES)
                                    if v is not None:
                                        try:
                                            excel_date_num = float(v.text)
                                            # Excel 基准是 1899-12-30
                                            dt = datetime(1899, 12, 30) + timedelta(days=excel_date_num)
                                            categories.append(dt)
                                        except (ValueError, OverflowError):
                                            categories.append(v.text)  # 回退到字符串
                                
                                if categories:
                                    print(f"  → 成功将 {len(categories)} 个 Excel 序列号转换为 datetime")
                
                # 提取每个系列的数据
                series_elements = plot_element.findall('.//c:ser', namespaces=NAMESPACES)
                for ser in series_elements:
                    # 获取系列名称
                    name = self._extract_series_name(ser)
                    
                    # 获取数值
                    values = self._extract_series_values(ser)
                    
                    if name and values:
                        data_dict[name] = values
                        max_data_points = max(max_data_points, len(values))
        
        # 如果没有分类数据，使用最大数据点数量生成索引
        if not categories and max_data_points > 0:
            categories = [f"项目{i+1}" for i in range(max_data_points)]
            print(f"  → 未找到分类数据，生成 {len(categories)} 个默认分类")
        elif not categories and not data_dict:
            # 如果完全没有数据，返回空 DataFrame
            print("  ⚠️ 未找到任何数据，返回空 DataFrame")
            return pd.DataFrame({categories_col: []}), categories_col
        
        # 确保所有系列的长度与分类数量一致
        target_length = len(categories)
        for name, values in data_dict.items():
            if len(values) < target_length:
                # 补齐 None
                data_dict[name] = values + [None] * (target_length - len(values))
            elif len(values) > target_length:
                # 截断
                data_dict[name] = values[:target_length]
                print(f"  → 系列 '{name}' 数据点从 {len(values)} 截断到 {target_length}")
        
        # 构建 DataFrame
        df_data = {categories_col: categories}
        df_data.update(data_dict)
        df = pd.DataFrame(df_data)
        
        print(f"  → 从 XML 提取了 {len(categories)} 行数据，{len(data_dict)} 个系列")
        
        return df, categories_col
    
    def _parse_series_config(self) -> List[Dict]:
        """
        解析系列配置
        
        Returns:
            系列配置列表
        """
        series_config = []
        
        # 图表类型映射
        plot_type_map = {
            'barChart': 'bar',
            'lineChart': 'line',
            'areaChart': 'area',
            'scatterChart': 'scatter'
        }
        
        # 遍历所有图表类型
        for plot_tag, chart_type in plot_type_map.items():
            plot_elements = self.plotArea.findall(f'.//c:{plot_tag}', namespaces=NAMESPACES)
            
            for plot_element in plot_elements:
                # ⭐ 关键修复：在绘图元素级别获取轴引用
                plot_axis_type = self._extract_plot_axis(plot_element)
                
                # 遍历该绘图中的所有系列
                series_elements = plot_element.findall('.//c:ser', namespaces=NAMESPACES)
                
                for ser in series_elements:
                    # 解析单个系列
                    config_entry = self._parse_single_series(ser, chart_type, plot_axis_type)
                    if config_entry:
                        series_config.append(config_entry)
        
        return series_config
    
    def _extract_plot_axis(self, plot_element) -> str:
        """
        提取绘图元素使用的坐标轴（主轴/次轴）
        
        Args:
            plot_element: 绘图元素 (<c:barChart>, <c:lineChart> 等)
            
        Returns:
            'primary' 或 'secondary'
        """
        # 查找绘图元素引用的轴 ID
        ax_id_refs = plot_element.findall('.//c:axId', namespaces=NAMESPACES)
        
        for ax_id_ref in ax_id_refs:
            ax_id = ax_id_ref.get('val')
            
            if ax_id in self.axis_map:
                axis_info = self.axis_map[ax_id]
                axis_type = axis_info['type']
                
                # 如果是值轴（主轴或次轴），返回它
                if axis_type in ('primary', 'secondary'):
                    return axis_type
        
        # 默认返回主轴
        return 'primary'
    
    def _parse_single_series(self, ser, chart_type: str, axis_type: str) -> Optional[Dict]:
        """
        解析单个系列
        
        Args:
            ser: 系列元素 (<c:ser>)
            chart_type: 图表类型
            axis_type: 轴类型（从父绘图元素获取）
            
        Returns:
            系列配置字典，或 None
        """
        config_entry = {}
        
        # A. 获取系列名称
        name = self._extract_series_name(ser)
        if not name:
            return None
        
        config_entry["name"] = name
        config_entry["key"] = name  # 默认使用 name 作为 key
        
        # B. 设置图表类型
        config_entry["type"] = chart_type
        
        # C. 使用传入的轴分配
        config_entry["axis"] = axis_type
        
        return config_entry
    
    def _extract_series_name(self, ser) -> Optional[str]:
        """
        提取系列名称
        
        Args:
            ser: 系列元素 (<c:ser>)
            
        Returns:
            系列名称，或 None
        """
        # 方法1: 从 <c:tx> 中提取
        tx_elem = ser.find('.//c:tx', namespaces=NAMESPACES)
        if tx_elem is not None:
            # 尝试从 strRef 中提取
            v_elem = tx_elem.find('.//c:v', namespaces=NAMESPACES)
            if v_elem is not None and v_elem.text:
                return v_elem.text
        
        return None
    
    def _extract_series_values(self, ser) -> List:
        """
        提取系列数值
        
        Args:
            ser: 系列元素 (<c:ser>)
            
        Returns:
            数值列表
        """
        values = []
        
        # 尝试从 <c:val> 提取（柱状图、折线图、面积图）
        val_elem = ser.find('.//c:val', namespaces=NAMESPACES)
        if val_elem is not None:
            num_cache = val_elem.find('.//c:numCache', namespaces=NAMESPACES)
            if num_cache is not None:
                for pt in num_cache.findall('.//c:pt', namespaces=NAMESPACES):
                    v = pt.find('.//c:v', namespaces=NAMESPACES)
                    if v is not None:
                        try:
                            values.append(float(v.text))
                        except (ValueError, TypeError):
                            values.append(0)
        
        # 尝试从 <c:yVal> 提取（散点图）
        if not values:
            yVal_elem = ser.find('.//c:yVal', namespaces=NAMESPACES)
            if yVal_elem is not None:
                num_cache = yVal_elem.find('.//c:numCache', namespaces=NAMESPACES)
                if num_cache is not None:
                    for pt in num_cache.findall('.//c:pt', namespaces=NAMESPACES):
                        v = pt.find('.//c:v', namespaces=NAMESPACES)
                        if v is not None:
                            try:
                                values.append(float(v.text))
                            except (ValueError, TypeError):
                                values.append(0)
        
        return values


def parse_chart_from_pptx(pptx_path: str, slide_idx: int = 0, shape_idx: int = 0) -> Tuple[List[Dict], pd.DataFrame, str]:
    """
    从 PPTX 文件解析图表
    
    Args:
        pptx_path: PPTX 文件路径
        slide_idx: 幻灯片索引（0-based）
        shape_idx: 图表形状索引（0-based）
        
    Returns:
        (series_config, df, categories_col)
        
    Example:
        >>> series_config, df, categories_col = parse_chart_from_pptx("template.pptx")
        >>> # 现在可以用这些配置来复现图表
        >>> builder.build(series_config)
    """
    # 加载 PPTX
    prs = Presentation(pptx_path)
    
    if slide_idx >= len(prs.slides):
        raise ValueError(f"幻灯片索引 {slide_idx} 超出范围（共 {len(prs.slides)} 张）")
    
    slide = prs.slides[slide_idx]
    
    # 查找图表
    chart = None
    chart_shape = None
    
    for shape in slide.shapes:
        if hasattr(shape, 'chart'):
            if shape_idx == 0:
                chart = shape.chart
                chart_shape = shape
                break
            shape_idx -= 1
    
    if chart is None:
        raise ValueError(f"在幻灯片 {slide_idx} 上未找到图表")
    
    # 解析图表
    parser = ChartParser(chart)
    return parser.parse()


def parse_all_charts_from_pptx(pptx_path: str) -> List[Dict]:
    """
    解析 PPTX 文件中的所有图表
    
    Args:
        pptx_path: PPTX 文件路径
        
    Returns:
        图表信息列表
        [
            {
                "slide_idx": 0,
                "shape_idx": 0,
                "shape_name": "图表 1",
                "series_config": [...],
                "df": DataFrame,
                "categories_col": "日期"
            },
            ...
        ]
    """
    prs = Presentation(pptx_path)
    all_charts = []
    
    print(f"\n🔍 扫描 PPTX 文件: {pptx_path}")
    print(f"  - 共 {len(prs.slides)} 张幻灯片\n")
    
    for slide_idx, slide in enumerate(prs.slides):
        chart_count = 0
        
        for shape_idx, shape in enumerate(slide.shapes):
            # 使用 has_chart 属性检查，而不是直接访问 chart
            if hasattr(shape, 'has_chart') and shape.has_chart:
                chart_count += 1
                
                try:
                    # 解析图表
                    parser = ChartParser(shape.chart)
                    series_config, df, categories_col, layout_info = parser.parse()
                    
                    chart_info = {
                        "slide_index": slide_idx,
                        "chart_index": chart_count - 1,  # 0-based
                        "shape_idx": shape_idx,
                        "shape_name": shape.name if hasattr(shape, 'name') else f"图表 {chart_count}",
                        "series_config": series_config,
                        "df": df,
                        "categories_col": categories_col,
                        "layout_info": layout_info  # 新增：布局信息
                    }
                    
                    all_charts.append(chart_info)
                    
                    print(f"✅ 幻灯片 {slide_idx + 1}, 图表 {chart_count}: {chart_info['shape_name']}")
                    print(f"   - {len(series_config)} 个系列")
                    
                except Exception as e:
                    print(f"⚠️ 幻灯片 {slide_idx + 1}, 图表 {chart_count} 解析失败: {e}")
        
        if chart_count > 0:
            print(f"  → 幻灯片 {slide_idx + 1}: 发现 {chart_count} 个图表\n")
    
    print(f"\n📊 总计解析 {len(all_charts)} 个图表")
    
    return all_charts