import gradio as gr
import pandas as pd
from PIL import Image
import easyocr
import numpy as np
import re
import traceback

class FinalOCRProcessor:
    def __init__(self):
        self.reader = None
        self.lang_config = "未初始化"
        self.initialize_reader()
    
    def initialize_reader(self):
        """安全初始化EasyOCR"""
        print("正在初始化OCR引擎...")
        
        # 嘗試不同的語言配置
        configs = [
            (['ch_tra', 'en'], "繁體中文+英文"),
            (['ch_sim', 'en'], "簡體中文+英文"), 
            (['en'], "僅英文")
        ]
        
        for lang_list, description in configs:
            try:
                print(f"嘗試配置：{description}")
                self.reader = easyocr.Reader(lang_list, gpu=False)
                self.lang_config = description
                print(f"✓ 成功初始化：{description}")
                return
            except Exception as e:
                print(f"✗ 配置 {description} 失敗：{str(e)}")
                continue
        
        print("❌ 所有配置都失敗了")
        self.reader = None
        self.lang_config = "初始化失敗"
    
    def process_image(self, image):
        """處理圖片的主要函數"""
        if self.reader is None:
            return self.create_error_result("OCR引擎未正確初始化")
        
        if image is None:
            return self.create_error_result("請上傳圖片")
        
        try:
            # 轉換圖片格式
            if isinstance(image, np.ndarray):
                img_array = image
            else:
                img_array = np.array(image)
            
            print("開始OCR識別...")
            
            # 執行OCR
            results = self.reader.readtext(img_array)
            
            print(f"識別完成，找到 {len(results)} 個文字區塊")
            
            if not results:
                return self.create_empty_result("圖片中未找到任何文字")
            
            # 處理OCR結果
            data_list = []
            for i, (bbox, text, confidence) in enumerate(results):
                # 計算位置信息
                x_coords = [point[0] for point in bbox]
                y_coords = [point[1] for point in bbox]
                
                left = int(min(x_coords))
                right = int(max(x_coords))
                top = int(min(y_coords))
                bottom = int(max(y_coords))
                center_x = int((left + right) / 2)
                center_y = int((top + bottom) / 2)
                
                data_list.append({
                    '序號': i + 1,
                    '識別文字': text.strip(),
                    '信心度': round(float(confidence), 3),
                    '中心X': center_x,
                    '中心Y': center_y,
                    '左': left,
                    '上': top,
                    '右': right,
                    '下': bottom,
                    '寬度': right - left,
                    '高度': bottom - top,
                    '類型': self.get_text_type(text.strip())
                })
            
            # 創建DataFrame並排序
            df = pd.DataFrame(data_list)
            df = df.sort_values(['中心Y', '中心X']).reset_index(drop=True)
            
            # 重新編號
            df['序號'] = range(1, len(df) + 1)
            
            # 生成摘要
            summary = self.generate_summary(df)
            
            # 嘗試重建表格
            table_df = self.reconstruct_table(df)
            
            return df, table_df, summary
            
        except Exception as e:
            error_msg = f"處理圖片時出錯：{str(e)}"
            print(error_msg)
            print(traceback.format_exc())
            return self.create_error_result(error_msg)
    
    def get_text_type(self, text):
        """判斷文字類型"""
        if not text:
            return "空白"
        
        # 日期格式
        if re.search(r'\d+月\d+日', text):
            return "日期"
        
        # 時間格式
        if re.search(r'\d{1,2}[-:]\d{1,2}', text):
            return "時間"
        
        # 純數字
        if re.match(r'^\d+$', text):
            return "數字"
        
        # 中文姓名或詞語
        if re.match(r'^[\u4e00-\u9fff]{1,6}$', text):
            return "中文"
        
        # 包含中文的混合內容
        if re.search(r'[\u4e00-\u9fff]', text):
            return "中文混合"
        
        # 英文
        if re.match(r'^[a-zA-Z\s]+$', text):
            return "英文"
        
        return "其他"
    
    def generate_summary(self, df):
        """生成識別摘要"""
        if df.empty:
            return "沒有識別到任何內容"
        
        total = len(df)
        avg_conf = df['信心度'].mean()
        high_conf = (df['信心度'] >= 0.8).sum()
        
        summary = f"""🔍 OCR識別報告
{'='*40}
📊 基本統計
• 引擎配置：{self.lang_config}
• 識別區塊：{total} 個
• 平均信心度：{avg_conf:.3f}
• 高信心度區塊：{high_conf} 個 ({high_conf/total*100:.1f}%)

📝 文字類型統計
"""
        
        # 統計各類型數量
        type_counts = df['類型'].value_counts()
        for text_type, count in type_counts.items():
            percentage = count / total * 100
            summary += f"• {text_type}：{count} 個 ({percentage:.1f}%)\n"
        
        # 品質評估
        if avg_conf >= 0.8:
            quality = "優秀 ✅"
        elif avg_conf >= 0.6:
            quality = "良好 ⚠️"
        else:
            quality = "待改進 ❌"
        
        summary += f"\n🎯 整體品質：{quality}"
        
        return summary
    
    def reconstruct_table(self, df):
        """重建表格結構"""
        if df.empty:
            return pd.DataFrame()
        
        try:
            # 按Y座標分組形成行
            rows = []
            sorted_df = df.sort_values('中心Y')
            
            current_row = []
            last_y = None
            y_threshold = 25  # Y座標閾值
            
            for _, item in sorted_df.iterrows():
                if last_y is None or abs(item['中心Y'] - last_y) <= y_threshold:
                    current_row.append(item)
                else:
                    if current_row:
                        # 按X座標排序
                        current_row.sort(key=lambda x: x['中心X'])
                        rows.append(current_row)
                    current_row = [item]
                last_y = item['中心Y']
            
            # 添加最後一行
            if current_row:
                current_row.sort(key=lambda x: x['中心X'])
                rows.append(current_row)
            
            if not rows:
                return pd.DataFrame()
            
            # 構建表格數據
            table_data = []
            max_cols = max(len(row) for row in rows)
            
            for row_idx, row_items in enumerate(rows):
                row_dict = {'行號': row_idx + 1}
                
                for col_idx, item in enumerate(row_items):
                    col_name = f'第{col_idx + 1}列'
                    row_dict[col_name] = item['識別文字']
                
                # 填充空列
                for col_idx in range(len(row_items), max_cols):
                    col_name = f'第{col_idx + 1}列'
                    row_dict[col_name] = ""
                
                table_data.append(row_dict)
            
            return pd.DataFrame(table_data)
            
        except Exception as e:
            print(f"重建表格失敗：{e}")
            return pd.DataFrame([{"錯誤": "表格重建失敗"}])
    
    def create_error_result(self, message):
        """創建錯誤結果"""
        error_df = pd.DataFrame([{"錯誤信息": message}])
        return error_df, error_df, f"❌ {message}"
    
    def create_empty_result(self, message):
        """創建空結果"""
        empty_df = pd.DataFrame()
        return empty_df, empty_df, f"ℹ️ {message}"

