import os
import time
import re
import io
from urllib.parse import urljoin
import pandas as pd
import openpyxl
from openpyxl.styles import Font, Alignment, Border, Side, PatternFill
from bs4 import BeautifulSoup
from playwright.sync_api import sync_playwright
import gradio as gr
from PIL import Image

import torch
import onnx
import onnxruntime as rt
from torchvision import transforms as T
from pathlib import Path
from huggingface_hub import hf_hub_download

from utils.tokenizer_base import Tokenizer

# ==========================================
# 1. 初始化 OCR 模型 (從 Hugging Face 雲端載入)
# ==========================================
cwd = Path(__file__).parent.resolve()
model_file = os.path.join(cwd, hf_hub_download("toandev/OCR-for-Captcha", "model.onnx"))

img_size = (32, 128)
vocab = r"0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ!\"#$%&'()*+,-./:;<=>?@[\]^_`{|}~"
tokenizer = Tokenizer(vocab)

def to_numpy(tensor):
    return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()

def get_transform(img_size):
    transforms = [
        T.Resize(img_size, T.InterpolationMode.BICUBIC),
        T.ToTensor(),
        T.Normalize(0.5, 0.5),
    ]
    return T.Compose(transforms)

transform, s = get_transform(img_size), rt.InferenceSession(model_file)

def predict_captcha(img: Image.Image):
    x = transform(img.convert("RGB")).unsqueeze(0)
    ort_inputs = {s.get_inputs()[0].name: to_numpy(x)}
    logits = s.run(None, ort_inputs)[0]
    probs = torch.tensor(logits).softmax(-1)
    preds, _ = tokenizer.decode(probs)
    return preds[0]

# ==========================================
# 2. 資料擷取與整理邏輯 (正式最終版)
# ==========================================
def extract_detail_data(html_content):
    soup = BeautifulSoup(html_content, 'html.parser')
    
    def get_value(label_name):
        # 1. 找標題
        tags = soup.find_all(lambda tag: tag.name in ['label', 'p'] and label_name in tag.text)
        for tag in tags:
            container = tag.find_parent('div', class_=re.compile('MuiGrid-container'))
            if container:
                # 2. 根據特徵，真正的資料都放在 1.125em 的 label 裡
                val_label = container.find('label', style=lambda s: s and '1.125em' in s)
                if val_label:
                    val = val_label.text.strip()
                    # 3. 防呆：過濾掉網頁還沒載入完的 '--' 或 '/' 預設值
                    if val and val not in ['--', '/', '-']:
                        return val
                        
                # 備用：如果沒有大字體，抓旁邊的資料區塊
                val_div = container.find('div', class_=re.compile('MuiGrid-grid-(sm|md)-9'))
                if val_div:
                    val = val_div.text.replace(tag.text, '').replace('：', '').replace(':', '').strip()
                    if val and val not in ['--', '/', '-']:
                        return val
        return ""

    lic_tag = soup.find('h6')
    lic_no = lic_tag.text.strip() if lic_tag else ""
    
    # 💡 將許可證字號的「字第」拆成兩行
    lic_no = lic_no.replace('字第', '字\n第')
    
    # 處理雙日期：以斜線分割，並取第 [0] 個 (即前者日期)
    raw_date = get_value("原始發證日期")
    issue_date = raw_date.split('/')[0].strip() if raw_date else ""

    return {
        "許可證字號": lic_no,
        "發證日期": issue_date,
        "中文品名": get_value("中文品名"),
        "英文品名": get_value("英文品名"),
        "成分": get_value("主成分略述"),
        "申請商": get_value("藥商名稱"),
        "製造廠": get_value("製造廠名稱"),
        "適應症": get_value("適應症")
    }

def format_english_title_case(text):
    if not text: return ""
    words = re.split(r'(\s+|,)', text)
    return "".join([w.capitalize() if w.isalpha() else w for w in words])

