"""
表格OCR处理模块
-------------
处理图片并提取表格内容，保存为Excel文件。
"""

import os
import time
import base64
from concurrent.futures import ThreadPoolExecutor
from typing import Dict, List, Optional, Tuple, Callable

from app.core.utils.log_utils import get_logger
from app.core.utils.file_utils import (
    ensure_dir, 
    get_file_extension, 
    get_files_by_extensions, 
    generate_timestamp_filename,
    is_file_size_valid,
    load_json,
    save_json
)
from app.config.settings import ConfigManager
from .baidu_ocr import BaiduOCRClient

logger = get_logger(__name__)

class ProcessedRecordManager:
    """处理记录管理器，用于跟踪已处理的文件"""
    
    def __init__(self, record_file: str):
        """
        初始化处理记录管理器
        
        Args:
            record_file: 记录文件路径
        """
        self.record_file = record_file
        self.processed_files = self._load_record()
    
    def _load_record(self) -> Dict[str, str]:
        """
        加载处理记录
        
        Returns:
            处理记录字典，键为输入文件路径，值为输出文件路径
        """
        return load_json(self.record_file, {})
    
    def save_record(self) -> None:
        """保存处理记录"""
        save_json(self.processed_files, self.record_file)
    
    def is_processed(self, image_file: str) -> bool:
        """
        检查图片是否已处理
        
        Args:
            image_file: 图片文件路径
            
        Returns:
            是否已处理
        """
        return image_file in self.processed_files
    
    def mark_as_processed(self, image_file: str, output_file: str) -> None:
        """
        标记图片为已处理
        
        Args:
            image_file: 图片文件路径
            output_file: 输出文件路径
        """
        self.processed_files[image_file] = output_file
        self.save_record()
    
    def get_output_file(self, image_file: str) -> Optional[str]:
        """
        获取图片的输出文件路径
        
        Args:
            image_file: 图片文件路径
            
        Returns:
            输出文件路径，如果不存在则返回None
        """
        return self.processed_files.get(image_file)
    
    def get_unprocessed_files(self, files: List[str]) -> List[str]:
        """
        获取未处理的文件列表
        
        Args:
            files: 文件列表
            
        Returns:
            未处理的文件列表
        """
        return [file for file in files if not self.is_processed(file)]

