import os
import tempfile
import torch
import gradio as gr
import spaces
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
from diffusers import LTXPipeline, LTXVideoTransformer3DModel
from diffusers.utils import export_to_video

BASE_MODEL = "Lightricks/LTX-Video"
FINE_TUNED_MODEL = "DollasAndSpence/LTX2.3-ICEdit-Insight"
WEIGHT_NAME = "ltx-2.3-edit-insight-dev-fp8.safetensors"

HF_TOKEN = os.environ.get("HF_TOKEN")

print("Downloading checkpoint weights...")
ckpt_path = hf_hub_download(
    repo_id=FINE_TUNED_MODEL,
    filename=WEIGHT_NAME,
    token=HF_TOKEN,
)

print("Configuring transformer model...")
transformer_config = dict(
    LTXVideoTransformer3DModel.load_config(
        BASE_MODEL,
        subfolder="transformer",
        token=HF_TOKEN,
    )
)
# Adjust dimensions for LTX 2.3 / IC-Edit architecture
transformer_config["attention_head_dim"] = 128
transformer_config["caption_channels"] = 4096

# Instantiate model architecture on CPU
transformer = LTXVideoTransformer3DModel.from_config(transformer_config)

print("Loading state dict with shape-mismatch handling...")
state_dict = load_file(ckpt_path)

cleaned_state_dict = {}
for k, v in state_dict.items():
    new_key = k.replace("model.diffusion_model.", "").replace("diffusion_model.", "")
    cleaned_state_dict[new_key] = v

model_state_dict = transformer.state_dict()
loaded_state_dict = {}

for name, param in model_state_dict.items():
    if name in cleaned_state_dict:
        ckpt_param = cleaned_state_dict[name]
        if ckpt_param.shape == param.shape:
            loaded_state_dict[name] = ckpt_param
        else:
            # Handle shape mismatch gracefully (e.g. scale_shift_table or dimensions)
            min_dim0 = min(ckpt_param.shape[0], param.shape[0])
            if len(param.shape) == 2:
                min_dim1 = min(ckpt_param.shape[1], param.shape[1])
                new_tensor = param.clone()
                new_tensor[:min_dim0, :min_dim1] = ckpt_param[:min_dim0, :min_dim1]
                loaded_state_dict[name] = new_tensor
            else:
                new_tensor = param.clone()
                new_tensor[:min_dim0] = ckpt_param[:min_dim0]
                loaded_state_dict[name] = new_tensor
    else:
        # Keep default model parameter if not in checkpoint
        loaded_state_dict[name] = param

transformer.load_state_dict(loaded_state_dict, strict=False)
print("Transformer weights loaded successfully.")

transformer = transformer.to(dtype=torch.bfloat16)

print("Loading LTXPipeline...")
pipe = LTXPipeline.from_pretrained(
    BASE_MODEL,
    transformer=transformer,
    torch_dtype=torch.bfloat16,
    token=HF_TOKEN,
)

@spaces.GPU(duration=120)
def process_video(prompt, input_video_path, inference_steps, guidance_scale):
    if not input_video_path:
        raise gr.Error("Please upload an input video.")

    pipe.to("cuda")

    kwargs = {
        "prompt": prompt,
        "video": input_video_path,
        "num_inference_steps": inference_steps,
        "guidance_scale": guidance_scale,
        "output_type": "np",
    }

    try:
        output = pipe(**kwargs)
    except Exception as e:
        raise gr.Error(f"Inference failed: {str(e)}")

    out_file = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
    export_to_video(output.frames[0], out_file, fps=24)

    return out_file

with gr.Blocks(title="LTX2.3 ICEdit-Insight Video Restoration") as demo:
    gr.Markdown("# 🎥 LTX2.3 ICEdit-Insight Video Restoration")
    gr.Markdown("Upload an input video and provide an instruction prompt to restore or edit it.")
    
    with gr.Row():
        with gr.Column():
            input_vid = gr.Video(label="Input Video")
            prompt = gr.Textbox(
                label="Instruction Prompt",
                placeholder="e.g., Clean up AI inpainting glitches, remove artifacts...",
            )
            steps = gr.Slider(minimum=1, maximum=50, value=30, step=1, label="Inference Steps")
            cfg = gr.Slider(minimum=1.0, maximum=15.0, value=4.0, step=0.1, label="Guidance Scale")
            run_btn = gr.Button("Generate Fix", variant="primary")

        with gr.Column():
            output_vid = gr.Video(label="Restored Output Video")

    run_btn.click(
        fn=process_video,
        inputs=[prompt, input_vid, steps, cfg],
        outputs=output_vid,
    )

if __name__ == "__main__":
    demo.launch()