#!/usr/bin/env python3
"""Generate Wan2.2 I2V prompt plans for mechanism shot plans."""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from autoshorts.keyframe_prompt import write_i2v_prompt_plans
from autoshorts.mechanism_pipeline import build_mechanism_shot_plan
from autoshorts.retention_studio import PILOT_IDEAS, build_mechanism_explainer


def _default_shot_plans() -> list[dict]:
    return [build_mechanism_shot_plan(build_mechanism_explainer(PILOT_IDEAS[0]))]


def _load_shot_plans(path: Path) -> list[dict]:
    data = json.loads(path.read_text(encoding="utf-8"))
    if isinstance(data, list):
        return data
    if isinstance(data, dict) and "scenes" in data:
        return [data]
    if isinstance(data, dict) and "plans" in data:
        return data["plans"]
    raise ValueError(f"Unsupported shot plan JSON shape: {path}")


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--shot-plans-json",
        type=Path,
        help="Optional JSON shot-plan file. Defaults to the current mechanism pilot.",
    )
    parser.add_argument(
        "--output-dir",
        type=Path,
        default=Path("data/i2v_prompt_plans/mechanism_pilots"),
        help="Directory for generated I2V prompt plans.",
    )
    args = parser.parse_args()

    shot_plans = _load_shot_plans(args.shot_plans_json) if args.shot_plans_json else _default_shot_plans()
    paths = write_i2v_prompt_plans(shot_plans, args.output_dir)
    for path in paths:
        print(path)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())