# ==========================================
# 3. Excel 產出專用函式
# ==========================================
def generate_excel(results, factory_name, start_year, end_year):
    df = pd.DataFrame(results)
    if df.empty:
        return None, "沒有資料可以產出 Excel。"

    df['英文品名'] = df['英文品名'].apply(format_english_title_case)
    df['成分'] = df['成分'].apply(format_english_title_case)
    
    columns_order = ['許可證字號', '發證日期', '中文品名', '英文品名', '成分', '申請商', '製造廠', '適應症']
    df = df[columns_order]
    df.insert(0, 'NO', range(1, len(df) + 1))
    
    base_title = f"{factory_name}_{start_year}至{end_year}年領證品項"
    output_filename = f"{base_title}.xlsx"
    df.to_excel(output_filename, index=False)
    
    wb = openpyxl.load_workbook(output_filename)
    ws = wb.active
    
    # 第一列大標題
    ws.insert_rows(1)
    ws.merge_cells('A1:I1')
    title_cell = ws['A1']
    title_cell.value = base_title
    title_cell.font = Font(name='微軟正黑體', size=14, bold=True)
    title_cell.alignment = Alignment(horizontal='center', vertical='center')
    title_cell.fill = PatternFill(start_color="FFCC00", end_color="FFCC00", fill_type="solid")
    ws.row_dimensions[1].height = 30
    
    font_style = Font(name='微軟正黑體', size=11)
    header_font = Font(name='微軟正黑體', size=11, bold=True)
    border_style = Border(left=Side(style='thin'), right=Side(style='thin'),
                          top=Side(style='thin'), bottom=Side(style='thin'))
    
    alignment_style = Alignment(horizontal='left', vertical='center', wrap_text=True)
    header_fill = PatternFill(start_color="FFFFCC", end_color="FFFFCC", fill_type="solid")
    
    ws.row_dimensions[2].height = 25

    for row_idx, row in enumerate(ws.iter_rows(min_row=1, max_row=ws.max_row), start=1):
        for cell in row:
            if row_idx == 1:
                cell.border = border_style
            elif row_idx == 2:
                cell.font = header_font
                cell.border = border_style
                cell.alignment = alignment_style
                cell.fill = header_fill
            else:
                cell.font = font_style
                cell.border = border_style
                cell.alignment = alignment_style

    # E 欄 (成分) 寬度為 30
    column_widths = {'A': 5, 'B': 16, 'C': 12, 'D': 25, 'E': 30, 'F': 40, 'G': 30, 'H': 30, 'I': 30}
    for col, width in column_widths.items():
        ws.column_dimensions[col].width = width

    ws.page_setup.orientation = ws.ORIENTATION_LANDSCAPE
    ws.page_setup.fitToWidth = 1
    ws.page_setup.fitToHeight = 0
    ws.sheet_properties.pageSetUpPr.fitToPage = True

    wb.save(output_filename)
    return output_filename, "🎉 Excel 檔案產出成功！"

# ==========================================
# 3.5 假資料測試模式 (一秒產出 Excel)
# ==========================================
def test_excel_layout():
    fake_results = [{
        "許可證字號": "衛部藥製字\n第123456號", "發證日期": "112-05-20",
        "中文品名": "測試用頭痛藥", "英文品名": "TEST HEADACHE PILL",
        "成分": "ACETAMINOPHEN", "申請商": "健亞生物科技股份有限公司",
        "製造廠": "健亞生物科技股份有限公司新竹廠", "適應症": "退燒、止痛"
    }, {
        "許可證字號": "衛部藥輸字\n第654321號", "發證日期": "113-01-10",
        "中文品名": "很長很長很長的一串測試品名", "英文品名": "VERY LONG ENGLISH NAME",
        "成分": "VITAMIN C", "申請商": "另一家生技股份有限公司",
        "製造廠": "台灣不知道哪裡的製造廠區", "適應症": "這是一個非常長的適應症測試確保換行正常"
    }]
    
    file_path, log_msg = generate_excel(fake_results, "排版測試廠", 110, 115)
    return file_path, f"✅ 測試模式執行完畢 (未啟動爬蟲)！\n{log_msg}"


