import os
import sys
import argparse
from pathlib import Path
from PIL import Image, ImageGrab  # ImageGrabを追加

# 添付ライブラリの読み込み
import tkai_lib
from get_from_ai import read_ini

def get_image_from_clipboard():
    """Windows 11のクリップボードから画像を安定して取得する"""
    print("--- クリップボードを確認中 ---")
    try:
        # Windowsのクリップボードから画像を直接取得
        img = ImageGrab.grabclipboard()
        
        if isinstance(img, Image.Image):
            target_path = "clipboard_ocr_input.png"
            img.save(target_path, "PNG")
            print(f"✅ 画像を取得しました: {target_path} ({img.size[0]}x{img.size[1]})")
            return target_path
        else:
            print("❌ クリップボードに画像データが見つかりませんでした。")
            print("   (Win + Shift + S などで画像をコピーしてから実行してください)")
            return None
    except Exception as e:
        print(f"❌ クリップボード取得エラー: {e}")
        return None

def call_ai_ocr(prompt, image_path, api="gemini", model=None):
    img = Image.open(image_path)
    
    if api == "gemini":
        # エラーが出た場合は 'gemini-3.1-pro' または 'gemini-1.5-pro' (安定版) を試してください
        model_name = model or os.getenv("gemini_model", "gemini-3.1-pro")
        api_key = os.getenv("GOOGLE_API_KEY")
        
        print(f"🚀 {api.upper()} [{model_name}] へリクエスト送信中...")
        tkai_lib.genai.configure(api_key=api_key)
        
        # モデルの存在確認と呼び出し
        try:
            m = tkai_lib.genai.GenerativeModel(model_name)
            response = m.generate_content([prompt, img])
            return response.text
        except Exception as e:
            if "404" in str(e):
                print(f"⚠️ モデル '{model_name}' が見つかりません。")
                print("   最新の有効なモデル名を確認してください（例: gemini-1.5-pro など）")
            raise e

def main():
    # 1. AI設定の読み込み (ai.env)
    tkai_lib.read_ai_config("ai.env")

    parser = argparse.ArgumentParser(description="AI OCR to Markdown ツール")
    parser.add_argument("input", help="画像ファイルパス、または 'clip' (クリップボード)")
    parser.add_argument("--output", "-o", default=None, help="出力Markdownファイル名")
    parser.add_argument("--ini", "-i", default="ai_ocr2md.ini", help="プロンプト設定ファイル")
    parser.add_argument("--api", "-a", choices=["gemini", "openai", "openai5"], default="gemini", help="使用するAPI")
    parser.add_argument("--model", "-m", default="gemini-3.1-pro-preview", help="モデル名 (default: gemini-3.1-pro)")
    args = parser.parse_args()

    if args.output is None:
        input_path = Path(args.input)
        args.output = f"{input_path.stem}.md"
    
    print()
    print(f"入力ファイル名: {args.input}")
    print(f"出力ファイル名: {args.output}")
    print(f"プロンプトファイル名: {args.ini}")
    print(f"AI API: {args.api}")
    print(f"AI model: {args.model}")

    # 2. プロンプトの準備
    if os.path.exists(args.ini):
        print(f"プロンプト読み込み: {args.ini}")
        prompt = read_ini(args.ini).get("PROMPT", "画像をOCRしてください。")
    else:
        prompt = "画像のテキストを正確に抽出し、Markdown形式で整えてください。数式はLaTeXを使用してください。"

    # 3. 入力画像の確定
    target_image_path = ""
    if args.input.lower() == "clip":
        target_image_path = get_image_from_clipboard()
        if not target_image_path:
            sys.exit(1)
    else:
        if os.path.exists(args.input):
            target_image_path = args.input
            print(f"入力ファイルを確認: {target_image_path}")
        else:
            print(f"❌ ファイルが見つかりません: {args.input}")
            sys.exit(1)

    # 4. AI実行
    try:
        result_text = call_ai_ocr(prompt, target_image_path, api=args.api, model=args.model)
        
        # 5. 結果保存
        Path(args.output).write_text(result_text.strip(), encoding="utf-8")
        print(f"✅ 成功: {args.output} に書き込みました。")
        
    except Exception as e:
        print(f"❌ AI呼び出し中にエラーが発生しました: {e}")

if __name__ == "__main__":
    main()