#!/usr/bin/env python3
"""
PDF OCR & Markdown変換ツール (v5.latest)

最新のgoogle.genai SDKを使用したPDFからMarkdownへの変換アプリ
"""

import os
import socket
import tempfile
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path
from typing import List, Tuple

import gradio as gr
import fitz  # PyMuPDF
from google import genai
from google.genai import types

from presets import SYSTEM_PROMPTS, list_presets


# 環境変数から設定を読み込む
GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY", "")

if not GEMINI_API_KEY:
    print("⚠️ 警告: GEMINI_API_KEYが設定されていません。")
    print("環境変数またはHugging Face Spacesのシークレットに設定してください。")


def split_pdf(pdf_path: str, output_dir: str, pages_per_chunk: int) -> List[Tuple[int, str]]:
    """
    PDFを指定ページ数ごとに分割する

    Args:
        pdf_path: 元のPDFファイルパス
        output_dir: 分割PDFの保存先ディレクトリ
        pages_per_chunk: 何ページごとに分割するか

    Returns:
        (開始ページ番号, 分割PDFパス)のリスト
    """
    pdf_document = fitz.open(pdf_path)
    total_pages = len(pdf_document)

    split_pdfs = []

    for start_page in range(0, total_pages, pages_per_chunk):
        end_page = min(start_page + pages_per_chunk - 1, total_pages - 1)

        # 新しいPDFドキュメントを作成
        output_pdf = fitz.open()

        # 指定範囲のページを新しいPDFに追加
        for page_num in range(start_page, end_page + 1):
            output_pdf.insert_pdf(pdf_document, from_page=page_num, to_page=page_num)

        # 分割したPDFを保存
        output_path = os.path.join(
            output_dir, f"split_{start_page+1}_to_{end_page+1}.pdf"
        )
        output_pdf.save(output_path)
        output_pdf.close()

        split_pdfs.append((start_page, output_path))

    pdf_document.close()
    return split_pdfs


def process_pdf_chunk(
    pdf_path: str,
    system_instruction: str,
    model_name: str,
    enable_thinking: bool,
) -> str:
    """
    1つのPDFチャンクをGemini APIで処理

    Args:
        pdf_path: PDFファイルパス
        system_instruction: システムインストラクション
        model_name: 使用するモデル名
        enable_thinking: Thinking Mode有効化

    Returns:
        生成されたMarkdown文字列
    """
    # PDFをバイナリ読み込み
    pdf_bytes = Path(pdf_path).read_bytes()

    # Geminiクライアント初期化
    client = genai.Client(api_key=GEMINI_API_KEY)

    # リクエストコンテンツを構築
    contents = [
        types.Content(
            role="user",
            parts=[
                types.Part.from_text(
                    text="添付のPDFを人間が読む順序に従ってMarkdownへ変換してください。"
                ),
                types.Part.from_bytes(
                    data=pdf_bytes,
                    mime_type="application/pdf"
                ),
            ],
        ),
    ]

    # 生成設定
    config = types.GenerateContentConfig(
        system_instruction=[
            types.Part.from_text(text=system_instruction),
        ],
    )

    # Thinking Mode設定
    if enable_thinking:
        try:
            config.thinking_config = types.ThinkingConfig()
        except (TypeError, ValueError) as thinking_error:
            print(f"⚠️ ThinkingConfig設定時にエラーが発生しました: {thinking_error}")

    # ストリーミング処理
    collected_chunks: List[str] = []

    try:
        for chunk in client.models.generate_content_stream(
            model=model_name,
            contents=contents,
            config=config,
        ):
            if hasattr(chunk, "text") and chunk.text:
                collected_chunks.append(chunk.text)

        return "".join(collected_chunks)

    except Exception as e:
        return f"<!-- エラー: {e} -->\n\n⚠️ このチャンクの処理でエラーが発生しました。"


