import openpyxl
from typing import Dict, Any, Optional
import os


class ExcelWriter:
    def __init__(self, xlsx_path: str, template_path: str = 'E:\\template.xlsx', data_mapper=None):
        self.template_path = template_path
        self.ref_cache = {}  # 参考号缓存字典
        self._rpa_fill = data_mapper
        self.xlsx_path = xlsx_path

        # 检查文件是否存在，不存在则使用模板
        if not os.path.exists(xlsx_path):
            self.workbook = openpyxl.load_workbook(self.template_path)
        else:
            self.workbook = openpyxl.load_workbook(xlsx_path)
            
        # 初始化或更新参考号缓存
        self._init_ref_cache(self.workbook.active)
        self.sheet = self.workbook.active

    def _init_ref_cache(self, sheet):
        """初始化参考号缓存"""
        self.ref_cache.clear()
        for idx, cell in enumerate(sheet['B'][1:], start=2):  # B列是参考号列，从第2行开始
            if cell.value:
                self.ref_cache[cell.value] = idx
    
    def _find_row_by_ref(self, sheet, ref_no: str) -> Optional[int]:
        """查找参考号对应的行号"""
        # 优先从缓存查找
        if ref_no in self.ref_cache:
            return self.ref_cache[ref_no]
        return None
        
    def write_to_excel(self, task: Dict[str, Any]) -> bool:
        """
        将字典数据写入Excel文件
        :param data: 输入字典数据
        :param xlsx_path: xlsx文件路径
        :return: 是否成功
        """
        try:
            meta = task['meta']
            # 查找参考号对应的行
            ref_no = meta.get('ums_receipt_info', {}).get('参考号', '')
            row_idx = self._find_row_by_ref(self.sheet, ref_no) if ref_no else None
            
            # 准备数据映射
            row_data = self._rpa_fill._data_mapper(task)
            
            # 写入数据
            if row_idx is not None:
                # 更新现有行
                for col_idx, value in enumerate(row_data, 1):
                    self.sheet.cell(row=row_idx, column=col_idx, value=value)
            else:
                # 新增行
                self.sheet.append(list(row_data))
            return True
        except Exception as e:
            print(f"写入Excel失败: {str(e)}")
            return False

    def close(self):
        """关闭Excel文件"""
        self.workbook.save(self.xlsx_path)
        self.workbook.close()
        