"""
Submit a generation job to a LingBot-World server.

Sends a POST to /generate, then polls /status/<id> until the video is ready.

Usage:
    python caller.py --url http://localhost:8080 \\
        --image examples/00/image.jpg \\
        --prompt "A woman walks forward slowly" \\
        --action_dir examples/00/

    # Through a RunPod proxy:
    python caller.py --url https://YOUR-POD-8080.proxy.runpod.net \\
        --image photo.jpg --prompt "..."
"""

import argparse
import base64
import json
import sys
import time
from pathlib import Path

import requests


def encode_file_b64(path: str) -> str:
    with open(path, "rb") as f:
        return base64.b64encode(f.read()).decode("utf-8")


def build_payload(args) -> dict:
    payload = {
        "image": encode_file_b64(args.image),
        "prompt": args.prompt,
        "frame_num": args.frame_num,
        "size": args.size,
        "seed": args.seed,
        "guide_scale": args.guide_scale,
        "sampling_steps": args.sampling_steps,
    }

    if args.action_dir:
        action_dir = Path(args.action_dir)
        poses_path = action_dir / "poses.npy"
        intrinsics_path = action_dir / "intrinsics.npy"

        if not poses_path.exists() or not intrinsics_path.exists():
            print(f"ERROR: Expected poses.npy and intrinsics.npy in {action_dir}")
            sys.exit(1)

        payload["action_poses"] = encode_file_b64(str(poses_path))
        payload["action_intrinsics"] = encode_file_b64(str(intrinsics_path))
        print(f"Camera actions loaded from {action_dir}")

    return payload


def submit(url: str, payload: dict, poll_interval: int = 10, output_path: str | None = None):
    """POST the job, then poll until completion."""
    submit_url = f"{url.rstrip('/')}/generate"

    print(f"Submitting to {submit_url} ...")
    resp = requests.post(submit_url, json=payload, timeout=30)

    if resp.status_code not in (200, 202):
        print(f"ERROR {resp.status_code}: {resp.text[:500]}")
        sys.exit(1)

    result = resp.json()
    if "error" in result:
        print(f"ERROR: {result['error']}")
        sys.exit(1)

    job_id = result.get("id")
    print(f"Job accepted — ID: {job_id}")
    print("Polling for completion ...")

    status_url = f"{url.rstrip('/')}/status/{job_id}"
    while True:
        time.sleep(poll_interval)
        try:
            resp = requests.get(status_url, timeout=15)
            data = resp.json()
        except Exception as e:
            print(f"  Poll error (will retry): {e}")
            continue

        status = data.get("status")
        print(f"  Status: {status}")

        if status == "COMPLETED":
            output = data.get("output", {})
            print(f"\nDone!")
            if output.get("video_url"):
                print(f"  Video URL : {output['video_url']}")
            if output.get("video_path"):
                print(f"  Video path: {output['video_path']}")
                # Auto-download the video from the server
                download_url = f"{url.rstrip('/')}/download/{job_id}"
                dest = output_path or f"{job_id}.mp4"
                print(f"  Downloading to {dest} ...")
                try:
                    dl = requests.get(download_url, stream=True, timeout=300)
                    dl.raise_for_status()
                    with open(dest, "wb") as f:
                        for chunk in dl.iter_content(chunk_size=1024 * 1024):
                            f.write(chunk)
                    print(f"  Saved to {dest}")
                except Exception as e:
                    print(f"  Download failed: {e}")
            print(f"  Seed      : {output.get('seed')}")
            print(f"  Duration  : {output.get('duration_sec', 0):.1f}s")
            return
        elif status == "FAILED":
            output = data.get("output", {})
            print(f"\nJob failed: {output.get('error', 'Unknown error')}")
            sys.exit(1)


def main():
    parser = argparse.ArgumentParser(
        description="Submit a video generation job to a LingBot-World server"
    )
    parser.add_argument("--url", required=True,
                        help="Server URL (e.g. http://localhost:8080 or https://xyz-8080.proxy.runpod.net)")
    parser.add_argument("--image", required=True, help="Path to input image (JPEG/PNG)")
    parser.add_argument("--prompt", required=True, help="Text prompt")
    parser.add_argument("--action_dir", default=None,
                        help="Directory containing poses.npy and intrinsics.npy")
    parser.add_argument("--frame_num", type=int, default=81, help="Number of frames (default: 81)")
    parser.add_argument("--size", default="480*832", help="Output size: 480*832 or 720*1280")
    parser.add_argument("--seed", type=int, default=-1, help="RNG seed (-1 for random)")
    parser.add_argument("--guide_scale", type=float, default=5.0)
    parser.add_argument("--sampling_steps", type=int, default=40)
    parser.add_argument("--poll_interval", type=int, default=10,
                        help="Seconds between status polls (default: 10)")
    parser.add_argument("--output", default=None,
                        help="Local path to save downloaded video (default: <job_id>.mp4)")

    args = parser.parse_args()

    print(f"  Prompt: {args.prompt[:80]}...")
    print(f"  Frames: {args.frame_num}, Size: {args.size}, Steps: {args.sampling_steps}")

    payload = build_payload(args)
    submit(args.url, payload, poll_interval=args.poll_interval, output_path=args.output)


if __name__ == "__main__":
    main()