#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
tkai_movie.py

AI動画生成・講義動画支援向け 共通ライブラリ。

主な機能:
- OpenAI / Gemini API 呼び出し
- OpenAI Responses API 用 input_file アップロード
- make_textbook5.ini 互換の key=val / 三連引用符 ini 読み込み
- プロンプトテンプレート展開
- messages / raw response ログ保存
- [SCRIPT_START] ... [SCRIPT_END] などのブロック抽出
"""

from __future__ import annotations

import os
import re
import sys
import json
from pathlib import Path
from dataclasses import dataclass, field
from typing import Optional, Iterable, Any


try:
    import google.generativeai as genai
    from openai import OpenAI
except ImportError:
    print("必要なライブラリがインストールされていません。", file=sys.stderr)
    print("pip install google-generativeai openai python-dotenv", file=sys.stderr)
    input("\nPress ENTER to terminate>>\n")
    raise


DEFAULT_GEMINI_MODEL = "gemini-3.1-pro-preview"
DEFAULT_OPENAI_MODEL = "gpt-4o"
DEFAULT_OPENAI_MODEL5 = "gpt-5.4"


@dataclass
class AIConfig:
    api: str = "openai5"

    openai_key: Optional[str] = None
    gemini_key: Optional[str] = None

    openai_model: str = DEFAULT_OPENAI_MODEL
    openai_model5: str = DEFAULT_OPENAI_MODEL5
    gemini_model: str = DEFAULT_GEMINI_MODEL

    pause: int = 0

    @classmethod
    def from_env(cls, api: str = "openai5") -> "AIConfig":
        return cls(
            api=api,
            openai_key=os.getenv("OPENAI_API_KEY"),
            gemini_key=os.getenv("GOOGLE_API_KEY"),
            openai_model=os.getenv("OPENAI_MODEL", DEFAULT_OPENAI_MODEL),
            openai_model5=os.getenv("OPENAI_MODEL5", DEFAULT_OPENAI_MODEL5),
            gemini_model=os.getenv("GOOGLE_MODEL", DEFAULT_GEMINI_MODEL),
        )


@dataclass
class AIRequest:
    system_prompt: str
    text_inputs: dict[str, str] = field(default_factory=dict)
    file_inputs: list[str] = field(default_factory=list)
    output_blocks: list[str] = field(default_factory=lambda: ["SCRIPT"])
    final_instruction: Optional[str] = None


@dataclass
class AIResult:
    raw_text: str
    blocks: dict[str, str]

    def get_block(self, name: str) -> str:
        return self.blocks.get(name, "")

    def write_block(self, name: str, outfile: str | os.PathLike) -> None:
        text = self.get_block(name)
        Path(outfile).write_text(text, encoding="utf-8")
        print(f"✅ wrote {name}: {outfile}")


def search_file(
    infile: str | os.PathLike | None = None,
    *,
    default_name: str,
    script_file: str | os.PathLike | None = None,
) -> str | None:
    """
    ファイルを検索する。

    検索順序:
    - infile が None:
        1. カレントディレクトリ/default_name
        2. script_file と同じディレクトリ/default_name
    - infile が指定あり:
        1. 指定パスそのもの
        2. script_file と同じディレクトリ/infile
    """
    if script_file is None:
        script_dir = Path.cwd()
    else:
        script_dir = Path(script_file).expanduser().resolve().parent

    if infile is None:
        candidates = [
            Path.cwd() / default_name,
            script_dir / default_name,
        ]
    else:
        inpath = Path(infile).expanduser()
        candidates = [
            inpath if inpath.is_absolute() else Path.cwd() / inpath,
            script_dir / inpath,
        ]

    for candidate in candidates:
        if candidate.is_file():
            return str(candidate.resolve())

    return None


def read_ini_compat(inifile: str | os.PathLike) -> dict[str, str]:
    """
    make_textbook5.py互換の簡易iniを読み込む。

    対応:
    - key = value
    - # / ; コメント
    - 空行無視
    - 三連引用符による複数行値
    - 単一行三連引用符
    - $KEY 形式の変数展開
    """
    path = Path(inifile)
    if not path.is_file():
        raise FileNotFoundError(f"INI file not found: {path}")

    result: dict[str, str] = {}
    variables: dict[str, str] = {}

    current_key: str | None = None
    multiline_val: list[str] = []
    multiline_delim: str | None = None

    with path.open("r", encoding="utf-8") as f:
        for raw_line in f:
            line = raw_line.rstrip("\n").rstrip("\r")

            if multiline_delim:
                if line.strip() == multiline_delim:
                    val = "\n".join(multiline_val)
                    if current_key is not None:
                        result[current_key] = val
                        variables[current_key] = val
                    current_key = None
                    multiline_val = []
                    multiline_delim = None
                else:
                    multiline_val.append(line)
                continue

            stripped = line.strip()

            if not stripped or stripped.startswith("#") or stripped.startswith(";"):
                continue

            if "=" not in line:
                continue

            key, val = map(str.strip, line.split("=", 1))
            val = val.strip()

            if (
                val == '\"\"\"'
                or val == "'''"
                or (val.startswith('\"\"\"') and not val.endswith('\"\"\"'))
                or (val.startswith("'''") and not val.endswith("'''"))
            ):
                multiline_delim = val[:3]
                content = val[3:]
                multiline_val = [content] if content else []
                current_key = key
                continue

            if (val.startswith('\"\"\"') and val.endswith('\"\"\"')) or (
                val.startswith("'''") and val.endswith("'''")
            ):
                val = val[3:-3]

            result[key] = val
            variables[key] = val

    if multiline_delim:
        raise ValueError(f"Unclosed multiline value for key={current_key}")

    for key, val in list(result.items()):
        def expand_var(match: re.Match[str]) -> str:
            var_name = match.group(1)
            return variables.get(var_name, match.group(0))

        result[key] = re.sub(r"\$(\w+)\b", expand_var, val)

    return result


def load_prompt_ini(
    ini: str | os.PathLike | None,
    *,
    default_name: str,
    script_file: str | os.PathLike,
    required: bool = False,
) -> tuple[dict[str, str], str | None]:
    path = search_file(ini, default_name=default_name, script_file=script_file)

    if path is None:
        if required:
            raise FileNotFoundError(f"INI file not found: {ini or default_name}")
        return {}, None

    return read_ini_compat(path), path


def get_prompt_template(
    ini_data: dict[str, str],
    *,
    lang: str,
    default_template: str,
) -> str:
    lang_key = lang.lower()

    if lang_key in ("en",):
        return ini_data.get("PROMPT_TEMPLATE_EN") or ini_data.get("PROMPT_TEMPLATE_JA") or default_template

    if lang_key in ("zh", "cn"):
        return ini_data.get("PROMPT_TEMPLATE_ZH") or ini_data.get("PROMPT_TEMPLATE_JA") or default_template

    if lang_key in ("ko", "kr"):
        return ini_data.get("PROMPT_TEMPLATE_KO") or ini_data.get("PROMPT_TEMPLATE_JA") or default_template

    return ini_data.get("PROMPT_TEMPLATE_JA") or default_template


class _SafeFormatDict(dict):
    def __missing__(self, key: str) -> str:
        return "{" + key + "}"


def render_template(template: str, **kwargs: Any) -> str:
    return template.format_map(_SafeFormatDict(**kwargs)).strip()


def save_json_log(data: Any, logfile: str | os.PathLike) -> None:
    try:
        with open(logfile, "w", encoding="utf-8") as f:
            json.dump(data, f, ensure_ascii=False, indent=2)
        print(f"📝 log saved: {logfile}")
    except Exception as e:
        print(f"⚠ log save error: {e}", file=sys.stderr)


def build_final_instruction(output_blocks: Iterable[str]) -> str:
    parts = []

    for block in output_blocks:
        parts.append(
            f"""[{block}_START]
