"""
OmniLottie Gradio Demo - Hugging Face Compatible Version

This version uses decoder_hf.py with from_pretrained() to load models,
supporting automatic downloading from Hugging Face Hub.

Environment Variables:
    MODEL_PATH: Model path or HF Hub ID (default: "OmniLottie/OmniLottie")
    PROCESSOR_PATH: Processor path (default: "Qwen/Qwen2.5-VL-3B-Instruct")

Usage:
    # Using HF Hub model (automatic download)
    python app_hf.py

    # Using local model
    MODEL_PATH=/path/to/local/model python app_hf.py

    # Custom configuration
    MODEL_PATH=your-org/your-model PROCESSOR_PATH=/path/to/processor python app_hf.py
"""

import gradio as gr
import json
import torch
import os
import numpy as np
import random
import re
import tempfile
import base64
import threading
import time
from PIL import Image as PILImage
from decord import VideoReader, cpu

from decoder_hf import LottieDecoder  # Use HF decoder with from_pretrained
from transformers import AutoProcessor
from qwen_vl_utils import process_vision_info
from lottie.objects.lottie_tokenize import LottieTensor
from lottie.objects.lottie_param import (
    from_sequence, ShapeLayer, NullLayer, PreCompLayer, TextLayer,
    SolidColorLayer, Font, Chars,
    shape_layer_to_json, null_layer_to_json, precomp_layer_to_json,
    text_layer_to_json, solid_layer_to_json, font_to_json, char_to_json
)

SYSTEM_PROMPT = "You are a Lottie animation expert."
VIDEO_PROMPT = "Turn this video into Lottie code."
LOTTIE_BOS = 192398
LOTTIE_EOS = 192399
PAD_TOKEN = 151643

model = None
processor = None
device = None

generation_lock = threading.Lock()

def load_model_once():
    """
    Load OmniLottie model using from_pretrained() for HuggingFace compatibility
    Supports automatic downloading from HF Hub or loading from local path
    """
    global model, processor, device

    if model is not None:
        return model, processor, device

    # Model path - can be HF Hub ID or local path
    # Examples: "OmniLottie/OmniLottie" or "/path/to/local/model"
    model_path = os.environ.get("MODEL_PATH", "OmniLottie/OmniLottie")

    device = torch.device("cuda:0" if torch.cuda.is_available() else "xpu:0" if torch.xpu.is_available() else "cpu")

    print(f"Loading model from {model_path}...")

    # Load model using from_pretrained (supports HF Hub and local paths)
    model = LottieDecoder.from_pretrained(
        model_path,
        torch_dtype=torch.bfloat16,
        trust_remote_code=True
    )

    model = model.to(device).eval()

    # Load processor
    processor_path = os.environ.get("PROCESSOR_PATH", "Qwen/Qwen2.5-VL-3B-Instruct")
    processor = AutoProcessor.from_pretrained(
        processor_path,
        padding_side="left",
        trust_remote_code=True
    )

    print(f"✅ Model loaded on {device}")

    return model, processor, device

def simplify_to_animation_description(text):
    if not text or not isinstance(text, str):
        return text

    prefixes = [
        r'^The video features?\s+', r'^The scene shows?\s+',
        r'^An animation of\s+', r'^There is\s+', r'^It shows?\s+'
    ]

    for pattern in prefixes:
        text = re.sub(pattern, '', text, flags=re.IGNORECASE)

    if text:
        text = text[0].upper() + text[1:]

    return text.strip()

def add_random_background(img):
    if img.mode != 'RGBA':
        return img.convert('RGB')

    light_colors = [(255, 255, 255), (245, 245, 245), (250, 250, 250)]
    bg_color = random.choice(light_colors)
    background = PILImage.new('RGB', img.size, bg_color)
    background.paste(img, (0, 0), img)
    return background

