import gradio as gr
import os
import tempfile
import shutil
import torch
import numpy as np
from PIL import Image

print("Loading OVI model...")

device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Device: {device}")

model = None
tokenizer = None

try:
    from transformers import AutoTokenizer
    from huggingface_hub import hf_hub_download, snapshot_download
    import sys

    # Download full repo
    repo_path = snapshot_download("chetwinlow1/Ovi")
    sys.path.insert(0, repo_path)

    # Try importing model directly from repo
    try:
        from modeling_ovi import OviModel
        from processing_ovi import OviProcessor

        processor = OviProcessor.from_pretrained("chetwinlow1/Ovi")
        model = OviModel.from_pretrained(
            "chetwinlow1/Ovi",
            torch_dtype=torch.float16 if device == "cuda" else torch.float32,
        ).to(device)
        model.eval()
        print("✅ OVI loaded via custom classes!")

    except ImportError:
        # Fallback - try pipeline
        from transformers import pipeline
        pipe = pipeline(
            "image-to-video",
            model="chetwinlow1/Ovi",
            device=0 if device == "cuda" else -1,
        )
        model = pipe
        processor = None
        print("✅ OVI loaded via pipeline!")

except Exception as e:
    print(f"❌ Model error: {e}")
    model = None
    processor = None


def generate_video(image, prompt):
    if image is None:
        raise gr.Error("Please upload an image!")
    if not prompt or prompt.strip() == "":
        raise gr.Error("Please enter a text prompt!")
    if model is None:
        raise gr.Error("Model not loaded!")

    try:
        if isinstance(image, str):
            pil_image = Image.open(image).convert("RGB")
        else:
            pil_image = Image.fromarray(image).convert("RGB")

        output_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name

        if processor is not None:
            # Custom processor path
            inputs = processor(
                text=prompt,
                images=pil_image,
                return_tensors="pt"
            ).to(device)
            with torch.no_grad():
                outputs = model.generate(**inputs)
        else:
            # Pipeline path
            outputs = model(pil_image, prompt)

        # Save video
        if hasattr(outputs, 'video'):
            video_frames = outputs.video[0].cpu().numpy()
            import cv2
            h, w = video_frames.shape[1:3]
            writer = cv2.VideoWriter(
                output_path,
                cv2.VideoWriter_fourcc(*'mp4v'),
                24, (w, h)
            )
            for frame in video_frames:
                frame_bgr = cv2.cvtColor(
                    (frame * 255).astype(np.uint8),
                    cv2.COLOR_RGB2BGR
                )
                writer.write(frame_bgr)
            writer.release()
        elif isinstance(outputs, str) and os.path.exists(outputs):
            shutil.copy(outputs, output_path)
        elif isinstance(outputs, list) and len(outputs) > 0:
            out = outputs[0]
            if isinstance(out, str) and os.path.exists(out):
                shutil.copy(out, output_path)
            elif hasattr(out, 'get'):
                v = out.get('video') or out.get('path')
                if v:
                    shutil.copy(v, output_path)

        return output_path

    except Exception as e:
        raise gr.Error(f"Error: {str(e)}")


with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
# 🎬 OVI — Talking Avatar Generator
**Free & Open Source** | No login required
""")
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(
                label="📸 Upload Image",
                type="filepath",
                height=300,
            )
            prompt_input = gr.Textbox(
                label="💬 Text Prompt",
                lines=3,
                placeholder="A person speaks. <S>Hello world!<E> <AUDCAP>Clear voice<ENDAUDCAP>",
            )
            generate_btn = gr.Button("🎬 Generate Video", variant="primary", size="lg")
            clear_btn = gr.Button("🗑️ Clear", variant="secondary")

        with gr.Column():
            video_output = gr.Video(
                label="🎥 Generated Video",
                height=300,
                autoplay=True,
            )

    gr.Markdown("""
### 💡 Tips:
- `<S>speech here<E>` — what avatar says
- `<AUDCAP>voice style<ENDAUDCAP>` — voice description
""")

    generate_btn.click(
        fn=generate_video,
        inputs=[image_input, prompt_input],
        outputs=[video_output],
    )
    clear_btn.click(
        fn=lambda: (None, "", None),
        outputs=[image_input, prompt_input, video_output],
    )

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