({block} content)
[{block}_END]
"""
        )

    return "以下の形式で出力してください:\n\n" + "\n".join(parts)


def parse_output_blocks(text: str, output_blocks: Iterable[str]) -> dict[str, str]:
    results: dict[str, str] = {}

    for block in output_blocks:
        pattern = rf"\[{re.escape(block)}_START\](.*?)\[{re.escape(block)}_END\]"
        m = re.search(pattern, text, re.DOTALL)
        results[block] = m.group(1).strip() if m else ""

    return results


def upload_openai_file(client: OpenAI, filepath: str | os.PathLike) -> str:
    path = Path(filepath)
    print(f"📤 Uploading file: {path}")

    with open(path, "rb") as f:
        uploaded = client.files.create(file=f, purpose="user_data")

    print(f"  file_id = {uploaded.id}")
    return uploaded.id


def call_openai_responses_api(request: AIRequest, config: AIConfig) -> str:
    if not config.openai_key:
        raise ValueError("OPENAI_API_KEY is not set")

    client = OpenAI(api_key=config.openai_key)

    input_items: list[dict[str, Any]] = []

    input_items.append({
        "role": "system",
        "content": request.system_prompt,
    })

    for title, text in request.text_inputs.items():
        input_items.append({
            "role": "user",
            "content": f"# {title}\n\n{text}",
        })

    for filepath in request.file_inputs:
        file_id = upload_openai_file(client, filepath)
        input_items.append({
            "role": "user",
            "content": [
                {
                    "type": "input_file",
                    "file_id": file_id,
                }
            ],
        })

    final_instruction = request.final_instruction
    if final_instruction is None:
        final_instruction = build_final_instruction(request.output_blocks)

    input_items.append({
        "role": "user",
        "content": final_instruction,
    })

    print()
    print("🚀 OpenAI Responses API")
    print(f"  model = {config.openai_model5}")
    print()

    response = client.responses.create(
        model=config.openai_model5,
        input=input_items,
    )

    return response.output_text or ""


def call_gemini_api(request: AIRequest, config: AIConfig) -> str:
    if not config.gemini_key:
        raise ValueError("GOOGLE_API_KEY is not set")

    genai.configure(api_key=config.gemini_key)
    model = genai.GenerativeModel(config.gemini_model)

    parts: list[str] = [request.system_prompt]

    for title, text in request.text_inputs.items():
        parts.append(f"# {title}\n\n{text}")

    for filepath in request.file_inputs:
        parts.append(f"[FILE]\n{filepath}")

    final_instruction = request.final_instruction
    if final_instruction is None:
        final_instruction = build_final_instruction(request.output_blocks)

    parts.append(final_instruction)

    content = "\n\n".join(parts)

    print()
    print("🚀 Gemini API")
    print(f"  model = {config.gemini_model}")
    print()

    response = model.generate_content(content)
    return response.text


def run_ai_generation(
    request: AIRequest,
    config: AIConfig,
    *,
    save_log: bool = True,
    logfile: Optional[str] = None,
) -> AIResult:
    if config.api in ("openai5",):
        raw = call_openai_responses_api(request, config)
    elif config.api in ("gemini", "google"):
        raw = call_gemini_api(request, config)
    else:
        raise ValueError(f"unsupported api: {config.api}")

    blocks = parse_output_blocks(raw, request.output_blocks)

    if save_log:
        if logfile is None:
            logfile = "ai_messages.log"

        logdata = {
            "system_prompt": request.system_prompt,
            "text_inputs": request.text_inputs,
            "file_inputs": request.file_inputs,
            "output_blocks": request.output_blocks,
            "raw_response": raw,
        }

        save_json_log(logdata, logfile)

    return AIResult(raw_text=raw, blocks=blocks)