# ==========================================
# 4. 爬蟲主程式
# ==========================================
def scrape_fda(factory_name, start_year, end_year, start_record=1):
    os.system("playwright install chromium")
    
    results = []
    debug_logs = []
    
    def log(msg):
        print(msg)
        debug_logs.append(msg)
    
    log(f"=== 🕵️ 開始執行西藥許可證抓取 ===")
    log(f"條件：製造廠【{factory_name}】, 區間【{start_year}~{end_year}年】")
    
    with sync_playwright() as p:
        browser = p.chromium.launch(headless=True)
        page = browser.new_page()
        
        alert_messages = []
        page.on("dialog", lambda dialog: (alert_messages.append(dialog.message), dialog.accept()))
        
        log("🌐 前往 TFDA 查詢網站...")
        page.goto("https://lmspiq.fda.gov.tw/web/DRPIQ/license-search", wait_until="networkidle")
        page.locator('input[name="factoryName"]').fill(factory_name)
        
        log("📸 正在破解驗證碼...")
        captcha_locator = page.locator('img[alt="驗證碼"]')
        captcha_bytes = captcha_locator.screenshot()
        captcha_text = predict_captcha(Image.open(io.BytesIO(captcha_bytes)))
        log(f"🔍 辨識驗證碼為：【{captcha_text}】")
        
        page.locator('input[name="code"]').fill(captcha_text)
        page.locator('button.MuiButton-containedPrimary:has-text("查詢")').click()
        
        time.sleep(5)
        if alert_messages:
            log(f"🚨 錯誤！網頁跳出警告：【{alert_messages[0]}】 (驗證碼可能猜錯，請重新查詢)")
            browser.close()
            return None, "\n".join(debug_logs)
            
        # 💡【全新加入：加速引擎】自動將每頁顯示數量切換為 100 筆
        log("⚡ 嘗試將每頁顯示數量切換為 100 筆以加速收集...")
        try:
            # 確保列表已初步載入
            page.wait_for_selector('a[href*="DRPIQ1000Result"]', timeout=30000)
            
            # 尋找下拉選單 (combobox) 並展開
            combobox = page.locator('div[role="combobox"]')
            if combobox.count() > 0:
                combobox.first.click()
                time.sleep(1) # 等待選單動畫展開
                
                # 點選 100 筆的選項
                page.locator('li[role="option"]:has-text("100")').click()
                log("  ✅ 成功切換為每頁 100 筆！等待畫面更新...")
                
                # 給伺服器一點時間重新吐出 100 筆資料
                time.sleep(4) 
                page.wait_for_selector('a[href*="DRPIQ1000Result"]', timeout=30000)
            else:
                log("  ⚠️ 找不到下拉選單，維持預設數量。")
        except Exception as e:
            log(f"  ⚠️ 切換數量失敗，將以預設數量繼續 ({str(e)})")

        log("⏳ 正在收集所有分頁的資料連結...")
        all_hrefs = []
        page_num = 1
        start_idx = int(start_record) - 1
        
        while True:
            try:
                page.wait_for_selector('a[href*="DRPIQ1000Result"]', timeout=30000)
            except:
                log("⚠️ 等不到列表資料，可能查無此廠資料或網頁過慢。")
                break
                
            time.sleep(1)
            current_hrefs = page.locator('a[href*="DRPIQ1000Result"]').evaluate_all(
                "elements => elements.map(e => e.getAttribute('href'))"
            )
            all_hrefs.extend(current_hrefs)
            
            if len(all_hrefs) <= start_idx:
                log(f"  ⏩ 尋找目標中... 快速跳過第 {page_num} 頁 (目前進度: {len(all_hrefs)} / 目標: {int(start_record)} 筆)")
            else:
                log(f"  => 收集列表資料中... 第 {page_num} 頁，目前累計 {len(all_hrefs)} 筆")
            
            next_btn = page.locator('button[aria-label="Go to next page"]')
            if next_btn.count() == 0 or next_btn.is_disabled():
                if len(all_hrefs) > start_idx:
                    log("  => 已到達最後一頁，列表收集完畢。")
                else:
                    log(f"  ⚠️ 警告：網站總資料只有 {len(all_hrefs)} 筆，不到您設定的第 {int(start_record)} 筆。")
                break
                
            next_btn.click()
            page_num += 1
            time.sleep(3)
            
        if not all_hrefs:
            browser.close()
            return None, "\n".join(debug_logs)
            
        if start_idx > 0:
            log(f"🚀 準備就緒！跳過前 {start_idx} 筆，直接從第 {int(start_record)} 筆開始抓取詳細資料...")
            hrefs_to_process = all_hrefs[start_idx:]
        else:
            hrefs_to_process = all_hrefs

        base_url = "https://lmspiq.fda.gov.tw/web/DRPIQ/license-search"
        
        for i, href in enumerate(hrefs_to_process):
            real_num = start_idx + i + 1
            full_url = urljoin(base_url, href)
            log(f"  🔗 [{real_num}/{len(all_hrefs)}] 嘗試進入: {full_url}")
            
            try:
                page.goto(full_url, timeout=15000, wait_until="domcontentloaded")
                page.wait_for_function("""
                    () => {
                        const labels = Array.from(document.querySelectorAll('label[style*="1.125em"]'));
                        return labels.some(l => l.innerText.trim().length > 2 && l.innerText.trim() !== '--');
                    }
                """, timeout=15000)
                time.sleep(1.5)
                
            except Exception as e:
                log(f"  👉 第 {real_num} 筆：網頁讀取超時！({str(e)})")
                error_img_path = "timeout_error.png"
                page.screenshot(path=error_img_path, full_page=True)
                log("📸 發生超時！程式已中斷。請從下方的「下載 Excel」按鈕處下載錯誤截圖查看問題。")
                browser.close()
                return error_img_path, "\n".join(debug_logs)
                
            detail_data = extract_detail_data(page.content())
            date_str = detail_data.get("發證日期", "")
            lic_no = detail_data.get("許可證字號", "未知")
            
            year_match = re.match(r'^(\d+)-', date_str)
            if year_match:
                year = int(year_match.group(1))
                if start_year <= year <= end_year:
                    results.append(detail_data)
                    log(f"  ✅ [{lic_no.replace(chr(10), '')}] 年份 {year}，已收錄！")
                else:
                    log(f"  ⏭️ [{lic_no.replace(chr(10), '')}] 年份 {year} 不在區間，跳過。")
            else:
                log(f"  ❌ 無法解析發證日期 [{date_str}]")
            
        browser.close()
    
    log(f"=== 🏁 執行完畢：成功收錄 {len(results)} 筆符合年份的資料 ===")
    
    if not results:
        return None, "\n".join(debug_logs)
    
    file_path, excel_log = generate_excel(results, factory_name, start_year, end_year)
    
    return file_path, "\n".join(debug_logs) + "\n" + excel_log

