#!/usr/bin/env python
"""Run one local Action Viz generation without opening the browser UI."""

from __future__ import annotations

import argparse
import json
import os
import shutil
from pathlib import Path
from typing import Any

from cosmos_framework.data.vfm.action.action_viz.adapters import build_adapter, sample_action_to_numpy
from cosmos_framework.data.vfm.action.action_viz.local_worker import default_local_worker_from_env
from cosmos_framework.data.vfm.action.action_viz.state import (
    GenerationRequest,
    control_points_from_action,
    make_generation_id,
)
from cosmos_framework.data.vfm.action.urdf_visualizer.viewer import _build_datasets, _create_dataset


def main() -> None:
    _configure_cache_env()
    args = _parse_args()
    datasets = _build_datasets()
    if args.dataset not in datasets:
        raise ValueError(f"Unknown dataset {args.dataset!r}; expected one of {sorted(datasets)}")

    entry = datasets[args.dataset]
    sample_index = int(args.sample_index)
    if sample_index < 0:
        sample_index = int(entry.initial_index)

    dataset = _create_dataset(entry, int(args.chunk_length))
    sample = dataset[sample_index]
    adapter = build_adapter(args.dataset, entry)
    baked_action = sample_action_to_numpy(sample).astype("float32", copy=True)
    generation_id = args.generation_id or make_generation_id()
    generation_dir = Path(args.output_root) / generation_id
    if generation_dir.exists():
        shutil.rmtree(generation_dir)

    request = GenerationRequest(
        generation_id=generation_id,
        model_mode=args.model_mode,
        dataset=args.dataset,
        sample_index=sample_index,
        experiment_name="",
        s3_checkpoint_dir=args.checkpoint,
        checkpoint_cache_dir=None,
        output_dir=str(generation_dir),
        seed=int(args.seed),
        num_steps=int(args.num_steps),
        guidance=float(args.guidance),
        control_points=control_points_from_action(baked_action, baked_action.shape[1]),
        baked_action=baked_action.astype(float).tolist(),
        prompt_description=_extract_prompt_description(sample.get("ai_caption", "")),
        dataset_split="full",
        dataset_selector=adapter.dataset_selector,
        dataset_kwargs=entry.dataset_kwargs,
        use_torch_compile=False,
    )

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

    def _progress(percent: int, message: str) -> None:
        progress.append({"percent": int(percent), "message": str(message)})
        print(f"progress {percent:3d}% {message}", flush=True)

    worker = default_local_worker_from_env()
    try:
        result = worker.run(request, progress_callback=_progress, queue_callback=lambda state: print(f"queue {state}"))
    finally:
        worker.close()

    summary = {
        "status": result.status,
        "message": result.message,
        "generation_id": result.generation_id,
        "result_path": result.result_path,
        "video_path": result.video_path,
        "generated_action_path": result.generated_action_path,
        "progress": progress,
    }
    print(json.dumps(summary, indent=2, sort_keys=True))
    if result.status != "success":
        raise RuntimeError(f"Generation failed: {result.message}")
    if result.video_path is None or not Path(result.video_path).is_file():
        raise FileNotFoundError(f"Generation did not produce a video file: {result.video_path}")
    if args.model_mode == "policy" and (result.generated_action_path is None or not Path(result.generated_action_path).is_file()):
        raise FileNotFoundError(f"Policy generation did not produce an action file: {result.generated_action_path}")
    if not args.keep_output:
        shutil.rmtree(generation_dir, ignore_errors=True)


def _parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--dataset", default="bridge")
    parser.add_argument("--sample-index", type=int, default=-1)
    parser.add_argument("--chunk-length", type=int, default=16)
    parser.add_argument("--model-mode", choices=("forward_dynamics", "policy"), default="forward_dynamics")
    parser.add_argument("--checkpoint", default="nvidia/Cosmos3-Nano")
    parser.add_argument("--output-root", default="/tmp/action_viz_generation_smoke")
    parser.add_argument("--generation-id", default="")
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--num-steps", type=int, default=1)
    parser.add_argument("--guidance", type=float, default=1.0)
    parser.add_argument("--keep-output", action="store_true")
    return parser.parse_args()


def _configure_cache_env() -> None:
    app_data_root = Path(os.environ.get("ACTION_VIZ_APP_DATA_ROOT", "/app_data"))
    hf_home = Path(os.environ.setdefault("HF_HOME", str(app_data_root / "huggingface")))
    hf_hub_cache = Path(os.environ.setdefault("HF_HUB_CACHE", str(hf_home / "hub")))
    hf_home.mkdir(parents=True, exist_ok=True)
    hf_hub_cache.mkdir(parents=True, exist_ok=True)


def _extract_prompt_description(prompt: object) -> str:
    if isinstance(prompt, dict):
        prompt_obj = prompt
    elif isinstance(prompt, str):
        try:
            prompt_obj = json.loads(prompt)
        except json.JSONDecodeError:
            return prompt
    else:
        return ""
    value = prompt_obj.get("description", "")
    return value if isinstance(value, str) else ""


if __name__ == "__main__":
    main()