import os
import sys
import argparse
import base64
import mimetypes
from pathlib import Path
from PIL import Image, ImageGrab

# 添付ライブラリの読み込み
try:
    import tkai_lib
    from get_from_ai import read_ini
except ImportError:
    print("Error: 'tkai_lib.py' または 'get_from_ai.py' が見つかりません。")
    sys.exit(1)

# LiteLLMの読み込み
try:
    from litellm import completion, supports_vision
except ImportError:
    print("Error: 'litellm' がインストールされていません。")
    print("Install: pip install litellm")
    sys.exit(1)

# OCR_general.py の MODEL_MAP を継承
MODEL_MAP = {
    "openai": "openai/gpt-5.5",
    "gpt": "openai/gpt-5.5",
    "gemini": "gemini/gemini-3.1-pro-preview",
    "claude": "anthropic/claude-sonnet-4-5",
}

def image_to_data_url(path: Path) -> str:
    """画像をBase64 Data URLに変換（LiteLLM用）"""
    mime, _ = mimetypes.guess_type(str(path))
    if mime is None: mime = "image/png"
    b64 = base64.b64encode(path.read_bytes()).decode("utf-8")
    return f"data:{mime};base64,{b64}"

def get_image_from_clipboard():
    """Windowsクリップボードから画像を取得"""
    print("--- クリップボードを確認中 ---")
    img = ImageGrab.grabclipboard()
    if isinstance(img, Image.Image):
        target_path = Path("clipboard_ocr_input.png")
        img.save(target_path, "PNG")
        print(f"✅ 画像を取得しました: {target_path}")
        return target_path
    print("❌ クリップボードに画像が見つかりませんでした。")
    return None

def main():
    # 1. 環境設定の読み込み (tkai_libの活用)
    tkai_lib.read_ai_config("ai.env")

    parser = argparse.ArgumentParser(description="Universal AI OCR Tool (LiteLLM Edition)")
    parser.add_argument("input", help="画像パス、または 'clip'")
    parser.add_argument("--output", "-o", default=None, help="出力Markdownファイル名")
    parser.add_argument("--provider", "-p", default="gemini", choices=MODEL_MAP.keys(), help="AIプロバイダー")
    parser.add_argument("--model", "-m", default=None, help="LiteLLMモデル名を直接指定")
    parser.add_argument("--ini", "-i", default="ai_ocr2md.ini", help="プロンプト設定ファイル")
    args = parser.parse_args()

    # モデルの決定
    selected_model = args.model if args.model else MODEL_MAP[args.provider]
    
    # 2. 出力ファイル名の決定
    input_path = Path(args.input)
    output_path = Path(args.output) if args.output else Path(f"{input_path.stem}.md")

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

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

    # 5. LiteLLMによるOCR実行
    print(f"🚀 {selected_model} でOCR実行中...")
    try:
        data_url = image_to_data_url(img_file)
        
        # LiteLLMの呼び出し
        response = completion(
            model=selected_model,
            messages=[{
                "role": "user",
                "content": [
                    {"type": "text", "text": prompt},
                    {"type": "image_url", "image_url": {"url": data_url}}
                ]
            }]
        )
        result_text = response.choices[0].message.content

        # 6. 結果保存
        output_path.write_text(result_text.strip(), encoding="utf-8")
        print(f"✅ 成功: {output_path} に書き込みました。")

    except Exception as e:
        print(f"❌ エラーが発生しました: {e}")

if __name__ == "__main__":
    main()