import spaces
import gradio as gr
import sys
import time
import os
import random
from PIL import Image 
 # os.environ["CUDA_VISIBLE_DEVICES"] = ""

# Create the gr.State component *outside* the gr.Blocks context
predictor_state = gr.State(None)

def init_predictor(task_type: str):
    import torch
    from skyreelsinfer import TaskType
    from skyreelsinfer.offload import OffloadConfig
    from skyreelsinfer.skyreels_video_infer import SkyReelsVideoInfer
    from huggingface_hub.utils import RepositoryNotFoundError, RevisionNotFoundError, EntryNotFoundError

    try:
        predictor = SkyReelsVideoInfer(
            task_type=TaskType.I2V if task_type == "i2v" else TaskType.T2V,
            model_id="Skywork/skyreels-v1-Hunyuan-i2v",
            quant_model=True,
            is_offload=True,
            offload_config=OffloadConfig(
                high_cpu_memory=True,
                parameters_level=True,
            ),
            use_multiprocessing=False,
        )
        return "Model loaded successfully!", predictor  # Return predictor

    except (RepositoryNotFoundError, RevisionNotFoundError, EntryNotFoundError) as e:
        return f"Error: Model not found. Details: {e}", None
    except Exception as e:
        return f"Error loading model: {e}", None
        
init_predictor('i2v')

@spaces.GPU(duration=80)
def generate_video(prompt, image, predictor):
    from diffusers.utils import export_to_video
    from diffusers.utils import load_image
    import os

    if image == None:
        return "Error: For i2v, provide image path.", "{}"
    if not isinstance(prompt, str):
        return "Error: No prompt.", "{}"

    #if seed == -1:
    random.seed(time.time())
    seed = int(random.randrange(4294967294))

    kwargs = {
        "prompt": prompt,
        "height": 256,
        "width": 256,
        "num_frames": 24,
        "num_inference_steps": 30,
        "seed": int(seed),
        "guidance_scale": 7.0,
        "embedded_guidance_scale": 1.0,
        "negative_prompt": "bad quality, blur",
        "cfg_for": False,
    }

    kwargs["image"] = load_image(image=image)
    output = predictor.inference(kwargs)
    frames = output
    save_dir = f"./result/{task_type}"
    os.makedirs(save_dir, exist_ok=True)
    video_out_file = f"{save_dir}/{prompt[:100]}_{int(seed)}.mp4"
    print(f"Generating video: {video_out_file}")
    export_to_video(frames, video_out_file, fps=24)
    return video_out_file

 
def display_image(file):
    if file is not None:
        return Image.open(file.name)
    else:
        return None
        
# --- Minimal Gradio Interface ---
with gr.Blocks() as demo:

    image_file = gr.File(label="Image Prompt (Required)", file_types=["image"])
    image_file_preview = gr.Image(label="Image Prompt Preview", interactive=False)
    prompt_textbox = gr.Textbox(label="Prompt")
    generate_button = gr.Button("Generate")
    output_video = gr.Video(label="Output Video") # Just a textbox

    image_file.change(
        display_image,
        inputs=[image_file],
        outputs=[image_file_preview]
    )

    generate_button.click(
        fn=generate_video,
        inputs=[prompt_textbox, image_file, predictor_state],
        outputs=[output_video],
    )

demo.launch()