# ============================================================
# OCR + Email 合併工具（Gemini OCR → Excel → Resend 寄信）
# ============================================================
# 安裝依賴：pip install resend google-genai openpyxl ipywidgets
# ============================================================

import os
import re
import base64
import resend
import openpyxl
from openpyxl.styles import Font, Alignment, PatternFill
from google import genai
from google.genai import types
import ipywidgets as widgets
from IPython.display import display, clear_output, HTML

# ─────────────────────────────────────────
# UI 元件：API 金鑰 & 設定
# ─────────────────────────────────────────
style = {'description_width': '140px'}
layout = widgets.Layout(width='600px')

gemini_key_input = widgets.Password(
    description='Gemini API Key:',
    placeholder='貼上你的 Gemini API Key',
    style=style, layout=layout
)
resend_key_input = widgets.Password(
    description='Resend API Key:',
    placeholder='貼上你的 Resend API Key',
    style=style, layout=layout
)
sender_input = widgets.Text(
    description='寄件人:',
    value='Acme <onboarding@resend.dev>',
    style=style, layout=layout
)
recipient_input = widgets.Text(
    description='收件人 Email:',
    placeholder='example@gmail.com',
    style=style, layout=layout
)
subject_input = widgets.Text(
    description='郵件主旨:',
    value='OCR 辨識結果',
    style=style, layout=layout
)
prompt_input = widgets.Text(
    description='OCR Prompt:',
    value='幫我進行OCR，提取圖片中的所有文字，保留原始排版結構',
    style=style, layout=layout
)
upload = widgets.FileUpload(accept='image/*', multiple=True)
button = widgets.Button(
    description='🚀 開始辨識並寄信',
    button_style='primary',
    layout=widgets.Layout(width='220px')
)
output = widgets.Output()

# ─────────────────────────────────────────
# 工具函式
# ─────────────────────────────────────────
def sanitize_sheet_name(name: str, index: int) -> str:
    name = re.sub(r'[\\/*?:\[\]]', '_', name)
    return name[:31] if len(name) <= 31 else f"圖片_{index+1}"


def build_excel(files: list, prompt: str, client) -> tuple[bytes, str]:
    """執行 OCR 並建立 Excel，回傳 (bytes, 摘要文字)"""
    output_path = 'ocr_results.xlsx'
    wb = openpyxl.Workbook()
    wb.remove(wb.active)

    header_font  = Font(name='Arial', bold=True, size=12, color='FFFFFF')
    header_fill  = PatternFill('solid', start_color='2F5496')
    header_align = Alignment(horizontal='center', vertical='center')
    body_font    = Font(name='Arial', size=11)
    label_font   = Font(name='Arial', bold=True, size=11, color='2F5496')
    meta_fill    = PatternFill('solid', start_color='D9E1F2')

    summary_lines = []

    for idx, (filename, file_info) in enumerate(files):
        print(f'[{idx+1}/{len(files)}] 辨識：{filename} ...')
        img_bytes = bytes(file_info['content'])
        mime_type = file_info['metadata']['type']

        try:
            response = client.models.generate_content(
                model='gemini-2.5-flash',
                contents=[
                    types.Part.from_bytes(data=img_bytes, mime_type=mime_type),
                    prompt
                ]
            )
            ocr_text = response.text
        except Exception as e:
            ocr_text = f'❌ 辨識失敗：{e}'

        sheet_name = sanitize_sheet_name(filename, idx)
        ws = wb.create_sheet(title=sheet_name)
        ws.column_dimensions['A'].width = 10
        ws.column_dimensions['B'].width = 100

        ws['A1'] = '行號'
        ws['B1'] = '文字內容'
        for cell in [ws['A1'], ws['B1']]:
            cell.font = header_font
            cell.fill = header_fill
            cell.alignment = header_align
        ws.row_dimensions[1].height = 22

        for label, value in [('檔案名稱', filename), ('MIME 類型', mime_type)]:
            ws.append(['', f'【{label}】{value}'])
            row_idx = ws.max_row
            cell_b = ws.cell(row=row_idx, column=2)
            cell_b.font = label_font
            cell_b.fill = meta_fill
            cell_b.alignment = Alignment(vertical='center')
            ws.row_dimensions[row_idx].height = 18

        ws.append(['', ''])

        lines = ocr_text.splitlines()
        for line_no, line in enumerate(lines, start=1):
            ws.append([line_no, line])
            row_idx = ws.max_row
            ws.cell(row=row_idx, column=1).font = Font(name='Arial', size=10, color='888888')
            ws.cell(row=row_idx, column=1).alignment = Alignment(horizontal='center', vertical='center')
            ws.cell(row=row_idx, column=2).font = body_font
            ws.cell(row=row_idx, column=2).alignment = Alignment(vertical='center')
            ws.row_dimensions[row_idx].height = 18

        summary_lines.append(f'• {filename}：共 {len(lines)} 行')
        print(f'   ✅ 完成，共 {len(lines)} 行')

    wb.save(output_path)
    with open(output_path, 'rb') as f:
        excel_bytes = f.read()

    return excel_bytes, '\n'.join(summary_lines)