def load_frames_from_video(video_path, num_frames=8, max_size=336):
    import os
    ext = os.path.splitext(video_path)[1].lower()

    frames = []

    if ext in ('.gif', '.webp'):
        try:
            img = PILImage.open(video_path)
            total_frames = getattr(img, 'n_frames', 1)

            if total_frames < 1:
                raise ValueError(f"No frames in {ext.upper()}: {video_path}")

            indices = np.linspace(0, total_frames - 1, min(num_frames, total_frames)).astype(int)

            for idx in indices:
                img.seek(idx)
                frame = img.convert('RGB')
                if max(frame.size) > max_size:
                    frame.thumbnail((max_size, max_size), PILImage.LANCZOS)
                frames.append(frame)

            img.close()

        except Exception as e:
            raise ValueError(f"Failed to load {ext.upper()}: {str(e)}")

    else:
        try:
            vr = VideoReader(video_path, ctx=cpu(0))
            total_frames = len(vr)

            if total_frames < 1:
                raise ValueError(f"Video has no frames: {video_path}")

            indices = np.linspace(0, total_frames - 1, num_frames).astype(int)
            frames_np = vr.get_batch(indices).asnumpy()

            for f in frames_np:
                img = PILImage.fromarray(f)
                if max(img.size) > max_size:
                    img.thumbnail((max_size, max_size), PILImage.LANCZOS)
                frames.append(img)

        except Exception as e:
            raise ValueError(f"Failed to load video: {str(e)}")

    while len(frames) < num_frames:
        frames.append(frames[-1].copy())

    return frames

def build_messages(task_type, text_prompt=None, image=None, video_frames=None):
    messages = [{"role": "system", "content": SYSTEM_PROMPT}]

    if task_type == "text":
        text = simplify_to_animation_description(text_prompt)
        messages.append({
            "role": "user",
            "content": [{"type": "text", "text": f"Generate Lottie code: {text}"}]
        })

    elif task_type == "image":
        text = simplify_to_animation_description(text_prompt)
        messages.append({
            "role": "user",
            "content": [
                {"type": "image", "image": image},
                {"type": "text", "text": f"Animate this image: {text}"}
            ]
        })

    elif task_type == "video":
        messages.append({
            "role": "user",
            "content": [
                {"type": "video", "video": video_frames, "fps": 8.0},
                {"type": "text", "text": VIDEO_PROMPT}
            ]
        })

    return messages

def prepare_inference_input(processor, messages, device):
    text_input = processor.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )

    image_inputs, video_inputs = process_vision_info(messages)

    inputs = processor(
        text=[text_input],
        images=image_inputs if image_inputs else None,
        videos=video_inputs if video_inputs else None,
        padding=False,
        return_tensors="pt"
    )

    input_ids = inputs['input_ids']
    attention_mask = inputs['attention_mask']
    target_len = 1500

    if input_ids.shape[1] < target_len:
        pad_len = target_len - input_ids.shape[1]
        input_ids = torch.cat([
            torch.full((1, pad_len), PAD_TOKEN, dtype=torch.long),
            input_ids
        ], dim=1)
        attention_mask = torch.cat([
            torch.zeros((1, pad_len), dtype=torch.long),
            attention_mask
        ], dim=1)

    result = {
        'input_ids': input_ids.to(device),
        'attention_mask': attention_mask.to(device),
        'pixel_values': inputs.get('pixel_values').to(device) if inputs.get('pixel_values') is not None else None,
        'image_grid_thw': inputs.get('image_grid_thw').to(device) if inputs.get('image_grid_thw') is not None else None,
        'pixel_values_videos': inputs.get('pixel_values_videos').to(device) if inputs.get('pixel_values_videos') is not None else None,
        'video_grid_thw': inputs.get('video_grid_thw').to(device) if inputs.get('video_grid_thw') is not None else None,
    }

    return result

def generate_lottie(model, inputs, max_tokens, device, use_sampling=False, temperature=0.95, top_p=0.25, top_k=5):
    """生成 Lottie tokens"""
    model.transformer.rope_deltas = None
    position_ids, _ = model.transformer.get_rope_index(
        input_ids=inputs['input_ids'],
        attention_mask=inputs['attention_mask'],
        image_grid_thw=inputs.get('image_grid_thw'),
        video_grid_thw=inputs.get('video_grid_thw'),
    )
    position_ids = position_ids * inputs['attention_mask'][None, ]

    kwargs = {
        'input_ids': inputs['input_ids'],
        'attention_mask': inputs['attention_mask'],
        'pixel_values': inputs.get('pixel_values'),
        'image_grid_thw': inputs.get('image_grid_thw'),
        'pixel_values_videos': inputs.get('pixel_values_videos'),
        'video_grid_thw': inputs.get('video_grid_thw'),
        'position_ids': position_ids,
        'max_new_tokens': max_tokens,
        'eos_token_id': LOTTIE_EOS,
        'pad_token_id': PAD_TOKEN,
        'use_cache': True,
    }

    if use_sampling:
        kwargs.update({'do_sample': True, 'temperature': temperature, 'top_p': top_p, 'top_k': top_k})
    else:
        kwargs.update({'do_sample': False, 'num_beams': 1})

    with torch.no_grad():
        outputs = model.transformer.generate(**kwargs)

    input_len = inputs['input_ids'].shape[1]
    generated_ids = outputs[0][input_len:].tolist()

    del outputs, kwargs, position_ids

    if generated_ids and generated_ids[0] == LOTTIE_BOS:
        generated_ids = generated_ids[1:]
    if LOTTIE_EOS in generated_ids:
        generated_ids = generated_ids[:generated_ids.index(LOTTIE_EOS)]

    return generated_ids