# ==========================================
# 5. Gradio 網頁介面
# ==========================================
with gr.Blocks(title="製造廠年度領證品項查詢") as demo:
    gr.Markdown("## 🏥 製造廠年度領證品項查詢")
    gr.Markdown("輸入製造廠名稱與年份區間，系統將自動破解驗證碼、抓取資料並匯出格式化 Excel。")
    
    with gr.Row():
        factory_input = gr.Textbox(label="製造廠名稱", placeholder="例如：健亞")
        start_year_input = gr.Number(label="起始年 (民國)", value=110, precision=0)
        end_year_input = gr.Number(label="結束年 (民國)", value=115, precision=0)
        start_record_input = gr.Number(label="起始掃描筆數", value=1, precision=0, info="輸入 90 表示從第 90 筆開始")
        
    with gr.Row():
        submit_btn = gr.Button("開始查詢與匯出", variant="primary")
        test_btn = gr.Button("🚀 測試 Excel 排版 (1秒產出)", variant="secondary")
    
    with gr.Row():
        status_output = gr.Textbox(label="執行狀態日誌", lines=15)
        file_output = gr.File(label="下載 Excel 結果 (或超時錯誤截圖)")
    
    submit_btn.click(
        fn=scrape_fda,
        inputs=[factory_input, start_year_input, end_year_input, start_record_input],
        outputs=[file_output, status_output]
    )
    
    test_btn.click(
        fn=test_excel_layout,
        inputs=[],
        outputs=[file_output, status_output]
    )

if __name__ == "__main__":
    demo.launch()