#!/usr/bin/env python3
"""Generate illustration assets through the LeapSkills Runtime API."""

from __future__ import annotations

import argparse
import concurrent.futures
import json
import sys
import time
from pathlib import Path
from typing import Any

from leapskills_runtime import (
    IMAGE_GENERATIONS_PATH,
    LeapSkillsRuntimeError,
    call_image_generation,
    resolve_leapskills_origin,
    runtime_url,
    write_images,
)

SKILL_ROOT = Path(__file__).resolve().parents[1]
SKILL_SLUG = "design"
DEFAULT_MODEL = "image-gen/gemini-3.1-flash-image-preview"
DEFAULT_SIZE = "1024x1024"
DEFAULT_RESPONSE_FORMAT = "url"
DEFAULT_WORKERS = 3

def fail(message: str, code: int = 1) -> None:
    print(json.dumps({"ok": False, "error": message}, ensure_ascii=False, indent=2), file=sys.stderr)
    raise SystemExit(code)

def allowed_models() -> list[str]:
    meta_path = SKILL_ROOT / "_meta.json"
    try:
        meta = json.loads(meta_path.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError):
        return [DEFAULT_MODEL]
    models = meta.get("requires", {}).get("models", [])
    return [str(model) for model in models] or [DEFAULT_MODEL]

def build_payload(prompt: str, model: str, size: str) -> dict[str, Any]:
    models = allowed_models()
    if model not in models:
        fail(f"Model {model!r} is not declared in _meta.json.requires.models.")
    return {
        "model": model,
        "prompt": prompt,
        "size": size,
        "n": 1,
        "response_format": DEFAULT_RESPONSE_FORMAT,
    }

def load_batch_jobs(batch_path: Path) -> list[dict[str, Any]]:
    manifest = json.loads(batch_path.read_text(encoding="utf-8"))
    jobs = manifest.get("jobs") if isinstance(manifest, dict) else manifest
    if not isinstance(jobs, list) or not jobs:
        fail("Batch manifest must be a non-empty JSON array or an object with jobs array.")
    normalized: list[dict[str, Any]] = []
    for index, job in enumerate(jobs, start=1):
        if not isinstance(job, dict):
            fail(f"Batch job {index} must be an object.")
        if not job.get("prompt") or not job.get("out"):
            fail(f"Batch job {index} must include prompt and out.")
        if job.get("reference") or job.get("references"):
            fail("Reference-image generation is blocked by the current LeapSkills API surface.")
        normalized.append(job)
    return normalized

def generate_one(args: argparse.Namespace, prompt: str, out: Path, size: str | None = None) -> dict[str, Any]:
    payload = build_payload(prompt=prompt, model=args.model, size=size or args.size)
    runtime_origin = resolve_leapskills_origin(args.runtime_origin)
    endpoint = runtime_url(IMAGE_GENERATIONS_PATH, runtime_origin)
    if args.dry_run:
        return {
            "ok": True,
            "dry_run": True,
            "endpoint": endpoint,
            "endpoint_path": IMAGE_GENERATIONS_PATH,
            "runtime_origin": runtime_origin,
            "skill_slug": SKILL_SLUG,
            "model": payload["model"],
            "payload": payload,
        }
    started = time.time()
    try:
        response = call_image_generation(
            skill_slug=SKILL_SLUG,
            model=payload["model"],
            prompt=payload["prompt"],
            size=payload["size"],
            response_format=payload["response_format"],
            runtime_origin=runtime_origin,
            timeout=args.timeout,
        )
        paths = write_images(response, out, args.timeout)
    except LeapSkillsRuntimeError as exc:
        fail(str(exc))
    return {
        "ok": True,
        "endpoint_path": IMAGE_GENERATIONS_PATH,
        "runtime_origin": runtime_origin,
        "skill_slug": SKILL_SLUG,
        "model": payload["model"],
        "prompt": prompt,
        "paths": paths,
        "elapsed_seconds": round(time.time() - started, 2),
    }

