"""English ZeroGPU interface with separate CPU/GPU/CPU stages."""
import spaces  # Must precede any imports that can initialize CUDA.
import logging
import os

import gradio as gr

from cloud_backend import CloudEngine
from cloud_catalog import CloudCatalog
from studio_ui import build_studio, launch_style

logger = logging.getLogger(__name__)
ENGINE = CloudEngine(catalog=CloudCatalog())
ENGINE.eagerly_load(print)


def prepare_motion(source, style_name, reference_mode, reference_id, upload, caption, frames,
             seed, steps, text_weight, style_weight, style_guidance, content, progress=gr.Progress()):
    def report(message):
        progress(0, desc=message)
    try:
        if source == "Dataset reference":
            if reference_mode == "Specific ID" and not str(reference_id or "").strip():
                raise ValueError("Choose or enter a reference ID for the selected style.")
            if reference_mode not in {"Random reference", "Specific ID"}:
                raise ValueError("Choose a reference selection mode.")
            selected = ENGINE.catalog.select(style_name,
                reference_id if reference_mode == "Specific ID" else "", seed=int(seed))
            style_path, reference_info = selected.path, selected.metadata()
        elif source == "Upload .npy":
            style_path, reference_info = upload, None
        else:
            raise ValueError("Choose a style reference source.")
        prepared = ENGINE.prepare(
            style_path, caption, frames=frames, seed=seed, steps=steps,
            text_weight=text_weight, style_weight=style_weight,
            style_guidance=style_guidance, content_path=content, report=report,
            reference_info=reference_info,
        )
    except ValueError as error:
        raise gr.Error(str(error)) from error
    except Exception as error:
        logger.exception("Motion preparation failed")
        raise gr.Error(f"Reference preparation failed ({type(error).__name__}): {str(error)[:500]}") from error
    return prepared, "Reference validated. Waiting for ZeroGPU...", gr.Button(interactive=False)


@spaces.GPU(duration=60)
def generate_motion(prepared, progress=gr.Progress(track_tqdm=True)):
    """The registered callback performs actual CUDA sampling, never rendering."""
    def report(message):
        progress(0, desc=message)
    try:
        generated = ENGINE.infer(prepared, report=report)
    except ValueError as error:
        raise gr.Error(str(error)) from error
    except Exception as error:
        logger.exception("ZeroGPU motion inference failed")
        raise gr.Error(f"GPU inference failed ({type(error).__name__}): {str(error)[:500]}") from error
    return generated, "Inference finished. Rendering MP4 and BVH..."


def render_motion(generated, progress=gr.Progress()):
    def report(message):
        progress(0, desc=message)
    try:
        result = ENGINE.finish(generated, report=report)
    except ValueError as error:
        raise gr.Error(str(error)) from error
    except Exception as error:
        logger.exception("CPU motion export failed")
        raise gr.Error(f"Motion export failed ({type(error).__name__}): {str(error)[:500]}") from error
    info = result["info"]
    selected = info["style_reference"]
    reference_label = (f"{selected['style']} / {selected['id']}" if "id" in selected
                       else selected.get("filename", "uploaded reference"))
    message = (f"**Reference: {reference_label}**  \n"
               f"Motion: {info['frames'] / info['fps']:.1f} s · "
               f"Inference: {info['inference_seconds']:.2f} s · "
               f"Processing: {info['processing_seconds']:.2f} s · Seed: {info['seed']}  \n"
               "Generated with Hugging Face ZeroGPU.")
    return (result["video"], result["reference"], result["content_video"],
            [result["video"], result["bvh"], result["features"], result["motion"], result["metadata"]], message)


def build_demo():
    return build_studio(ENGINE.catalog, prepare_motion, generate_motion, render_motion)


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    build_demo().launch(server_name=os.environ.get("GRADIO_SERVER_NAME", "127.0.0.1"),
        server_port=int(os.environ.get("PORT", "7860")),
        share=False, max_file_size="20mb", show_error=False, **launch_style())