import json
import openpyxl
from openpyxl.styles import Font, Alignment, PatternFill, Border, Side
from openpyxl.utils import get_column_letter
from typing import List, Dict, Any, Set


def extract_all_fields(jsonl_data: List[Dict]) -> Set[str]:
    """提取JSONL文件中出现的所有字段名（只提取顶级字段）"""
    all_fields = set()
    
    for item in jsonl_data:
        # 只提取顶级字段
        all_fields.update(item.keys())
    
    return all_fields


def flatten_dict(obj: Dict, parent_key: str = '', sep: str = '.') -> Dict:
    """将嵌套字典展平为单层字典，但保留顶级字段"""
    items = []
    for k, v in obj.items():
        new_key = f"{parent_key}{sep}{k}" if parent_key else k
        if isinstance(v, dict):
            # 只展平嵌套的字典，保留其他类型
            if v:  # 如果字典不为空
                # 将字典转换为JSON字符串表示
                items.append((new_key, json.dumps(v, ensure_ascii=False)))
            else:
                items.append((new_key, ""))
        elif isinstance(v, list):
            # 将列表转换为JSON字符串表示
            items.append((new_key, json.dumps(v, ensure_ascii=False) if v else ""))
        else:
            items.append((new_key, v if v is not None else ""))
    return dict(items)

def save_xlsx_multi_sheet(
    data_sheets: List[tuple],  # [(data, sheet_name), (data, sheet_name), ...]
    output_file: str,
) -> str:
    """
    生成Excel格式的通用JSONL转换报告（多表版本）
    
    Args:
        data_sheets: 数据和对应工作表名的列表 [(data, sheet_name), ...]
        output_file: 输出的Excel文件路径
    
    Returns:
        Excel文件路径
    """
    try:
        # 创建工作簿
        wb = openpyxl.Workbook()
        
        # 删除默认创建的空工作表
        wb.remove(wb.active)
        
        # 定义样式
        header_font = Font(bold=True, size=12, color="FFFFFF")
        header_fill = PatternFill(start_color="4472C4", end_color="4472C4", fill_type="solid")
        thin_border = Border(
            left=Side(style='thin'),
            right=Side(style='thin'),
            top=Side(style='thin'),
            bottom=Side(style='thin')
        )
        
        total_records = 0
        total_sheets = 0
        
        # 遍历每个数据集和工作表名称
        for idx, (data, sheet_name) in enumerate(data_sheets):
            if not data:
                print(f"警告: 第{idx+1}个工作表 '{sheet_name}' 没有数据，跳过")
                continue
                
            # 提取所有字段
            all_fields = extract_all_fields(data)
            print(f"工作表 '{sheet_name}' 检测到字段: {all_fields}")
            
            # 创建工作表
            ws = wb.create_sheet(title=sheet_name)
            
            # 构建表头 - 按照第一行数据的字典顺序
            # 获取第一条数据的字段顺序
            first_row_fields = list(data[0].keys())
            # 其他字段按原顺序追加
            other_fields = [f for f in all_fields if f not in first_row_fields]
            headers = ["序号"] + first_row_fields + other_fields
            
            # 写入表头
            for col_idx, header in enumerate(headers, 1):
                cell = ws.cell(row=1, column=col_idx, value=header)
                cell.font = header_font
                cell.fill = header_fill
                cell.alignment = Alignment(horizontal='center', vertical='center')
                cell.border = thin_border
        
            # 写入数据
            for row_idx, item in enumerate(data, 2):
                # 序号列
                ws.cell(row=row_idx, column=1, value=row_idx-1)
                
                # 展平数据
                flat_item = flatten_dict(item)
                
                # 写入各字段数据
                for col_idx, field in enumerate(headers[1:], 2):  # 从第2列开始
                    value = flat_item.get(field, "")
                    
                    # 处理特殊数据类型
                    if isinstance(value, (list, dict)):
                        value = json.dumps(value, ensure_ascii=False, indent=2)
                    elif value is None:
                        value = ""
                    
                    cell = ws.cell(row=row_idx, column=col_idx, value=str(value))
                    cell.border = thin_border
                    cell.alignment = Alignment(vertical='top', wrap_text=True)
            
            # 自动调整列宽
            for col_idx, header in enumerate(headers, 1):
                col_letter = get_column_letter(col_idx)
                max_length = 0
                
                # 获取该列中最长内容的长度
                for row in ws.iter_rows(min_col=col_idx, max_col=col_idx):
                    for cell in row:
                        try:
                            if len(str(cell.value)) > max_length:
                                max_length = len(str(cell.value))
                        except:
                            pass
                
                # 设置列宽（限制最大宽度为50）
                adjusted_width = min(max_length + 2, 50)
                ws.column_dimensions[col_letter].width = adjusted_width
            
            total_records += len(data)
            total_sheets += 1
            print(f"成功创建工作表 '{sheet_name}'，包含 {len(data)} 条记录，{len(all_fields)} 个字段")
        
        # 保存文件
        wb.save(output_file)
        print(f"\n成功生成Excel文件: {output_file}")
        print(f"共创建 {total_sheets} 个工作表")
        print(f"共转换 {total_records} 条记录")
        
        return output_file
        
    except Exception as e:
        print(f"错误: {e}")
        raise



if __name__ == "__main__":
    save_format_xlsx([], "")