def send_email_with_attachment(
    api_key: str, sender: str, recipient: str,
    subject: str, excel_bytes: bytes, summary: str
):
    """透過 Resend 寄送含 Excel 附件的郵件"""
    resend.api_key = api_key

    b64_content = base64.b64encode(excel_bytes).decode()

    html_body = f"""
    <h2>📊 OCR 辨識結果</h2>
    <p>您好，以下是本次 OCR 辨識摘要：</p>
    <pre style="background:#f4f4f4;padding:12px;border-radius:6px;">{summary}</pre>
    <p>詳細結果請見附件 <strong>ocr_results.xlsx</strong>。</p>
    <hr>
    <small style="color:#888;">由 Gemini OCR 工具自動產生</small>
    """

    params: resend.Emails.SendParams = {
        'from': sender,
        'to': [recipient],
        'subject': subject,
        'html': html_body,
        'attachments': [
            {
                'filename': 'ocr_results.xlsx',
                'content': b64_content,
            }
        ],
    }

    result = resend.Emails.send(params)
    return result


# ─────────────────────────────────────────
# 按鈕事件
# ─────────────────────────────────────────
def on_click(b):
    with output:
        clear_output()

        # 驗證輸入
        if not gemini_key_input.value.strip():
            print('⚠️ 請填入 Gemini API Key！')
            return
        if not resend_key_input.value.strip():
            print('⚠️ 請填入 Resend API Key！')
            return
        if not recipient_input.value.strip():
            print('⚠️ 請填入收件人 Email！')
            return
        if not upload.value:
            print('⚠️ 請先上傳至少一張圖片！')
            return

        # 初始化 Gemini client
        client = genai.Client(api_key=gemini_key_input.value.strip())
        files = list(upload.value.items())
        print(f'共 {len(files)} 張圖片，開始辨識...\n')

        # OCR + 建立 Excel
        try:
            excel_bytes, summary = build_excel(files, prompt_input.value, client)
        except Exception as e:
            print(f'❌ OCR 或 Excel 建立失敗：{e}')
            return

        # 提供下載連結
        b64 = base64.b64encode(excel_bytes).decode()
        link_html = (
            f'<a href="data:application/vnd.openxmlformats-officedocument'
            f'.spreadsheetml.sheet;base64,{b64}" download="ocr_results.xlsx">'
            f'📥 點此下載 ocr_results.xlsx</a>'
        )
        display(HTML(link_html))

        # 寄送郵件
        print(f'\n📧 正在寄送至 {recipient_input.value.strip()} ...')
        try:
            result = send_email_with_attachment(
                api_key=resend_key_input.value.strip(),
                sender=sender_input.value.strip(),
                recipient=recipient_input.value.strip(),
                subject=subject_input.value.strip(),
                excel_bytes=excel_bytes,
                summary=summary
            )
            print(f'✅ 郵件寄送成功！ID：{result.get("id", result)}')
        except Exception as e:
            print(f'❌ 郵件寄送失敗：{e}')


button.on_click(on_click)

# ─────────────────────────────────────────
# 顯示介面
# ─────────────────────────────────────────
display(
    widgets.HTML('<h3>🔑 API 設定</h3>'),
    gemini_key_input,
    resend_key_input,
    widgets.HTML('<h3>📧 郵件設定</h3>'),
    sender_input,
    recipient_input,
    subject_input,
    widgets.HTML('<h3>🖼️ 圖片辨識設定</h3>'),
    prompt_input,
    upload,
    button,
    output
)