def fix_lottie_json(anim):

    anim_ip = int(round(anim.get("ip", 0)))
    anim_op = int(round(anim.get("op", 16)))
    anim["ip"]  = anim_ip
    anim["op"]  = anim_op
    anim["fr"]  = int(round(anim.get("fr", 8)))
    anim["ddd"] = int(anim.get("ddd", 0))

    def fix_t_recursive(obj):
        if isinstance(obj, dict):
            if obj.get("a") == 1 and isinstance(obj.get("k"), list):
                for kf in obj["k"]:
                    if isinstance(kf, dict) and "t" in kf:
                        kf["t"] = int(round(kf["t"]))
            for v in obj.values():
                fix_t_recursive(v)
        elif isinstance(obj, list):
            for item in obj:
                fix_t_recursive(item)

    fix_t_recursive(anim)

    max_x = float(anim.get("w", 512))
    max_y = float(anim.get("h", 512))

    def collect_pos(layer):
        nonlocal max_x, max_y
        p = layer.get("ks", {}).get("p", {})
        if isinstance(p, dict):
            if p.get("a", 0) == 0:
                pv = p.get("k", [0, 0])
                if isinstance(pv, list) and len(pv) >= 2:
                    max_x = max(max_x, float(pv[0]))
                    max_y = max(max_y, float(pv[1]))
            else:
                for kf in p.get("k", []):
                    if isinstance(kf, dict):
                        for sv in (kf.get("s", []), kf.get("e", [])):
                            if isinstance(sv, list) and len(sv) >= 2:
                                max_x = max(max_x, float(sv[0]))
                                max_y = max(max_y, float(sv[1]))
        for sub in layer.get("layers", []):
            collect_pos(sub)

    for layer in anim.get("layers", []):
        collect_pos(layer)

    anim["w"] = max(512, int((max_x * 1.1 + 15) // 16 * 16))
    anim["h"] = max(512, int((max_y * 1.1 + 15) // 16 * 16))

    valid_inds = set()
    for layer in anim.get("layers", []):
        if "ind" in layer:
            valid_inds.add(int(layer["ind"]))

    def clean_shapes(shapes):
        if not isinstance(shapes, list):
            return shapes
        cleaned = []
        for sh in shapes:
            if not isinstance(sh, dict):
                continue
            if sh.get("ty") == "gr":
                sh["it"] = clean_shapes(sh.get("it", []))
                if not sh["it"]:
                    continue
                has_tr = any(item.get("ty") == "tr" for item in sh["it"] if isinstance(item, dict))
                if not has_tr:
                    sh["it"].append({
                        "ty": "tr", "nm": "",
                        "a": {"a": 0, "k": [0, 0], "ix": 1},
                        "p": {"a": 0, "k": [0, 0], "ix": 2},
                        "s": {"a": 0, "k": [100, 100], "ix": 3},
                        "r": {"a": 0, "k": 0, "ix": 6},
                        "o": {"a": 0, "k": 100, "ix": 7},
                        "sk": {"a": 0, "k": 0, "ix": 4},
                        "sa": {"a": 0, "k": 0, "ix": 5},
                        "hd": False
                    })
            cleaned.append(sh)
        return cleaned

    def fix_layer(layer):
        ip = int(round(layer.get("ip", anim_ip)))
        op = int(round(layer.get("op", anim_op)))
        layer["ip"] = max(anim_ip, ip)
        layer["op"] = min(anim_op, max(layer["ip"] + 1, op))
        layer["st"] = int(round(layer.get("st", anim_ip)))
        if "ind" in layer:
            layer["ind"] = int(layer["ind"])
        if "parent" in layer:
            p = int(layer["parent"])
            if p in valid_inds:
                layer["parent"] = p
            else:
                del layer["parent"]
        layer.pop("ct", None)
        if "shapes" in layer:
            layer["shapes"] = clean_shapes(layer["shapes"])
        for sub in layer.get("layers", []):
            fix_layer(sub)
        return layer

    fixed_layers = []
    for l in anim.get("layers", []):
        fix_layer(l)
        shapes = l.get("shapes", [])
        if l.get("ty") == 4 and not shapes:
            continue  
        fixed_layers.append(l)
    anim["layers"] = fixed_layers

    for asset in anim.get("assets", []):
        if "layers" in asset:
            fixed = []
            for l in asset["layers"]:
                fix_layer(l)
                if l.get("ty") == 4 and not l.get("shapes"):
                    continue
                fixed.append(l)
            asset["layers"] = fixed

    return anim

def tokens_to_lottie_json(generated_ids):
    reconstructed_tensor = LottieTensor.from_list(generated_ids)
    reconstructed_sequence = reconstructed_tensor.to_sequence()
    reconstructed = from_sequence(reconstructed_sequence)

    json_animation = {
        "v": reconstructed.get("v", "5.5.2"),
        "fr": reconstructed.get("fr", 8),
        "ip": reconstructed.get("ip", 0),
        "op": reconstructed.get("op", 16),
        "w": reconstructed.get("w", 512),
        "h": reconstructed.get("h", 512),
        "nm": reconstructed.get("nm", "Animation"),
        "ddd": reconstructed.get("ddd", 0),
        "assets": [],
        "layers": [],
    }

    if "fonts" in reconstructed and reconstructed["fonts"]:
        fonts_data = reconstructed["fonts"]
        if isinstance(fonts_data, dict) and "list" in fonts_data:
            fonts_json = {"list": []}
            for font in fonts_data["list"]:
                if isinstance(font, Font):
                    fonts_json["list"].append(font_to_json(font))
                else:
                    fonts_json["list"].append(font)
            json_animation["fonts"] = fonts_json

    if "chars" in reconstructed and reconstructed["chars"]:
        chars_json = []
        for char in reconstructed["chars"]:
            if isinstance(char, Chars):
                chars_json.append(char_to_json(char))
            else:
                chars_json.append(char)
        json_animation["chars"] = chars_json

    for asset in reconstructed.get("assets", []):
        asset_json = dict(asset)
        if "layers" in asset:
            asset_json["layers"] = []
            for layer in asset["layers"]:
                if isinstance(layer, ShapeLayer):
                    asset_json["layers"].append(shape_layer_to_json(layer))
                elif isinstance(layer, NullLayer):
                    asset_json["layers"].append(null_layer_to_json(layer))
                elif isinstance(layer, PreCompLayer):
                    asset_json["layers"].append(precomp_layer_to_json(layer))
                elif isinstance(layer, TextLayer):
                    asset_json["layers"].append(text_layer_to_json(layer))
                elif isinstance(layer, SolidColorLayer):
                    asset_json["layers"].append(solid_layer_to_json(layer))
                else:
                    asset_json["layers"].append(layer)
        json_animation["assets"].append(asset_json)

    for layer in reconstructed.get("layers", []):
        if isinstance(layer, ShapeLayer):
            json_animation["layers"].append(shape_layer_to_json(layer))
        elif isinstance(layer, NullLayer):
            json_animation["layers"].append(null_layer_to_json(layer))
        elif isinstance(layer, PreCompLayer):
            json_animation["layers"].append(precomp_layer_to_json(layer))
        elif isinstance(layer, TextLayer):
            json_animation["layers"].append(text_layer_to_json(layer))
        elif isinstance(layer, SolidColorLayer):
            json_animation["layers"].append(solid_layer_to_json(layer))
        else:
            json_animation["layers"].append(layer)

    json_animation = fix_lottie_json(json_animation)

    return json_animation

_lottie_js_cache = None

def get_lottie_js():
    global _lottie_js_cache
    if _lottie_js_cache is not None:
        return _lottie_js_cache

    local_path = "lottie.min.js"
    if os.path.exists(local_path):
        with open(local_path, 'r', encoding='utf-8') as f:
            _lottie_js_cache = f.read()
        print(f"✅ Loaded local lottie.min.js ({len(_lottie_js_cache)} bytes)")
    else:
        _lottie_js_cache = '<script src="https://cdnjs.cloudflare.com/ajax/libs/lottie-web/5.12.2/lottie.min.js"></script>'
        print("⚠️ Using CDN lottie.min.js (local file not found)")

    return _lottie_js_cache

def create_lottie_html(animation_data, height=600):
    bg_style = """
        background-image:
            linear-gradient(45deg, #666666 25%, transparent 25%),
            linear-gradient(-45deg, #666666 25%, transparent 25%),
            linear-gradient(45deg, transparent 75%, #666666 75%),
            linear-gradient(-45deg, transparent 75%, #666666 75%);
        background-size: 16px 16px;
        background-position: 0 0, 0 8px, 8px -8px, -8px 0px;
        background-color: #444444;
    """

    lottie_js = get_lottie_js()

    if lottie_js.startswith('<script'):
        lottie_script = lottie_js
    else:
        lottie_script = f"<script>{lottie_js}</script>"

    animation_json_escaped = json.dumps(animation_data).replace('\\', '\\\\').replace("'", "\\'")

    anim_width = animation_data.get('w', 512)
    anim_height = animation_data.get('h', 512)

    inner_html = f"""<!DOCTYPE html>
<html>
<head>
    <meta charset="UTF-8">
    {lottie_script}
    <style>
        * {{ margin: 0; padding: 0; box-sizing: border-box; }}
        html, body {{
            width: 100%;
            height: 100%;
            overflow: hidden;
            background: transparent;
        }}
        #lottie-container {{
            width: 100%;
            height: 100%;
            {bg_style}
            display: flex;
            align-items: center;
            justify-content: center;
            padding: 20px;
        }}
        #lottie-animation {{
            max-width: 100%;
            max-height: 100%;
            width: {anim_width}px;
            height: {anim_height}px;
        }}
        .error {{ color: white; text-align: center; padding: 20px; }}
    </style>
</head>
<body>
    <div id="lottie-container">
        <div id="lottie-animation"></div>
    </div>
    <script>
        function renderLottie() {{
            try {{
                if (typeof lottie === 'undefined') {{
                    setTimeout(renderLottie, 100);
                    return;
                }}
                var animationData = JSON.parse('{animation_json_escaped}');
                lottie.loadAnimation({{
                    container: document.getElementById('lottie-animation'),
                    renderer: 'svg',
                    loop: true,
                    autoplay: true,
                    animationData: animationData
                }});
            }} catch (e) {{
                console.error('Lottie render error:', e);
                document.getElementById('lottie-container').innerHTML =
                    '<p class="error">Failed: ' + e.message + '</p>';
            }}
        }}
        renderLottie();
    </script>
</body>
</html>"""

    inner_html_b64 = base64.b64encode(inner_html.encode('utf-8')).decode('utf-8')
    iframe_html = f'<iframe src="data:text/html;base64,{inner_html_b64}" style="width:100%; height:{height}px; border:none; border-radius:8px;"></iframe>'

    return iframe_html

def save_json_to_temp(lottie_json):
    if lottie_json is None:
        return None

    temp_file = tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False, prefix='lottie_')
    json.dump(lottie_json, temp_file, indent=2)
    temp_file.close()
    return temp_file.name

def process_text_to_lottie(text_prompt, max_tokens, use_sampling, temperature, top_p, top_k):
    with generation_lock:
        try:
            start_time = time.time()

            if not text_prompt or not text_prompt.strip():
                return None, "❌ Please enter a text description", None
            model, processor, device = load_model_once()

            messages = build_messages("text", text_prompt=text_prompt)
            inputs = prepare_inference_input(processor, messages, device)

            generated_ids = generate_lottie(
                model, inputs, max_tokens, device,
                use_sampling, temperature, top_p, top_k
            )

            del inputs
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
            if torch.xpu.is_available():
                torch.xpu.empty_cache()

            lottie_json = tokens_to_lottie_json(generated_ids)

            html = create_lottie_html(lottie_json, height=600)

            elapsed_time = time.time() - start_time

            status = f"✅ Generated {len(generated_ids)} tokens | Layers: {len(lottie_json.get('layers', []))} | {lottie_json.get('fr', 8)} fps | Time: {elapsed_time:.1f}s"

            temp_path = save_json_to_temp(lottie_json)

            return html, status, temp_path

        except Exception as e:
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
            if torch.xpu.is_available():
                torch.xpu.empty_cache()
            return None, f"❌ Error: {str(e)}", None

def load_image_from_file(file_path):

    ext = os.path.splitext(file_path)[1].lower()

    if ext == '.svg':
        import cairosvg
        import io
        png_bytes = cairosvg.svg2png(url=file_path, output_width=448, output_height=448)
        image = PILImage.open(io.BytesIO(png_bytes))
    else:
        image = PILImage.open(file_path)

    if image.mode == 'RGBA':
        image = add_random_background(image)
    else:
        image = image.convert('RGB')

    return image


def process_image_to_lottie(image_file, text_description, max_tokens, use_sampling, temperature, top_p, top_k):
    with generation_lock:
        try:
            start_time = time.time()

            if image_file is None:
                return None, "❌ Please upload an image", None

            model, processor, device = load_model_once()

            image = load_image_from_file(image_file)

            image = image.resize((448, 448), PILImage.LANCZOS)

            desc = text_description if text_description else "A simple animation"
            messages = build_messages("image", text_prompt=desc, image=image)
            inputs = prepare_inference_input(processor, messages, device)

            generated_ids = generate_lottie(
                model, inputs, max_tokens, device,
                use_sampling, temperature, top_p, top_k
            )

            del inputs
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
            if torch.xpu.is_available():
                torch.xpu.empty_cache()

            lottie_json = tokens_to_lottie_json(generated_ids)

            html = create_lottie_html(lottie_json, height=600)

            elapsed_time = time.time() - start_time

            status = f"✅ Generated {len(generated_ids)} tokens | Layers: {len(lottie_json.get('layers', []))} | {lottie_json.get('fr', 8)} fps | Time: {elapsed_time:.1f}s"

            temp_path = save_json_to_temp(lottie_json)

            return html, status, temp_path

        except Exception as e:
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
            if torch.xpu.is_available():
                torch.xpu.empty_cache()
            return None, f"❌ Error: {str(e)}", None

def process_video_to_lottie(video, max_tokens, use_sampling, temperature, top_p, top_k):
    with generation_lock:
        try:
            start_time = time.time()

            if video is None:
                return None, "❌ Please upload a video/GIF/WebP file", None

            import os
            ext = os.path.splitext(video)[1].lower() if isinstance(video, str) else ''
            if ext not in ('.mp4', '.avi', '.mov', '.gif', '.webp'):
                return None, f"❌ Unsupported format: {ext}. Please upload MP4/AVI/MOV/GIF/WebP", None

            model, processor, device = load_model_once()

            frames = load_frames_from_video(video, num_frames=8)

            messages = build_messages("video", video_frames=frames)
            inputs = prepare_inference_input(processor, messages, device)

            generated_ids = generate_lottie(
                model, inputs, max_tokens, device,
                use_sampling, temperature, top_p, top_k
            )

            del inputs
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
            if torch.xpu.is_available():
                torch.xpu.empty_cache()

            lottie_json = tokens_to_lottie_json(generated_ids)

            html = create_lottie_html(lottie_json, height=600)

            elapsed_time = time.time() - start_time

            status = f"✅ Generated {len(generated_ids)} tokens (from {len(frames)} frames) | Layers: {len(lottie_json.get('layers', []))} | {lottie_json.get('fr', 8)} fps | Time: {elapsed_time:.1f}s"

            temp_path = save_json_to_temp(lottie_json)

            return html, status, temp_path

        except Exception as e:
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
            if torch.xpu.is_available():
                torch.xpu.empty_cache()
            return None, f"❌ Error: {str(e)}", None

def create_gradio_interface():

    with gr.Blocks(title="OmniLottie Demo Page", theme=gr.themes.Soft()) as demo:
        gr.Markdown("# 🎨 OmniLottie Demo Page")
        gr.Markdown("Offical Demo Page of OmniLottie")
        gr.Markdown("Generate Lottie animations from text, images, or videos")

        with gr.Tabs() as tabs:
            with gr.Tab("📝 Text-to-Lottie"):
                gr.Markdown("""
                ### 💡 Prompt Tips for Better Results

                **Good prompts should describe:**
                1. **Main Object**: What is being animated (e.g., "a blue bird", "a yellow folder icon", "an orange piggy bank")
                2. **Motion Pattern**: How it moves (e.g., "appearing, pulsing while sliding", "fading in, floating toward", "bouncing up and down")
                3. **Direction**: Where it moves (e.g., "downward", "toward the top-left", "back to its start")
                4. **Loop Behavior**: How it repeats (e.g., "repeating seamlessly", "looping smoothly", "repeating continuously")

                **Example Patterns:**
                - 🔄 **Simple Loop**: "a red ball appearing, bouncing up and down, then fading out, repeating seamlessly"
                - 🎯 **Movement**: "a blue arrow sliding from left to right, then quickly returning to start, looping continuously"
                - 💫 **Transformation**: "a yellow star fading in while rotating 360 degrees, holds briefly, then fading out, repeating smoothly"
                - 🎨 **Static Icon**: "static illustration of a cartoon cat's face with a cute expression, light orange body, red inner ears"
                - 👤 **Character**: "animated cartoon figure dressed in a beige suit with a white shirt, holding a gray tablet"

                **Pro Tips:**
                - Be specific about colors, shapes, and movements
                - Describe motion phases clearly (appear → move → hold → return)
                - Use descriptive motion verbs: sliding, pulsing, drifting, bouncing, rotating, fading
                - For icons: include style details (outline, colors, decorations)
                """)

                with gr.Row():
                    with gr.Column(scale=1):
                        text_input = gr.Textbox(
                            label="Text Description",
                            placeholder="Example: a blue bird appearing, pulsing while sliding downward, lingers briefly, then growing back while sliding upward to reset, repeating seamlessly",
                            lines=5
                        )

                        with gr.Accordion("⚙️ Generation Settings", open=False):
                            gr.Markdown("""
                            **Parameter Guide:**
                            - **Max Tokens**: Higher = more complex animations (slower), Lower = simpler animations (faster)
                            - **Top-p & Top-k**: Higher = more random/creative, Lower = more stable/consistent
                            - **Temperature**: Higher = more diverse, Lower = more deterministic

                            💡 **Quick Tips:**
                            - For complex animations: increase max tokens to 5856
                            - For faster generation: reduce max tokens to 2048-3072
                            - For more creative results: increase top-p (0.5-0.8) and top-k (20-50)
                            - For consistent results: decrease top-p (0.1-0.25) and top-k (5-10)
                            """)
                            text_max_tokens = gr.Slider(512, 5856, value=5556, step=256, label="Max Tokens")
                            text_use_sampling = gr.Checkbox(label="Use Sampling", value=True)
                            text_temperature = gr.Slider(0.1, 2.0, value=0.9, step=0.1, label="Temperature")
                            text_top_p = gr.Slider(0.1, 1.0, value=0.25, step=0.1, label="Top-p")
                            text_top_k = gr.Slider(1, 100, value=5, step=1, label="Top-k")

                        text_generate_btn = gr.Button("🚀 Generate", variant="primary", size="lg")

                        # Generation time tips
                        gr.Markdown("""
                        ⏱️ **Generation Time:**
                        - Simple icons/shapes: ~30-60 seconds (1000-2000 tokens)
                        - Medium animations: ~1-2 minutes (2000-3500 tokens)
                        - Complex characters: ~4-5 minutes (4500-6000 tokens)

                        Please be patient! Complex animations take time to generate. ☕
                        """)

                        text_status = gr.Markdown()

                    with gr.Column(scale=1):
                        text_output = gr.HTML(label="Animation Preview")
                        text_json_file = gr.File(label="📥 Download JSON", visible=True)

                def get_text_examples():
                    examples = []
                    demo_txt_path = "./example/demo.txt"
                    if os.path.exists(demo_txt_path):
                        with open(demo_txt_path, 'r', encoding='utf-8') as f:
                            lines = [line.strip() for line in f.readlines() if line.strip()]
                            examples = [[line] for line in lines[:50]]  
                    return examples

                gr.Examples(
                    examples=get_text_examples(),
                    inputs=text_input,
                    label="📂 Example Prompts (Click to Load)",
                    examples_per_page=10,
                    cache_examples=False
                )

                text_generate_btn.click(
                    fn=process_text_to_lottie,
                    inputs=[text_input, text_max_tokens, text_use_sampling, text_temperature, text_top_p, text_top_k],
                    outputs=[text_output, text_status, text_json_file]
                )

            with gr.Tab("🖼️ Text+Image-to-Lottie"):
                with gr.Row():
                    with gr.Column(scale=1):
                        image_input = gr.Image(
                            label="Upload Image",
                            type="filepath",
                            sources=["upload"]
                        )
                        image_text_input = gr.Textbox(
                            label="Animation Description",
                            placeholder="Example: The object rotates 360 degrees",
                            lines=3
                        )

                        with gr.Accordion("⚙️ Generation Settings", open=False):
                            gr.Markdown("""
                            **Parameter Guide:**
                            - **Max Tokens**: Higher = more complex animations (slower), Lower = simpler (faster)
                            - **Top-p & Top-k**: Higher = more creative/random, Lower = more stable
                            - **Temperature**: Controls output diversity
                            """)
                            image_max_tokens = gr.Slider(512, 5556, value=5556, step=256, label="Max Tokens")
                            image_use_sampling = gr.Checkbox(label="Use Sampling", value=True)
                            image_temperature = gr.Slider(0.1, 2.0, value=0.9, step=0.1, label="Temperature")
                            image_top_p = gr.Slider(0.1, 1.0, value=0.25, step=0.05, label="Top-p")
                            image_top_k = gr.Slider(1, 100, value=5, step=1, label="Top-k")

                        image_generate_btn = gr.Button("🚀 Generate", variant="primary", size="lg")

                        # Generation time tips
                        gr.Markdown("""
                        ⏱️ **Generation Time:** ~1-5 minutes depending on complexity
                        """)

                        image_status = gr.Markdown()

                    with gr.Column(scale=1):
                        image_output = gr.HTML(label="Animation Preview")
                        image_json_file = gr.File(label="📥 Download JSON", visible=True)

                def get_image_text_examples():
                    examples = []
                    demo_images_dir = "./example/demo_images"
                    if os.path.exists(demo_images_dir):
                        png_files = sorted([f for f in os.listdir(demo_images_dir) if f.endswith('.png')])
                        for png_file in png_files[:50]:  
                            base_name = os.path.splitext(png_file)[0]
                            txt_file = os.path.join(demo_images_dir, f"{base_name}.txt")
                            png_path = os.path.join(demo_images_dir, png_file)
                            if os.path.exists(txt_file):
                                with open(txt_file, 'r', encoding='utf-8') as f:
                                    text_desc = f.read().strip()
                                examples.append([png_path, text_desc])
                    return examples

                gr.Examples(
                    examples=get_image_text_examples(),
                    inputs=[image_input, image_text_input],
                    label="📂 Example Images (Click to Load)",
                    examples_per_page=5,
                    cache_examples=False
                )

                image_generate_btn.click(
                    fn=process_image_to_lottie,
                    inputs=[image_input, image_text_input, image_max_tokens, image_use_sampling,
                           image_temperature, image_top_p, image_top_k],
                    outputs=[image_output, image_status, image_json_file]
                )

            # Tab 3: Video-to-Lottie
            with gr.Tab("🎥 Video-to-Lottie"):
                with gr.Row():
                    with gr.Column(scale=1):
                        video_input = gr.Video(
                            label="Upload Video / GIF / WebP",
                            sources=["upload"]
                        )

                        with gr.Accordion("⚙️ Generation Settings", open=False):
                            gr.Markdown("""
                            **Parameter Guide:**
                            - **Max Tokens**: Higher = more complex animations (slower), Lower = simpler (faster)
                            - **Top-p & Top-k**: Higher = more creative/random, Lower = more stable
                            - **Temperature**: Controls output diversity
                            """)
                            video_max_tokens = gr.Slider(512, 5556, value=5556, step=256, label="Max Tokens")
                            video_use_sampling = gr.Checkbox(label="Use Sampling", value=True)
                            video_temperature = gr.Slider(0.1, 2.0, value=0.9, step=0.1, label="Temperature")
                            video_top_p = gr.Slider(0.1, 1.0, value=0.25, step=0.05, label="Top-p")
                            video_top_k = gr.Slider(1, 100, value=5, step=1, label="Top-k")

                        video_generate_btn = gr.Button("🚀 Generate", variant="primary", size="lg")

                        # Generation time tips
                        gr.Markdown("""
                        ⏱️ **Generation Time:** ~2-5 minutes depending on video complexity
                        """)

                        video_status = gr.Markdown()

                    with gr.Column(scale=1):
                        video_output = gr.HTML(label="Animation Preview")
                        video_json_file = gr.File(label="📥 Download JSON", visible=True)

                def get_video_examples():
                    examples = []
                    demo_video_dir = "./example/demo_video"
                    if os.path.exists(demo_video_dir):
                        video_files = sorted([os.path.join(demo_video_dir, f)
                                            for f in os.listdir(demo_video_dir)
                                            if f.endswith('.mp4')])
                        examples = [[vf] for vf in video_files[:50]]  
                    return examples

                gr.Examples(
                    examples=get_video_examples(),
                    inputs=video_input,
                    label="📂 Example Videos (Click to Load)",
                    examples_per_page=5,
                    cache_examples=False
                )

                video_generate_btn.click(
                    fn=process_video_to_lottie,
                    inputs=[video_input, video_max_tokens, video_use_sampling,
                           video_temperature, video_top_p, video_top_k],
                    outputs=[video_output, video_status, video_json_file]
                )

        gr.Markdown("---")
        gr.Markdown("""
        ⚠️ **Important Note:** This demo processes one request at a time.
        If another user is generating, your request will wait in queue until the current one completes.
        """)

    return demo

if __name__ == "__main__":
    demo = create_gradio_interface()
    demo.launch(
        server_name="0.0.0.0",
        server_port=7861,
        share=False,
        show_error=True
    )