def process_pdf(
    pdf_file,
    preset_name: str,
    custom_prompt: str,
    pages_per_chunk: int,
    model_name: str,
    enable_thinking: bool,
    progress=gr.Progress(),
) -> str:
    """
    PDFファイルを処理してMarkdownに変換

    Args:
        pdf_file: アップロードされたPDFファイル
        preset_name: プリセット名
        custom_prompt: カスタムシステムプロンプト
        pages_per_chunk: 何ページごとに分割するか
        model_name: 使用するモデル名
        enable_thinking: Thinking Mode有効化
        progress: Gradio進捗バー

    Returns:
        変換されたMarkdown文字列
    """
    if not GEMINI_API_KEY:
        return "❌ エラー: GEMINI_API_KEYが設定されていません。\n\n環境変数またはHugging Face Spacesのシークレットに設定してください。"

    if not pdf_file:
        return "❌ エラー: PDFファイルをアップロードしてください。"

    instruction_text = custom_prompt.strip()

    # システムインストラクション決定
    if preset_name == "カスタム":
        if not instruction_text:
            return "❌ エラー: カスタムプロンプトを入力してください。"
        system_instruction = instruction_text
    else:
        # プリセット選択時もテキストボックスを編集可能にする
        system_instruction = instruction_text or SYSTEM_PROMPTS[preset_name]

    progress_cb = progress if callable(progress) else (lambda *args, **kwargs: None)

    log_progress(progress_cb, 0.1, "PDFを読み込み中...")

    with tempfile.TemporaryDirectory() as temp_dir:
        # PDFファイルパス
        temp_pdf_path = resolve_pdf_path(pdf_file)
        if not temp_pdf_path or not os.path.exists(temp_pdf_path):
            return (
                "❌ エラー: アップロードされたPDFファイルを読み込めませんでした。\n"
                "再度アップロードしてから実行してください。"
            )

        # PDFを分割
        log_progress(progress_cb, 0.2, "PDFを分割中...")
        split_pdf_paths = split_pdf(temp_pdf_path, temp_dir, pages_per_chunk)

        total_chunks = len(split_pdf_paths)
        log_progress(progress_cb, 0.3, f"{total_chunks}個のチャンクに分割完了")

        # 各チャンクを最大3並列で処理
        markdown_results = {}
        max_workers = min(3, total_chunks)

        with ThreadPoolExecutor(max_workers=max_workers) as executor:
            future_map = {}
            for start_page, pdf_path in split_pdf_paths:
                future = executor.submit(
                    process_pdf_chunk,
                    pdf_path=pdf_path,
                    system_instruction=system_instruction,
                    model_name=model_name,
                    enable_thinking=enable_thinking,
                )
                future_map[future] = start_page

            completed = 0
            for future in as_completed(future_map):
                start_page = future_map[future]
                try:
                    result = future.result()
                except Exception as chunk_error:
                    result = (
                        f"<!-- エラー: {chunk_error} -->\n\n⚠️ このチャンクの処理でエラーが発生しました。"
                    )

                markdown_results[start_page] = result
                completed += 1
                progress_value = 0.3 + 0.6 * (completed / total_chunks)
                log_progress(
                    progress_cb,
                    progress_value,
                    f"チャンク {completed}/{total_chunks} を処理中...",
                )

        # 結果を結合
        log_progress(progress_cb, 0.9, "結果を結合中...")

        combined_markdown = "\n\n".join(
            markdown_results[page] for page in sorted(markdown_results.keys())
        )

        log_progress(progress_cb, 1.0, "完了！")

        return combined_markdown


def update_instruction_text(preset_name: str):
    """プリセットに対応するシステムインストラクションをテキストエリアへ反映"""

    if preset_name == "カスタム":
        return gr.update(value="")

    return gr.update(value=SYSTEM_PROMPTS[preset_name])


def log_progress(progress_cb, value: float, desc: str) -> None:
    """進捗バー更新と同時にログへも出力"""

    try:
        progress_cb(value, desc=desc)
    except TypeError:
        progress_cb(value)

    print(f"[Progress] {value * 100:.0f}% - {desc}", flush=True)


def resolve_pdf_path(uploaded_pdf) -> str:
    """
    GradioのFileコンポーネントから渡される値を安全にファイルパスへ変換
    """
    if uploaded_pdf is None:
        return ""

    if isinstance(uploaded_pdf, str):
        return uploaded_pdf

    if isinstance(uploaded_pdf, dict):
        # Gradio 5では辞書形式で渡されるケースもある
        return uploaded_pdf.get("name") or uploaded_pdf.get("path") or ""

    return getattr(uploaded_pdf, "name", "")


