import argparse
import base64
import json
import os
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path
import re

import requests


def read_openrouter_api_key():
    env_key = os.getenv("OPENROUTER_API_KEY")
    if env_key:
        return env_key.strip()
    root = Path(__file__).resolve().parents[4]
    key_path = root / ".ai/skills/.config/openrouter_api_key.txt"
    if key_path.exists():
        return key_path.read_text(encoding="utf-8").strip()
    
    # 尝试读取通用配置文件
    common_key_path = root / ".ai/skills/invoking-llm/.config/openrouter_api_key.txt"
    if common_key_path.exists():
        return common_key_path.read_text(encoding="utf-8").strip()
        
    return ""


def parse_size(value):
    parts = value.lower().split("x")
    if len(parts) != 2:
        raise ValueError("size must be WIDTHxHEIGHT")
    return int(parts[0]), int(parts[1])


def build_prompt(narration, ratio, width, height):
    narration_text = narration or ""
    return (
        "现代科技风PPT，用于口播视频中的一个展示，"
        f"PPT的内容是讲解我正在念的解说词：“{narration_text}”。"
        f"美观漂亮PPT，长宽比为{ratio}，分辨率{width}x{height}。"
        "禁止出现任何小字，只使用简单大字、结构线条和背景图、示意图等。"
    )


def extract_image_data_url(message):
    images = message.get("images") or []
    if images:
        return images[0]["image_url"]["url"]
    content = message.get("content")
    if isinstance(content, str):
        match = re.search(r"data:image\/[a-zA-Z0-9.+-]+;base64,[A-Za-z0-9+/=]+", content)
        if match:
            return match.group(0)
    if isinstance(content, list):
        for item in content:
            if item.get("type") == "image_url":
                return item["image_url"]["url"]
    raise ValueError("no image url found in response")


def decode_data_url(data_url):
    if not data_url.startswith("data:"):
        raise ValueError("image url is not a data url")
    b64_data = data_url.split(",", 1)[1]
    return base64.b64decode(b64_data)


def generate_one(shot, output_dir, ratio, width, height, model, api_key, base_url, referer, title):
    index = shot.get("index")
    if not index:
        raise ValueError("shot is missing index")
    narration = shot.get("text") or shot.get("image_prompt") or ""
    prompt = build_prompt(narration, ratio, width, height)
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
        "HTTP-Referer": referer,
        "X-Title": title,
    }
    payload = {
        "model": model,
        "messages": [{"role": "user", "content": prompt}],
        "modalities": ["image", "text"],
    }
    last_error = None
    for _ in range(3):
        response = requests.post(base_url, headers=headers, json=payload, timeout=180)
        response.raise_for_status()
        result = response.json()
        message = result["choices"][0]["message"]
        try:
            data_url = extract_image_data_url(message)
            image_bytes = decode_data_url(data_url)
            output_path = output_dir / f"shot_{int(index):03d}.png"
            output_path.write_bytes(image_bytes)
            return output_path
        except Exception as exc:
            last_error = exc
    raise last_error


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--shots", required=True)
    parser.add_argument("--out", required=True)
    parser.add_argument("--ratio", default="16:9")
    parser.add_argument("--size", default="1920x1080")
    parser.add_argument("--workers", type=int, default=4)
    parser.add_argument("--model", default="google/gemini-3-pro-image-preview")
    parser.add_argument("--base-url", default="https://openrouter.ai/api/v1/chat/completions")
    parser.add_argument("--referer", default="https://github.com/cgerike/aihot")
    parser.add_argument("--title", default="AIHot Scraper")
    args = parser.parse_args()

    api_key = read_openrouter_api_key()
    if not api_key:
        raise SystemExit("missing OPENROUTER api key")

    width, height = parse_size(args.size)
    shots_path = Path(args.shots)
    output_dir = Path(args.out)
    output_dir.mkdir(parents=True, exist_ok=True)

    shots = json.loads(shots_path.read_text(encoding="utf-8"))
    with ThreadPoolExecutor(max_workers=max(1, args.workers)) as executor:
        futures = [
            executor.submit(
                generate_one,
                shot,
                output_dir,
                args.ratio,
                width,
                height,
                args.model,
                api_key,
                args.base_url,
                args.referer,
                args.title,
            )
            for shot in shots
        ]
        for future in as_completed(futures):
            future.result()


if __name__ == "__main__":
    main()