class OCRProcessor:
    """
    OCR处理器，负责协调OCR识别和结果处理
    """
    
    def __init__(self, config: Optional[ConfigManager] = None):
        """
        初始化OCR处理器
        
        Args:
            config: 配置管理器
        """
        self.config = config or ConfigManager()
        self.ocr_client = None
        self._ensure_ocr_client()
        
        # 修复ConfigParser对象没有get_path方法的问题
        try:
            # 获取输入和输出目录
            self.input_folder = config.get_path('Paths', 'input_folder', fallback='data/input', create=True) if hasattr(config, 'get_path') else os.path.abspath('data/input')
            self.output_folder = config.get_path('Paths', 'output_folder', fallback='data/output', create=True) if hasattr(config, 'get_path') else os.path.abspath('data/output')
            self.temp_folder = config.get_path('Paths', 'temp_folder', fallback='data/temp', create=True) if hasattr(config, 'get_path') else os.path.abspath('data/temp')
            
            # 确保目录存在
            os.makedirs(self.input_folder, exist_ok=True)
            os.makedirs(self.output_folder, exist_ok=True)
            os.makedirs(self.temp_folder, exist_ok=True)
            
            # 获取文件类型列表
            allowed_extensions_str = config.get('File', 'allowed_extensions', fallback='.jpg,.jpeg,.png,.bmp')
            self.file_types = [ext.strip() for ext in allowed_extensions_str.split(',') if ext.strip()]
            if not self.file_types:
                self.file_types = ['.jpg', '.jpeg', '.png', '.bmp', '.gif', '.tif', '.tiff']
            
            # 初始化OCR客户端
            self.ocr_client = BaiduOCRClient(self.config)
            
            # 记录实际路径
            logger.info(f"使用输入目录: {os.path.abspath(self.input_folder)}")
            logger.info(f"使用输出目录: {os.path.abspath(self.output_folder)}")
            logger.info(f"使用临时目录: {os.path.abspath(self.temp_folder)}")
            logger.info(f"允许的文件类型: {self.file_types}")
            
            # 初始化processed_files_json和record_manager
            self.processed_files_json = os.path.join(self.output_folder, 'processed_files.json')
            self.record_manager = ProcessedRecordManager(self.processed_files_json)
            
            # 加载已处理文件记录
            self.processed_files = self._load_processed_files()
            
            logger.info(f"初始化OCRProcessor完成：输入目录={self.input_folder}, 输出目录={self.output_folder}")
        except Exception as e:
            logger.error(f"初始化OCRProcessor失败: {e}")
            raise
    
    def _load_processed_files(self) -> Dict[str, str]:
        """
        加载已处理的文件记录
        
        Returns:
            已处理的文件记录字典，键为输入文件路径，值为输出文件路径
        """
        return load_json(self.processed_files_json, {})
    
    def get_unprocessed_images(self) -> List[str]:
        """
        获取未处理的图片列表
        
        Returns:
            未处理的图片文件路径列表
        """
        # 获取所有图片文件
        image_files = get_files_by_extensions(self.input_folder, self.file_types)
        
        # 如果需要跳过已存在的文件
        skip_existing = True
        try:
            skip_existing = self.config.getboolean('Performance', 'skip_existing', fallback=True)
        except Exception:
            pass

        if skip_existing:
            # 过滤已处理的文件
            unprocessed_files = self.record_manager.get_unprocessed_files(image_files)
            logger.info(f"找到 {len(image_files)} 个图片文件，其中 {len(unprocessed_files)} 个未处理")
            return unprocessed_files
        
        logger.info(f"找到 {len(image_files)} 个图片文件（不跳过已处理的文件）")
        return image_files
    
    def validate_image(self, image_path: str) -> bool:
        """
        验证图片是否有效
        
        Args:
            image_path: 图片文件路径
            
        Returns:
            图片是否有效
        """
        # 检查文件是否存在
        if not os.path.exists(image_path):
            logger.warning(f"图片文件不存在: {image_path}")
            return False
        
        # 检查文件扩展名
        ext = get_file_extension(image_path)
        if ext not in self.file_types:
            logger.warning(f"不支持的文件类型: {ext}, 文件: {image_path}")
            return False
        
        # 检查文件大小
        max_size_mb = 4.0
        try:
            max_size_mb = float(self.config.get('File', 'max_file_size_mb', fallback='4.0'))
        except Exception:
            pass
        
        if not is_file_size_valid(image_path, max_size_mb):
            logger.warning(f"文件大小超过限制 ({max_size_mb}MB): {image_path}")
            return False
        
        return True
    
    def process_image(self, image_path: str) -> Optional[str]:
        """
        处理单个图片
        
        Args:
            image_path: 图片文件路径
            
        Returns:
            输出Excel文件路径，如果处理失败则返回None
        """
        # 验证图片
        if not self.validate_image(image_path):
            return None
        
        # 获取是否跳过已处理文件的配置
        skip_existing = True
        try:
            skip_existing = self.config.getboolean('Performance', 'skip_existing', fallback=True)
        except Exception:
            pass
        
        # 如果需要跳过已处理的文件
        if skip_existing and self.record_manager.is_processed(image_path):
            output_file = self.record_manager.get_output_file(image_path)
            logger.info(f"图片已处理，跳过: {image_path}, 输出文件: {output_file}")
            return output_file
        
        logger.info(f"开始处理图片: {image_path}")
        
        try:
            # 获取Excel扩展名
            excel_extension = '.xlsx'
            try:
                excel_extension = self.config.get('File', 'excel_extension', fallback='.xlsx')
            except Exception:
                pass
            
            # 生成输出文件路径
            file_name = os.path.splitext(os.path.basename(image_path))[0]
            output_file = os.path.join(self.output_folder, f"{file_name}{excel_extension}")
            
            # 检查是否已存在对应的Excel文件
            if os.path.exists(output_file) and skip_existing:
                logger.info(f"已存在对应的Excel文件，跳过处理: {os.path.basename(image_path)} -> {os.path.basename(output_file)}")
                # 记录处理结果
                self.record_manager.mark_as_processed(image_path, output_file)
                return output_file
            
            # 进行OCR识别
            ocr_result = self.ocr_client.recognize_table(image_path)
            if not ocr_result:
                logger.error(f"OCR识别失败: {image_path}")
                return None
                
            # 保存Excel文件 - 按照v1版本逻辑提取Excel数据
            excel_base64 = None
            
            # 从不同可能的字段中尝试获取Excel数据
            if 'excel_file' in ocr_result:
                excel_base64 = ocr_result['excel_file']
                logger.debug("从excel_file字段获取Excel数据")
            elif 'result' in ocr_result:
                if 'result_data' in ocr_result['result']:
                    excel_base64 = ocr_result['result']['result_data']
                    logger.debug("从result.result_data字段获取Excel数据")
                elif 'excel_file' in ocr_result['result']:
                    excel_base64 = ocr_result['result']['excel_file']
                    logger.debug("从result.excel_file字段获取Excel数据")
                elif 'tables_result' in ocr_result['result'] and ocr_result['result']['tables_result']:
                    for table in ocr_result['result']['tables_result']:
                        if 'excel_file' in table:
                            excel_base64 = table['excel_file']
                            logger.debug("从tables_result中获取Excel数据")
                            break
                    
            # 如果还是没有找到Excel数据，尝试通过get_excel_result获取
            if not excel_base64:
                logger.info("无法从直接返回中获取Excel数据，尝试通过API获取...")
                excel_data = self.ocr_client.get_excel_result(ocr_result)
                if not excel_data:
                    logger.error(f"获取Excel结果失败: {image_path}")
                    return None
                    
                # 保存Excel文件
                os.makedirs(os.path.dirname(output_file), exist_ok=True)
                with open(output_file, 'wb') as f:
                    f.write(excel_data)
            else:
                # 解码并保存Excel文件
                try:
                    excel_data = base64.b64decode(excel_base64)
                    os.makedirs(os.path.dirname(output_file), exist_ok=True)
                    with open(output_file, 'wb') as f:
                        f.write(excel_data)
                except Exception as e:
                    logger.error(f"解码或保存Excel数据时出错: {e}")
                    return None
            
            logger.info(f"图片处理成功: {image_path}, 输出文件: {output_file}")
            
            # 标记为已处理
            self.record_manager.mark_as_processed(image_path, output_file)
            
            return output_file
            
        except Exception as e:
            logger.error(f"处理图片时出错: {image_path}, 错误: {e}")
            return None
    
    def process_images_batch(self, batch_size: int = None, max_workers: int = None, progress_cb: Optional[Callable[[int], None]] = None) -> Tuple[int, int]:
        """
        批量处理图片
        
        Args:
            batch_size: 批处理大小，如果为None则使用配置值
            max_workers: 最大线程数，如果为None则使用配置值
            
        Returns:
            (总处理数, 成功处理数)元组
        """
        # 使用配置值或参数值
        if batch_size is None:
            try:
                batch_size = self.config.getint('Performance', 'batch_size', fallback=5)
            except Exception:
                batch_size = 5

        if max_workers is None:
            try:
                # 强制设为 1 以严格遵守百度 QPS=2 的限制（一张图调两个接口就满了）
                max_workers = self.config.getint('Performance', 'max_workers', fallback=1)
            except Exception:
                max_workers = 1
    
        # 获取未处理的图片
        unprocessed_images = self.get_unprocessed_images()
        if not unprocessed_images:
            logger.warning("没有需要处理的图片")
            return 0, 0
    
        total = len(unprocessed_images)
        success_count = 0
    
        # 按批次处理
        for i in range(0, total, batch_size):
            batch = unprocessed_images[i:i+batch_size]
            logger.info(f"处理批次 {i//batch_size+1}/{(total+batch_size-1)//batch_size}: {len(batch)} 个文件")
            try:
                if progress_cb:
                    # 以批次为单位估算进度（0-90%），保留10%给后续阶段
                    percent = int(10 + (i / max(total, 1)) * 80)
                    progress_cb(min(percent, 90))
            except Exception:
                pass
        
            # 使用多线程处理批次
            with ThreadPoolExecutor(max_workers=max_workers) as executor:
                results = list(executor.map(self.process_image, batch))
            
                # 统计成功数
                success_count += sum(1 for result in results if result is not None)
    
        logger.info(f"所有图片处理完成, 总计: {total}, 成功: {success_count}")
        try:
            if progress_cb:
                progress_cb(90)
        except Exception:
            pass
        return total, success_count 
