app.py

import base64
import io
import json
import logging
import mimetypes
import os
import shutil
import tempfile
import traceback
from datetime import datetime
from pathlib import Path
from typing import List, Tuple

import gradio as gr
from dotenv import load_dotenv
from google import genai
from google.genai import types
from PIL import Image
from pypdf import PdfReader

# 設定 logging
logging.basicConfig(
    level=logging.DEBUG,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

load_dotenv()

# 文字理解與分析使用 gemini-2.5-flash
TEXT_MODEL = "gemini-2.5-flash"
# 圖片生成使用 gemini-3-pro-image-preview  
IMAGE_MODEL = "gemini-3-pro-image-preview"
DEFAULT_API_KEY = os.getenv("GEMINI_API_KEY", "")

# 答案卷儲存路徑
ANSWERS_DIR = Path(__file__).parent / "data" / "answers"
ANSWERS_DIR.mkdir(parents=True, exist_ok=True)
# 答案卷圖片儲存路徑
ANSWER_IMAGES_DIR = Path(__file__).parent / "data" / "answer_images"
ANSWER_IMAGES_DIR.mkdir(parents=True, exist_ok=True)


def get_safe_name(name: str) -> str:
    """產生安全的檔名"""
    safe_name = "".join(c for c in name if c.isalnum() or c in " _-").strip()
    if not safe_name:
        safe_name = datetime.now().strftime("%Y%m%d_%H%M%S")
    return safe_name


def load_saved_answers() -> Tuple[dict, dict]:
    """從檔案系統載入已儲存的答案卷，返回 (文字內容, 圖片路徑列表)"""
    answers = {}
    answer_images = {}
    try:
        for json_file in ANSWERS_DIR.glob("*.json"):
            try:
                with open(json_file, "r", encoding="utf-8") as f:
                    data = json.load(f)
                    name = data.get("name", json_file.stem)
                    answers[name] = data.get("content", "")
                    # 載入對應的圖片路徑
                    image_paths = data.get("image_paths", [])
                    if image_paths:
                        answer_images[name] = image_paths
            except Exception as e:
                logger.warning(f"載入答案卷失敗 {json_file}: {e}")
    except Exception as e:
        logger.error(f"掃描答案卷目錄失敗: {e}")
    return answers, answer_images


def save_answer_to_file(name: str, content: str, image_paths: List[str] = None) -> bool:
    """儲存答案卷到檔案，包含原始圖片"""
    try:
        safe_name = get_safe_name(name)
        
        # 儲存圖片到專用目錄
        saved_image_paths = []
        if image_paths:
            answer_img_dir = ANSWER_IMAGES_DIR / safe_name
            answer_img_dir.mkdir(parents=True, exist_ok=True)
            
            for idx, img_path in enumerate(image_paths):
                src_path = Path(img_path)
                if src_path.exists():
                    # 複製圖片到永久目錄
                    ext = src_path.suffix or ".png"
                    dest_path = answer_img_dir / f"page_{idx+1}{ext}"
                    
                    # 複製檔案
                    shutil.copy2(src_path, dest_path)
                    saved_image_paths.append(str(dest_path))
                    logger.info(f"答案卷圖片已儲存: {dest_path}")
        
        file_path = ANSWERS_DIR / f"{safe_name}.json"
        data = {
            "name": name,
            "content": content,
            "image_paths": saved_image_paths,
            "updated_at": datetime.now().isoformat(),
        }
        with open(file_path, "w", encoding="utf-8") as f:
            json.dump(data, f, ensure_ascii=False, indent=2)
        logger.info(f"答案卷已儲存: {file_path}")
        return True
    except Exception as e:
        logger.error(f"儲存答案卷失敗: {e}")
        logger.error(traceback.format_exc())
        return False


def delete_answer_file(name: str) -> bool:
    """刪除答案卷檔案和圖片"""
    try:
        safe_name = get_safe_name(name)
        
        # 刪除 JSON 檔案
        file_path = ANSWERS_DIR / f"{safe_name}.json"
        if file_path.exists():
            file_path.unlink()
            logger.info(f"答案卷已刪除: {file_path}")
        
        # 刪除圖片目錄
        img_dir = ANSWER_IMAGES_DIR / safe_name
        if img_dir.exists():
            shutil.rmtree(img_dir)
            logger.info(f"答案卷圖片目錄已刪除: {img_dir}")
        
        # 嘗試找到匹配的檔案（以 name 欄位匹配）
        for json_file in ANSWERS_DIR.glob("*.json"):
            try:
                with open(json_file, "r", encoding="utf-8") as f:
                    data = json.load(f)
                    if data.get("name") == name:
                        json_file.unlink()
                        logger.info(f"答案卷已刪除: {json_file}")
                        return True
            except:
                pass
        return True
    except Exception as e:
        logger.error(f"刪除答案卷失敗: {e}")
        return False


def get_client(api_key: str) -> genai.Client:
    """建立 Gemini Client"""
    key = api_key.strip() or DEFAULT_API_KEY
    if not key:
        raise ValueError("請先輸入 Gemini API Key，或在 .env 設定 GEMINI_API_KEY。")
    return genai.Client(api_key=key)


def load_image_from_path(path: Path) -> Image.Image:
    """從路徑載入 PIL Image"""
    return Image.open(path)


def extract_pdf_text(path: Path) -> str:
    """Lightweight PDF text extraction"""
    try:
        reader = PdfReader(path)
        text = "\n\n".join((page.extract_text() or "") for page in reader.pages)
        return text.strip()
    except Exception:
        return ""


def pdf_to_images(path: Path) -> List[Image.Image]:
    """將 PDF 轉換為圖片列表（使用 pdf2image 或 pypdfium2）"""
    try:
        import pdf2image
        images = pdf2image.convert_from_path(path)
        return images
    except ImportError:
        pass
    
    try:
        import pypdfium2 as pdfium
        pdf = pdfium.PdfDocument(path)
        images = []
        for page in pdf:
            bitmap = page.render(scale=2)
            pil_image = bitmap.to_pil()
            images.append(pil_image)
        return images
    except ImportError:
        pass
    
    return []


def build_file_contents(paths: List[Path]) -> Tuple[List, List[str]]:
    """
    建立要傳給 Gemini 的內容列表（文字與圖片）
    返回 (contents, warnings)
    """
    contents: List = []
    warnings: List[str] = []

    for path in paths:
        suffix = path.suffix.lower()
        if suffix in {".png", ".jpg", ".jpeg", ".webp", ".gif"}:
            img = load_image_from_path(path)
            contents.append(img)
        elif suffix == ".pdf":
            # 嘗試將 PDF 轉為圖片
            pdf_images = pdf_to_images(path)
            if pdf_images:
                contents.extend(pdf_images)
            else:
                # 如果無法轉圖，至少提取文字
                pdf_text = extract_pdf_text(path)
                if pdf_text:
                    contents.append(f"[PDF {path.name} 文字]\n{pdf_text}")
                else:
                    warnings.append(f"{path.name} 無法解析，建議轉成圖片上傳。")
        else:
            warnings.append(f"{path.name} 格式不支援，已略過。")

    return contents, warnings


def to_gallery_items(paths: List[str]) -> List[tuple]:
    """Convert file paths to gallery tuples (path, caption)."""
    items: List[tuple] = []
    for idx, p in enumerate(paths, start=1):
        items.append((p, f"第 {idx} 頁"))
    return items


def parse_answer_key(api_key: str, files: List[str]) -> Tuple[str, str]:
    """解析答案卷，輸出 JSON 格式的標準答案"""
    logger.info(f"=== parse_answer_key 開始 ===")
    logger.info(f"files: {files}")
    logger.info(f"api_key 長度: {len(api_key) if api_key else 0}")
    
    if not files:
        logger.warning("沒有上傳檔案")
        return "", "請先上傳答案卷 (PDF / 圖片)。"

    paths = [Path(f) for f in files]
    logger.info(f"paths: {paths}")
    
    file_contents, warnings = build_file_contents(paths)
    logger.info(f"file_contents 數量: {len(file_contents)}")
    logger.info(f"file_contents 類型: {[type(c).__name__ for c in file_contents]}")
    logger.info(f"warnings: {warnings}")

    if not file_contents:
        logger.error("無法解析上傳的檔案")
        return "", "無法解析上傳的檔案，請確認格式正確。"

    prompt = """
你是一名考卷答案解析助手。請仔細辨識這份解答卷中的所有題目，包括：
- 選擇題、填充題、計算題、應用題等
- 數學公式、手寫答案

請抽取題號、標準答案與每題分數，輸出 JSON 陣列：
[
  {"question_number": "一、1", "answer": "標準答案", "points": 5},
  {"question_number": "二、1", "answer": "數學公式或步驟", "points": 10},
  ...
]

規則：
1. 保持題號原有格式與排序
2. 數學公式用 LaTeX 表示
3. 若缺少分數資訊，依題目難度合理估計
4. 只輸出 JSON，不要其他說明
"""

    try:
        logger.info(f"正在建立 Gemini Client...")
        client = get_client(api_key)
        logger.info(f"Client 建立成功，準備呼叫 API")
        logger.info(f"使用模型: {TEXT_MODEL}")
        logger.info(f"內容數量: {len([prompt] + file_contents)}")
        
        response = client.models.generate_content(
            model=TEXT_MODEL,
            contents=[prompt] + file_contents,
            config=types.GenerateContentConfig(
                temperature=0.2,
            )
        )
        
        logger.info(f"API 回應成功")
        logger.info(f"response.parts 數量: {len(response.parts) if response.parts else 0}")
        
        result_text = ""
        for part in response.parts:
            if part.text:
                result_text += part.text
                logger.debug(f"取得文字長度: {len(part.text)}")

        logger.info(f"解析結果長度: {len(result_text)}")
        warning_text = "；".join(warnings) if warnings else "已自動解析，老師可在下方微調後直接使用。"
        return result_text.strip(), warning_text
    except Exception as exc:
        logger.error(f"解析失敗: {exc}")
        logger.error(traceback.format_exc())
        return "", f"解析失敗：{exc}"


def grade_exam_with_image(
    api_key: str,
    answer_text: str,
    answer_image_paths: List[str],
    student_files: List[str],
    generate_teacher_feedback: bool = True,
) -> Tuple[str, List[str], str]:
    """
    批改考卷並生成帶有批改標記的圖片
    
    Args:
        api_key: Gemini API Key
        answer_text: 解析後的答案 JSON 文字
        answer_image_paths: 原始答案卷圖片路徑列表
        student_files: 學生考卷圖片路徑列表
    
    返回 (評語文字, 批改圖片路徑列表, 狀態訊息)
    """
    if not answer_text.strip() and not answer_image_paths:
        return "", [], "請先選擇或建立一份答案卷，再上傳學生考卷。"
    if not student_files:
        return "", [], "請上傳學生考卷圖片 (建議使用清晰的 PNG/JPG)。"

    # 載入學生考卷圖片
    student_paths = [Path(f) for f in student_files]
    student_contents, warnings = build_file_contents(student_paths)

    if not student_contents:
        return "", [], "無法解析上傳的學生考卷。"
    
    # 載入答案卷圖片
    answer_images = []
    for img_path in answer_image_paths:
        p = Path(img_path)
        if p.exists():
            try:
                img = Image.open(p)
                answer_images.append(img)
                logger.info(f"載入答案卷圖片: {p}")
            except Exception as e:
                logger.warning(f"無法載入答案卷圖片 {p}: {e}")
    
    logger.info(f"答案卷圖片數量: {len(answer_images)}, 學生考卷數量: {len(student_contents)}")

    try:
        client = get_client(api_key)
        
        # 使用圖片生成模型，同時輸出批改圖片和文字結果
        # 這樣可以確保圖片標記和文字結果一致
        graded_images = []
        all_feedback_parts = []
        
        for idx, student_img in enumerate(student_contents):
            if not isinstance(student_img, Image.Image):
                continue
            
            # 組合批改提示 - 要求同時輸出文字結果和批改圖片
            grading_prompt = f"""你是一位經驗豐富、嚴謹的老師。請仔細批改這張學生考卷。

【重要】請嚴格按照標準答案進行批改，不要自行判斷答案！

我會提供：
1. 標準答案卷圖片（包含所有正確答案，可能有手寫內容）
2. 解析後的答案 JSON（供參考）
3. 學生考卷圖片

請完成以下兩件事：

【任務一】輸出批改結果文字：
請用以下格式輸出第 {idx+1} 頁的批改結果：

### 第 {idx+1} 頁批改結果
| 題號 | 學生答案 | 正確答案 | 得分 | 滿分 | 評語 |
|------|----------|----------|------|------|------|
| 1 | (學生寫的) | (正確答案) | X | X | 正確/錯誤原因 |
...

**本頁得分：XX / XX 分**

【任務二】生成批改後的考卷圖片：
在學生考卷圖片上標記：
- ✓ 綠色打勾：答案正確
- ✗ 紅色打叉：答案錯誤，旁邊寫正確答案
- 每題標註得分
- 底部寫總分

---
=== 以下是【標準答案卷圖片】（以此為準！）===
"""
            
            contents = [grading_prompt]
            
            # 加入答案卷圖片
            if answer_images:
                contents.extend(answer_images)
            
            # 加入 JSON 答案作為補充
            if answer_text.strip():
                contents.append(f"\n=== 以下是【解析後的答案 JSON】（輔助參考）===\n{answer_text.strip()}\n")
            
            # 加入學生考卷
            contents.append("=== 以下是【學生考卷圖片】（請批改這張）===")
            contents.append(student_img)
            
            logger.info(f"批改第 {idx+1} 頁，傳送內容數量: {len(contents)}")
            
            try:
                response = client.models.generate_content(
                    model=IMAGE_MODEL,
                    contents=contents,
                    config=types.GenerateContentConfig(
                        response_modalities=['TEXT', 'IMAGE'],
                        image_config=types.ImageConfig(
                            aspect_ratio="3:4",
                            image_size="2K",
                        ),
                        temperature=0.2,
                    )
                )
                
                # 從回應中提取文字和圖片
                page_feedback = ""
                for part in response.parts:
                    if hasattr(part, 'text') and part.text:
                        page_feedback += part.text
                    if hasattr(part, 'inline_data') and part.inline_data:
                        img = part.as_image()
                        temp_path = tempfile.mktemp(suffix=f"_graded_{idx+1}.png")
                        img.save(temp_path)
                        graded_images.append(temp_path)
                        logger.info(f"批改圖片已生成: {temp_path}")
                
                if page_feedback:
                    all_feedback_parts.append(page_feedback)
                    logger.info(f"第 {idx+1} 頁批改文字長度: {len(page_feedback)}")
                
            except Exception as img_exc:
                logger.error(f"第 {idx+1} 頁批改失敗: {img_exc}")
                warnings.append(f"第 {idx+1} 頁批改失敗：{img_exc}")

        # 組合所有頁面的批改結果
        combined_feedback = "\n\n".join(all_feedback_parts)
        
        # 如果有多頁，再生成總結評語
        if len(student_contents) > 0 and combined_feedback:
            try:
                summary_prompt = f"""根據以下批改結果，請以溫暖關懷的語氣寫一段「老師的評語」：

{combined_feedback}

請用以下格式：

---

## 老師的評語

親愛的同學：
（在這裡寫出具體的回饋，包括：
- 肯定做得好的部分
- 指出需要加強的觀念
- 給予鼓勵）
"""
                summary_response = client.models.generate_content(
                    model=TEXT_MODEL,
                    contents=[summary_prompt],
                    config=types.GenerateContentConfig(temperature=0.4)
                )
                
                teacher_comment = ""
                for part in summary_response.parts:
                    if part.text:
                        teacher_comment += part.text
                
                if teacher_comment:
                    combined_feedback += "\n\n" + teacher_comment
                    
            except Exception as e:
                logger.warning(f"生成老師評語失敗: {e}")

        warning_text = "；".join(warnings) if warnings else "批改完成！"
        return combined_feedback, graded_images, warning_text

    except Exception as exc:
        logger.error(f"批改失敗: {exc}")
        logger.error(traceback.format_exc())
        return "", [], f"批改失敗：{exc}"


def guess_exam_name(api_key: str, answer_text: str, file_names: List[str]) -> str:
    """Try to produce a friendly exam name; fallback to file stem."""
    fallback = Path(file_names[0]).stem if file_names else "未命名考卷"
    if not answer_text.strip():
        return fallback
    prompt = f"""請幫老師為這份考卷取一個簡短的中文名稱。

規則：
- 長度必須在 6-16 字之間
- 只回傳名稱本身，不要加引號、標點符號或任何說明
- 範例：「四年級數學期中考」、「五上自然第三單元」

可參考的答案內容：
{answer_text[:800]}

請直接輸出名稱（6-16字）："""
    try:
        client = get_client(api_key)
        response = client.models.generate_content(
            model=TEXT_MODEL,
            contents=[prompt],
            config=types.GenerateContentConfig(
                temperature=0.2,
                max_output_tokens=30,  # 限制輸出長度
            )
        )
        name = ""
        for part in response.parts:
            if part.text:
                name += part.text
        cleaned = name.replace("\n", "").replace("「", "").replace("」", "").replace('"', '').replace("'", "").strip()
        # 強制限制長度
        if len(cleaned) > 20:
            cleaned = cleaned[:20]
        return cleaned or fallback
    except Exception:
        return fallback


def parse_and_store_answer(
    api_key: str,
    files: List[str],
    provided_name: str,
    stored: dict,
    stored_images: dict,
):
    logger.info(f"parse_and_store_answer 開始執行，檔案數量: {len(files) if files else 0}")
    stored = stored or {}
    stored_images = stored_images or {}
    parsed, status = parse_answer_key(api_key, files)
    if not parsed:
        return "", status, provided_name, gr.update(choices=list(stored.keys()), value=None, visible=True), stored, stored_images, to_gallery_items(files)

    final_name = provided_name.strip() if provided_name and provided_name.strip() else guess_exam_name(
        api_key, parsed, [Path(f).name for f in files]
    )
    stored = dict(stored)
    stored[final_name] = parsed
    
    # 處理圖片路徑（從 PDF 轉換的圖片或原始圖片）
    image_paths = []
    for f in files:
        p = Path(f)
        if p.suffix.lower() == '.pdf':
            # PDF 轉成的圖片會暫存，我們要取得它們
            pdf_images = pdf_to_images(p)
            for idx, img in enumerate(pdf_images):
                temp_path = tempfile.mktemp(suffix=f"_answer_{idx}.png")
                img.save(temp_path)
                image_paths.append(temp_path)
        else:
            image_paths.append(f)
    
    # 自動儲存到檔案（包含圖片）
    save_answer_to_file(final_name, parsed, image_paths)
    
    # 更新 stored_images
    stored_images = dict(stored_images)
    # 讀取已儲存的圖片路徑
    safe_name = get_safe_name(final_name)
    json_path = ANSWERS_DIR / f"{safe_name}.json"
    if json_path.exists():
        with open(json_path, "r", encoding="utf-8") as f:
            data = json.load(f)
            stored_images[final_name] = data.get("image_paths", [])

    choices = sorted(stored.keys())
    msg = f"✅ 解答已確認「{final_name}」並已儲存（含原始圖片）。{status}"
    logger.info(f"答案解析完成，名稱: {final_name}，圖片數量: {len(stored_images.get(final_name, []))}")
    return (
        parsed,
        msg,
        final_name,
        gr.update(choices=choices, value=final_name, visible=True),
        stored,
        stored_images,
        to_gallery_items(files),
    )


def load_answer_from_storage(selected: str, stored: dict):
    stored = stored or {}
    if not selected or selected not in stored:
        return "", ""
    return stored[selected], selected


def save_manual_edit(selected: str, name_input: str, edited_text: str, stored: dict):
    stored = stored or {}
    if not edited_text.strip():
        return stored, "請先確認答案 JSON 內容再儲存。", gr.update()
    if not (selected or name_input):
        return stored, "請先輸入或選擇考卷名稱，再儲存。", gr.update()

    final_name = (name_input or selected or "未命名考卷").strip()
    stored = dict(stored)
    stored[final_name] = edited_text.strip()
    
    # 儲存到檔案
    save_answer_to_file(final_name, edited_text.strip())
    
    choices = sorted(stored.keys())
    status = f"✅ 已儲存「{final_name}」的修正版本至檔案。"
    return stored, status, gr.update(choices=choices, value=final_name)


def delete_answer_from_storage(selected: str, stored: dict):
    """刪除選中的答案卷"""
    stored = stored or {}
    if not selected:
        return stored, "請先選擇要刪除的答案卷。", gr.update(), "", ""
    
    stored = dict(stored)
    if selected in stored:
        del stored[selected]
    
    # 從檔案系統刪除
    delete_answer_file(selected)
    
    choices = sorted(stored.keys())
    status = f"🗑️ 已刪除「{selected}」。"
    return stored, status, gr.update(choices=choices, value=None), "", ""


def grade_exam_from_storage(
    api_key: str,
    selected_exam: str,
    stored: dict,
    stored_images: dict,
    files: List[str],
):
    """執行批改並返回結果，同時傳遞答案卷圖片"""
    logger.info(f"grade_exam_from_storage 開始執行")
    logger.info(f"選中的考卷: {selected_exam}")
    logger.info(f"stored keys: {list(stored.keys()) if stored else 'None'}")
    logger.info(f"學生考卷檔案數量: {len(files) if files else 0}")
    
    stored = stored or {}
    stored_images = stored_images or {}
    if not selected_exam:
        logger.warning("沒有選擇考卷")
        return "", [], "請先選擇一份已解析的答案卷。"
    answer_text = stored.get(selected_exam, "")
    if not answer_text:
        logger.warning(f"找不到 {selected_exam} 的內容")
        return "", [], "找不到這份答案卷的內容，請重新上傳或選擇其他考卷。"
    
    # 取得答案卷的原始圖片
    answer_image_paths = stored_images.get(selected_exam, [])
    logger.info(f"開始批改，答案長度: {len(answer_text)}，答案卷圖片數量: {len(answer_image_paths)}")
    
    return grade_exam_with_image(api_key, answer_text, answer_image_paths, files)


def update_gallery(files: List[str]):
    return to_gallery_items(files or [])


def update_graded_gallery(images: List[str]):
    """更新批改結果的圖片展示"""
    if not images:
        return []
    return [(img, f"批改結果 {idx+1}") for idx, img in enumerate(images)]


def clear_api_key():
    """移除 API Key"""
    return "", "🗑️ API Key 已移除"


CUSTOM_CSS = """
.card {
    background: linear-gradient(135deg, #f8fafc 0%, #f1f5f9 100%);
    border-radius: 16px;
    border: 1px solid #e2e8f0;
    padding: 20px;
    box-shadow: 0 4px 20px rgba(0, 0, 0, 0.05);
    margin-bottom: 16px;
}
.primary-button {
    font-weight: 700;
    height: 52px;
    font-size: 16px;
    background: linear-gradient(135deg, #6366f1 0%, #8b5cf6 100%);
    border: none;
    border-radius: 12px;
}
.primary-button:hover {
    background: linear-gradient(135deg, #4f46e5 0%, #7c3aed 100%);
}
.section-title {
    font-size: 18px;
    font-weight: 700;
    color: #1e293b;
    margin-bottom: 12px;
}
.status-ok {
    color: #16a34a;
    font-weight: 600;
    background: #dcfce7;
    padding: 8px 16px;
    border-radius: 8px;
    display: inline-block;
}
.status-warn {
    color: #ea580c;
    font-weight: 600;
    background: #fff7ed;
    padding: 8px 16px;
    border-radius: 8px;
}
.api-key-card {
    background: linear-gradient(135deg, #f0fdf4 0%, #dcfce7 100%);
    border: 1px solid #86efac;
}
.header-title {
    text-align: center;
    font-size: 28px;
    font-weight: 800;
    background: linear-gradient(135deg, #6366f1, #8b5cf6);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    margin-bottom: 8px;
}
.header-subtitle {
    text-align: center;
    color: #64748b;
    font-size: 16px;
    margin-bottom: 24px;
}
.feedback-box {
    background: #fffbeb;
    border: 1px solid #fcd34d;
    border-radius: 12px;
    padding: 16px;
    font-size: 15px;
    line-height: 1.8;
}
"""


def build_demo() -> gr.Blocks:
    # 啟動時載入已儲存的答案卷（包含圖片路徑）
    initial_answers, initial_images = load_saved_answers()
    initial_choices = sorted(initial_answers.keys())
    
    with gr.Blocks(title="老師閱卷助手") as demo:
        # 注入 CSS
        gr.HTML(f"<style>{CUSTOM_CSS}</style>")
        stored_answers = gr.State(initial_answers)
        stored_answer_images = gr.State(initial_images)  # 新增：儲存答案卷圖片路徑

        gr.HTML("""
        <div style="text-align: center; padding: 20px 0;">
            <h1 class="header-title">🎓 AI 智慧閱卷助手</h1>
            <p class="header-subtitle">上傳解答與考卷，讓 AI 為您自動批改與評分</p>
        </div>
        """)

        # Step 1: API Key
        with gr.Group(elem_classes="card api-key-card"):
            gr.Markdown("### 🔑 Gemini API Key", elem_classes="section-title")
            with gr.Row():
                api_key = gr.Textbox(
                    placeholder="請輸入您的 Gemini API Key",
                    type="password",
                    label="",
                    value=DEFAULT_API_KEY,
                    scale=4,
                )
                clear_key_btn = gr.Button("🗑️ 移除", scale=1, variant="secondary")
            api_status = gr.Markdown(
                "● API Key 已儲存" if DEFAULT_API_KEY else "請輸入 API Key",
                elem_classes="status-ok" if DEFAULT_API_KEY else ""
            )
            gr.Markdown(
                "💡 您的 API Key 僅會儲存在瀏覽器中，不會傳送至任何伺服器。",
                elem_classes="status-warn"
            )

        # Step 2: 答案卷管理
        with gr.Group(elem_classes="card"):
            gr.Markdown("### 📝 答案卷管理", elem_classes="section-title")
            
            # 已儲存的答案卷選擇器（顯示在最上方）
            with gr.Row():
                exam_picker = gr.Dropdown(
                    label="📂 已儲存的答案卷",
                    choices=initial_choices,
                    value=initial_choices[0] if initial_choices else None,
                    interactive=True,
                    visible=True,
                    scale=3,
                )
                delete_answer_btn = gr.Button("🗑️ 刪除", variant="secondary", scale=1)
            
            with gr.Row():
                with gr.Column(scale=2):
                    exam_name = gr.Textbox(
                        label="考卷名稱（可留空讓 AI 協助命名）",
                        placeholder="例：四年級下學期數學單元評量",
                    )
                    answer_upload = gr.File(
                        label="上傳新的答案卷 (PDF / 圖片，可多檔)",
                        file_count="multiple",
                        file_types=[".pdf", ".png", ".jpg", ".jpeg", ".webp"],
                        type="filepath",
                    )
                    with gr.Row():
                        modify_answer_btn = gr.Button("📝 修改解答", variant="secondary")
                        save_button = gr.Button("💾 儲存修正", variant="secondary")
                with gr.Column(scale=3):
                    answer_gallery = gr.Gallery(
                        label="答案卷預覽",
                        show_label=True,
                        height=200,
                        columns=4,
                    )
            answer_status = gr.Markdown("")
            answer_preview = gr.Textbox(
                label="標準答案 (老師可在此微調)",
                lines=8,
                placeholder="上傳答案卷後，AI 會自動解析答案...",
                visible=False,
            )

        # Step 3: 學生考卷上傳
        with gr.Group(elem_classes="card"):
            gr.Markdown("### 📷 拍攝/上傳學生考卷", elem_classes="section-title")
            with gr.Row():
                student_upload = gr.File(
                    label="上傳學生考卷 (可多頁拍攝)",
                    file_count="multiple",
                    file_types=[".png", ".jpg", ".jpeg", ".webp"],
                    type="filepath",
                )
            student_gallery = gr.Gallery(
                label="考卷頁面預覽",
                show_label=True,
                height=280,
                columns=4,
            )
            gr.Markdown("💡 支援多頁拍攝。請確保光線充足、字跡清晰。", elem_classes="status-warn")

        # 批改設定區塊
        with gr.Group(elem_classes="card"):
            gr.Markdown("### ⚙️ 批改設定", elem_classes="section-title")
            with gr.Row():
                grading_exam_picker = gr.Dropdown(
                    label="🎯 選擇要使用的答案卷（批改依據）",
                    choices=initial_choices,
                    value=initial_choices[0] if initial_choices else None,
                    interactive=True,
                    scale=3,
                )
                refresh_answers_btn = gr.Button("🔄 重新整理", variant="secondary", scale=1)
            
            gr.Markdown("⚠️ **請確認已選擇正確的答案卷再開始批改！**", elem_classes="status-warn")
        
        # 批改按鈕
        grade_button = gr.Button(
            "🚀 開始批改考卷",
            variant="primary",
            elem_classes="primary-button",
            size="lg",
        )

        # Step 4: 批改結果
        with gr.Group(elem_classes="card"):
            gr.Markdown("### 📊 批改結果", elem_classes="section-title")
            
            with gr.Row():
                with gr.Column(scale=1):
                    graded_gallery = gr.Gallery(
                        label="批改後的考卷",
                        show_label=True,
                        height=400,
                        columns=2,
                    )
                with gr.Column(scale=1):
                    grade_feedback = gr.Markdown(
                        value="*批改結果將顯示在這裡...*",
                        elem_classes="feedback-box",
                    )
            
            grade_status = gr.Markdown("")

        # Event handlers
        clear_key_btn.click(
            fn=clear_api_key,
            outputs=[api_key, api_status],
        )
        
        modify_answer_btn.click(
            fn=lambda: gr.update(visible=True),
            outputs=[answer_preview],
        )

        # 上傳答案卷後，同時更新兩個選擇器
        def on_answer_uploaded(api_key, files, name, stored, stored_images):
            result = parse_and_store_answer(api_key, files, name, stored, stored_images)
            # result: (preview, status, name, picker_update, stored, stored_images, gallery)
            # 需要同時更新 exam_picker 和 grading_exam_picker
            preview, status, final_name, picker_update, new_stored, new_images, gallery = result
            choices = sorted(new_stored.keys()) if new_stored else []
            return (
                preview, 
                status, 
                final_name, 
                gr.update(choices=choices, value=final_name),  # exam_picker
                gr.update(choices=choices, value=final_name),  # grading_exam_picker
                new_stored, 
                new_images, 
                gallery
            )

        answer_upload.upload(
            fn=on_answer_uploaded,
            inputs=[api_key, answer_upload, exam_name, stored_answers, stored_answer_images],
            outputs=[answer_preview, answer_status, exam_name, exam_picker, grading_exam_picker, stored_answers, stored_answer_images, answer_gallery],
        )

        # 當答案卷管理的選擇器變更時，同步更新批改選擇器
        def sync_pickers(selected, stored):
            content, name = load_answer_from_storage(selected, stored)
            return content, name, gr.update(value=selected)

        exam_picker.change(
            fn=sync_pickers,
            inputs=[exam_picker, stored_answers],
            outputs=[answer_preview, exam_name, grading_exam_picker],
        )
        
        # 當批改選擇器變更時，同步更新答案卷管理的選擇器
        grading_exam_picker.change(
            fn=lambda selected: gr.update(value=selected),
            inputs=[grading_exam_picker],
            outputs=[exam_picker],
        )
        
        # 重新整理按鈕
        def refresh_answer_choices(stored):
            choices = sorted(stored.keys()) if stored else []
            return gr.update(choices=choices), gr.update(choices=choices)
        
        refresh_answers_btn.click(
            fn=refresh_answer_choices,
            inputs=[stored_answers],
            outputs=[exam_picker, grading_exam_picker],
        )

        save_button.click(
            fn=save_manual_edit,
            inputs=[exam_picker, exam_name, answer_preview, stored_answers],
            outputs=[stored_answers, answer_status, exam_picker],
        )
        
        # 刪除答案卷按鈕 - 同時更新兩個選擇器
        def on_delete_answer(selected, stored):
            new_stored, status, picker_update, preview, name = delete_answer_from_storage(selected, stored)
            choices = sorted(new_stored.keys()) if new_stored else []
            return (
                new_stored, 
                status, 
                gr.update(choices=choices, value=None),  # exam_picker
                gr.update(choices=choices, value=None),  # grading_exam_picker
                preview, 
                name
            )
        
        delete_answer_btn.click(
            fn=on_delete_answer,
            inputs=[exam_picker, stored_answers],
            outputs=[stored_answers, answer_status, exam_picker, grading_exam_picker, answer_preview, exam_name],
        )

        student_upload.upload(
            fn=update_gallery,
            inputs=[student_upload],
            outputs=[student_gallery],
        )

        def do_grading(api_key, grading_picker, stored, stored_images, files):
            logger.info(f"do_grading 被呼叫，grading_picker={grading_picker}")
            logger.info(f"stored_images keys: {list(stored_images.keys()) if stored_images else 'None'}")
            if not grading_picker:
                return "**❌ 請先選擇要使用的答案卷！**", [], "請在「批改設定」中選擇答案卷"
            feedback, images, status = grade_exam_from_storage(api_key, grading_picker, stored, stored_images, files)
            gallery_items = update_graded_gallery(images)
            return feedback, gallery_items, status

        grade_button.click(
            fn=do_grading,
            inputs=[api_key, grading_exam_picker, stored_answers, stored_answer_images, student_upload],
            outputs=[grade_feedback, graded_gallery, grade_status],
        )

    return demo


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