def run_batch(args: argparse.Namespace) -> dict[str, Any]:
    jobs = load_batch_jobs(Path(args.batch).expanduser().resolve())
    workers = args.workers or DEFAULT_WORKERS
    if workers < 1:
        fail("--workers must be at least 1.")

    results: list[dict[str, Any] | None] = [None] * len(jobs)
    failures = 0

    def run_job(index: int, job: dict[str, Any]) -> dict[str, Any]:
        result = generate_one(
            args=args,
            prompt=str(job["prompt"]),
            out=Path(str(job["out"])).expanduser().resolve(),
            size=job.get("size") or args.size,
        )
        result["job_index"] = index + 1
        return result

    with concurrent.futures.ThreadPoolExecutor(max_workers=workers) as executor:
        future_to_index = {executor.submit(run_job, index, job): index for index, job in enumerate(jobs)}
        for future in concurrent.futures.as_completed(future_to_index):
            index = future_to_index[future]
            try:
                results[index] = future.result()
            except SystemExit:
                raise
            except Exception as exc:
                failures += 1
                results[index] = {"ok": False, "job_index": index + 1, "error": str(exc)}

    return {
        "ok": failures == 0,
        "endpoint_path": GATEWAY_PATH,
        "skill_slug": SKILL_SLUG,
        "job_count": len(jobs),
        "failure_count": failures,
        "results": results,
    }

def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Generate illustrations through the LeapSkills Model Gateway.")
    parser.add_argument("--prompt", help="Prompt used for single-image generation.")
    parser.add_argument("--out", help="Output image path for single-image generation.")
    parser.add_argument("--batch", help="JSON batch manifest with prompt/out jobs.")
    parser.add_argument("--workers", type=int, default=None, help="Concurrent workers for --batch.")
    parser.add_argument("--model", default=DEFAULT_MODEL, help="Declared model ID from _meta.json.requires.models.")
    parser.add_argument("--size", default=DEFAULT_SIZE, help="Gateway image size, default 1024x1024.")
    parser.add_argument("--runtime-origin", default="", help="LeapSkills Runtime origin override.")
    parser.add_argument("--timeout", type=int, default=180, help="HTTP timeout in seconds.")
    parser.add_argument("--dry-run", action="store_true", help="Print the gateway request without sending it.")
    parser.add_argument("--emit-json", action="store_true", help="Emit a machine-readable JSON summary.")
    parser.add_argument("--reference", action="append", default=[], help=argparse.SUPPRESS)
    parser.add_argument("--aspect-ratio", default=None, help=argparse.SUPPRESS)
    parser.add_argument("--image-size", default=None, help=argparse.SUPPRESS)
    parser.add_argument("--width", type=int, default=None, help=argparse.SUPPRESS)
    parser.add_argument("--height", type=int, default=None, help=argparse.SUPPRESS)
    return parser.parse_args()

def main() -> int:
    args = parse_args()
    if args.reference:
        fail("Reference-image generation is blocked by the current LeapSkills API surface.")
    if args.aspect_ratio or args.image_size or args.width or args.height:
        fail("Provider-specific aspect ratio and slot-size controls are not supported by the current LeapSkills API surface. Use --size only.")
    if args.batch and (args.prompt or args.out):
        fail("--batch cannot be combined with --prompt or --out.")
    if not args.batch and (not args.prompt or not args.out):
        fail("Single-image generation requires both --prompt and --out.")

    result = run_batch(args) if args.batch else generate_one(args, args.prompt, Path(args.out).expanduser().resolve())
    if args.emit_json:
        print(json.dumps(result, ensure_ascii=False, indent=2))
    elif "results" in result:
        for item in result["results"]:
            if item and item.get("ok"):
                print(f"Generated image: {', '.join(item.get('paths', []))}")
            elif item:
                print(f"Failed job {item.get('job_index')}: {item.get('error')}", file=sys.stderr)
    else:
        print(f"Generated image: {', '.join(result.get('paths', []))}")
    return 0 if result.get("ok") else 1

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