import os
import traceback
from paste.chart.pie import gen_pie


class ChartPieExample:
    """
    环形图测试管理器：封装对 paste.chart.pie.gen_pie 的调用。
    数据结构严格匹配函数参数要求，支持扩展更多测试用例。
    """

    def __init__(self, output_directory="./charts"):
        """
        初始化测试器，定义所有测试数据。
        数据结构完全匹配 gen_pie 函数的参数要求。
        """
        self.output_directory = output_directory
        os.makedirs(self.output_directory, exist_ok=True)

        # 构造符合 gen_pie 要求的 DataFrame 数据（模拟真实业务场景）
        # 假设业务场景：网络设备统计（服务器、交换机、路由器等）
        import pandas as pd
        self.data_df = pd.DataFrame({
            'device_count': [35, 28, 22, 15, 10],           # value_column
            'percentage': ['35.2%', '28.1%', '22.0%', '15.0%', '9.7%'],  # percentage_column
            'device_type': ['服务器', '交换机', '路由器', '防火墙', 'AP']  # legend_labels
        })

        # 测试参数
        self.value_column = 'device_count'
        self.percentage_column = 'percentage'
        self.legend_labels = 'device_type'
        self.color_palette = 'BuPu'   # 可尝试 'viridis', 'Set3', 'plasma'
        self.dpi = 128

        # 输出文件名
        self.filename = "pie_chart.svg"

    def generate_pie_chart(self) -> str:
        """
        调用 gen_pie 函数，参数完全一致。
        注意：gen_pie 接收的是 pandas.DataFrame，不是列表。
        """
        try:
            svg_data = gen_pie(
                data_df=self.data_df,
                value_column=self.value_column,
                percentage_column=self.percentage_column,
                legend_labels=self.legend_labels,
                color_palette=self.color_palette,
                dpi=self.dpi
            )
            if not svg_data or not isinstance(svg_data, str):
                raise ValueError("gen_pie 返回的 SVG 数据为空或类型错误")
            return svg_data
        except Exception as e:
            print(f"环形图生成失败: {e}")
            traceback.print_exc()
            raise

    def save_svg(self, svg_data: str, filename: str) -> None:
        """
        将 SVG 的 base64 Data URL 写入文件（保留原始 SVG 格式）。
        注意：svg_data 是 "data:image/svg+xml;base64,..."，需提取真实 SVG 内容。
        """
        if not svg_data or not isinstance(svg_data, str):
            print(f"生成的 SVG 数据无效（为空或非字符串）: {filename}")
            return

        # 提取 base64 编码部分（去除 data URL 前缀）
        if svg_data.startswith("data:image/svg+xml;base64,"):
            base64_content = svg_data[len("data:image/svg+xml;base64,"):]
            try:
                # 解码 base64 得到原始 SVG 字符串
                import base64
                svg_content = base64.b64decode(base64_content).decode('utf-8')
            except Exception as e:
                print(f"解码 base64 失败: {e}")
                svg_content = svg_data  # 退化为直接写入
        else:
            # 如果不是标准格式，直接写入（兼容调试）
            svg_content = svg_data

        filepath = os.path.join(self.output_directory, filename)
        with open(filepath, 'w', encoding='utf-8') as f:
            f.write(svg_content)
        print(f"已保存: {filepath}")

    def run(self) -> None:
        """按顺序执行图表生成与保存"""
        print("开始生成环形图...")
        try:
            print("正在生成环形图...")
            svg_data = self.generate_pie_chart()
            self.save_svg(svg_data, self.filename)

            print(f"\n环形图已成功生成。")
            print(f"输出目录: {self.output_directory}")
            print(f"文件列表:")
            print(f"  - {self.filename}")

        except Exception as e:
            print(f"\n测试失败: {e}")
            traceback.print_exc()


# 程序入口
if __name__ == "__main__":
    tester = ChartPieExample()
    tester.run()