def main():
    """主函數，創建Gradio應用"""
    
    # 初始化處理器
    processor = FinalOCRProcessor()
    
    def process_image_wrapper(image, min_confidence):
        """包裝處理函數"""
        try:
            # 處理圖片
            df, table_df, summary = processor.process_image(image)
            
            # 應用信心度過濾
            if not df.empty and '信心度' in df.columns:
                original_count = len(df)
                df = df[df['信心度'] >= min_confidence].reset_index(drop=True)
                filtered_count = len(df)
                
                if filtered_count < original_count:
                    summary += f"\n\n🔍 信心度過濾：保留 {filtered_count}/{original_count} 個結果"
            
            return summary, df, table_df
            
        except Exception as e:
            error_msg = f"處理失敗：{str(e)}"
            error_df = pd.DataFrame([{"錯誤": error_msg}])
            return error_msg, error_df, error_df
    
    # 創建Gradio界面
    with gr.Blocks(
        title="中文OCR識別系統",
        theme=gr.themes.Default()
    ) as app:
        
        # 標題和說明
        gr.HTML("""
        <div style="text-align: center; padding: 20px; background: #f0f2f6; border-radius: 10px; margin-bottom: 20px;">
            <h1 style="color: #1f2937; margin-bottom: 10px;">🔍 中文OCR識別系統</h1>
            <p style="color: #6b7280; font-size: 16px;">
                支持中文、英文、數字識別，自動重建表格結構
            </p>
        </div>
        """)
        
        # 主要操作區域
        with gr.Row():
            # 左側：圖片上傳和控制
            with gr.Column(scale=1):
                image_input = gr.Image(
                    label="📤 上傳圖片",
                    type="pil"
                )
                
                confidence_slider = gr.Slider(
                    minimum=0.0,
                    maximum=1.0,
                    value=0.3,
                    step=0.1,
                    label="🎯 最低信心度閾值"
                )
                
                process_button = gr.Button(
                    "🚀 開始識別",
                    variant="primary"
                )
            
            # 右側：摘要報告
            with gr.Column(scale=1):
                summary_output = gr.Textbox(
                    label="📊 識別報告",
                    lines=15,
                    max_lines=20
                )
        
        # 結果展示區域
        with gr.Row():
            with gr.Column():
                gr.Markdown("### 📋 詳細識別結果")
                detail_output = gr.Dataframe(
                    label="所有識別的文字內容",
                    interactive=True
                )
            
            with gr.Column():
                gr.Markdown("### 🔄 重建表格")
                table_output = gr.Dataframe(
                    label="按表格結構重新組織的數據",
                    interactive=True
                )
        
        # 使用說明
        with gr.Accordion("📖 使用說明", open=False):
            gr.Markdown("""
            ### 🚀 快速開始
            1. **上傳圖片**：點擊上傳包含文字的圖片
            2. **調整參數**：設置最低信心度（建議0.3-0.7）
            3. **開始識別**：點擊按鈕或直接上傳圖片自動處理
            4. **查看結果**：在兩個表格中查看識別結果
            
            ### 📊 結果說明
            - **詳細識別結果**：每個文字區塊的完整信息
            - **重建表格**：嘗試還原原始表格的行列結構
            - **識別報告**：統計信息和品質評估
            
            ### 💡 優化建議
            - 使用清晰、高對比度的圖片
            - 確保文字大小適中（不要太小）
            - 避免圖片傾斜或變形
            - 調整信心度閾值獲得最佳結果
            """)
        
        # 綁定事件
        process_button.click(
            fn=process_image_wrapper,
            inputs=[image_input, confidence_slider],
            outputs=[summary_output, detail_output, table_output]
        )
        
        # 圖片上傳時自動處理
        image_input.change(
            fn=process_image_wrapper,
            inputs=[image_input, confidence_slider],
            outputs=[summary_output, detail_output, table_output]
        )
    
    return app

if __name__ == "__main__":
    app = main()
    app.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=True
    )