def create_interface():
    """Gradio UIを作成"""

    with gr.Blocks(title="PDF OCR & Markdown変換") as demo:
        gr.Markdown("# 📄 PDF OCR & Markdown変換ツール (v5.latest)")
        gr.Markdown(
            "PDFをアップロードすると、Gemini APIでOCR処理してMarkdown形式に変換します。\n\n"
            "✨ **新機能**: システムプロンプトのプリセット選択、分割単位の調整、Thinking Mode"
        )

        with gr.Row():
            with gr.Column(scale=1):
                pdf_input = gr.File(
                    label="📎 PDFファイルをアップロード",
                    file_types=[".pdf"],
                    type="filepath"
                )

                preset_dropdown = gr.Dropdown(
                    choices=list_presets() + ["カスタム"],
                    value="汎用OCR",
                    label="🎯 システムプロンプトのプリセット",
                    info="用途に応じたプリセットを選択"
                )

                custom_prompt_text = gr.Textbox(
                    label="✏️ システムインストラクション",
                    placeholder="プリセットに応じた指示が自動で入力されます。自由に編集できます。",
                    lines=10,
                    value=SYSTEM_PROMPTS["汎用OCR"],
                )

                pages_slider = gr.Slider(
                    minimum=1,
                    maximum=20,
                    value=5,
                    step=1,
                    label="📑 PDFを何ページごとに分割？",
                    info="小さいほど並列処理が増えますが、API呼び出し数も増えます"
                )

                model_dropdown = gr.Dropdown(
                    choices=[
                        "gemini-2.5-flash-preview-tts",
                        "gemini-flash-latest",
                        "gemini-pro-latest"
                    ],
                    value="gemini-flash-latest",
                    label="🤖 使用するモデル",
                    info="Flash: 高速・低コスト、Pro: 高精度"
                )

                thinking_checkbox = gr.Checkbox(
                    label="🧠 Thinking Mode有効化",
                    value=True,
                    info="モデルの思考過程を活用（より正確な結果）"
                )

                convert_btn = gr.Button(
                    "🚀 変換開始",
                    variant="primary",
                    size="lg"
                )

            with gr.Column(scale=2):
                markdown_output = gr.Textbox(
                    label="📝 変換結果（Markdown）",
                    lines=20,
                    max_lines=30,
                    show_copy_button=True
                )

                with gr.Row():
                    download_btn = gr.Button("💾 Markdownをダウンロード")

        preset_dropdown.change(
            fn=update_instruction_text,
            inputs=preset_dropdown,
            outputs=custom_prompt_text
        )

        convert_btn.click(
            fn=process_pdf,
            inputs=[
                pdf_input,
                preset_dropdown,
                custom_prompt_text,
                pages_slider,
                model_dropdown,
                thinking_checkbox,
            ],
            outputs=markdown_output
        )

        download_btn.click(
            None,
            markdown_output,
            [],
            js="""(x) => {
                const blob = new Blob([x], {type: 'text/markdown;charset=utf-8'});
                const url = URL.createObjectURL(blob);
                const a = document.createElement('a');
                a.href = url;
                a.download = 'converted.md';
                document.body.appendChild(a);
                a.click();
                document.body.removeChild(a);
                URL.revokeObjectURL(url);
            }"""
        )

        with gr.Accordion("📖 使用方法", open=False):
            gr.Markdown("""
### 基本的な使い方

1. **PDFアップロード**: 変換したいPDFファイルを選択
2. **プリセット選択**: 用途に応じたプリセットを選択
   - **汎用OCR**: 一般的な文書
   - **教育・参考書**: 教科書や参考書（右→左、上→下の読み順対応）
   - **ビジネス文書**: ビジネスレポート、提案書
   - **スクリーンショット**: ソフトウェアのUI画像
   - **カスタム**: 独自のシステムプロンプトを入力
3. **分割設定**: PDFを何ページごとに分割するか調整（デフォルト: 5ページ）
4. **モデル選択**: Flash（高速）またはPro（高精度）
5. **変換開始**: ボタンをクリックして処理開始
6. **結果確認**: 変換されたMarkdownを確認・ダウンロード

### Tips

- **大きなPDF**: ページ数を減らして分割すると並列処理で高速化
- **高精度**: Proモデル + Thinking Mode ONで最高品質
- **低コスト**: Flashモデル + Thinking Mode OFFで高速・低コスト
- **カスタムプロンプト**: 特定の用途に最適化した独自の指示を作成可能

### API Key設定

Hugging Face Spacesで使用する場合:
1. Spaceの設定で「Secrets」を開く
2. `GEMINI_API_KEY`という名前でGoogle AI StudioのAPIキーを設定
            """)

    return demo


def find_open_port() -> int:
    """利用可能なポート番号を取得"""
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
        sock.bind(("", 0))
        return sock.getsockname()[1]


def launch_demo():
    """Gradioアプリを起動（ポート競合を自動回避）"""
    demo = create_interface()

    server_name = os.getenv("GRADIO_SERVER_NAME") or os.getenv("HOST", "0.0.0.0")
    share = os.getenv("GRADIO_SHARE", "").lower() == "true"

    candidate_ports: List[int | None] = []
    for env_var in ("GRADIO_SERVER_PORT", "PORT"):
        value = os.getenv(env_var)
        if value:
            try:
                candidate_ports.append(int(value))
            except ValueError:
                print(f"⚠️ 無効なポート番号が指定されています ({env_var}={value})")

    # デフォルト → 自動検出 → 最後にGradioのデフォルトへ
    candidate_ports.extend([7860, None])

    last_error: Exception | None = None

    for port in candidate_ports:
        kwargs = {
            "server_name": server_name,
            "share": share,
            "show_error": True,
        }

        if port is None:
            port_to_use = find_open_port()
        else:
            port_to_use = port

        kwargs["server_port"] = port_to_use

        try:
            print(f"ℹ️ Gradio UIを http://{server_name}:{port_to_use} で起動します。")
            demo.launch(**kwargs)
            return
        except OSError as error:
            print(f"⚠️ ポート {port_to_use} の起動に失敗しました: {error}")
            last_error = error

    if last_error:
        raise last_error


if __name__ == "__main__":
    launch_demo()