import os
import sys
import base64
import time
from pathlib import Path

# Dynamic monkeypatch for gradio_client schema parsing bug (TypeError: argument of type 'bool' is not iterable)
try:
    import gradio_client.utils
    
    _orig_json_schema_to_python_type = gradio_client.utils._json_schema_to_python_type
    def patched_json_schema_to_python_type(schema, defs=None):
        if isinstance(schema, bool):
            return "Any"
        try:
            return _orig_json_schema_to_python_type(schema, defs)
        except Exception:
            return "Any"
            
    _orig_get_type = getattr(gradio_client.utils, "get_type", None)
    if _orig_get_type:
        def patched_get_type(schema):
            if isinstance(schema, bool):
                return "Any"
            try:
                return _orig_get_type(schema)
            except Exception:
                return "Any"
        gradio_client.utils.get_type = patched_get_type
        
    gradio_client.utils._json_schema_to_python_type = patched_json_schema_to_python_type
    print("Successfully patched gradio_client schema parsing helper.")
except Exception as e:
    print(f"gradio_client patch failed: {e}")

import gradio as gr
from PIL import Image, ImageFile
from loguru import logger
from dotenv import load_dotenv

# Set paths before importing task_history and local_inference
os.environ["TASK_CACHE_DIR"] = str(Path(__file__).resolve().parent / "task_cache")
os.environ["TASK_HISTORY_PATH"] = str(Path(__file__).resolve().parent / "task_history.json")

# Import the xfuser mock first to register it in sys.modules
import xfuser

# Import spaces if available (Hugging Face ZeroGPU environment)
try:
    import spaces
except Exception:
    # Dummy decorator for local running/testing
    class spaces:
        @staticmethod
        def GPU(duration=120):
            def decorator(fn):
                return fn
            return decorator

from local_inference import (
    generate_dance_video,
    truncate_audio,
    crop_and_resize,
    FPS
)

import task_history

# Setup directories
TASK_CACHE_DIR = Path(os.environ["TASK_CACHE_DIR"])
TASK_CACHE_DIR.mkdir(parents=True, exist_ok=True)

# Load env vars (including HF_TOKEN)
load_dotenv()
# The HF_TOKEN should be configured in Hugging Face Space settings (secrets)
# and is read directly from the environment.

# Pre-download models on startup so the Space is ready for inference
from huggingface_hub import snapshot_download
try:
    print("Pre-downloading model weights from Hugging Face...")
    snapshot_download(
        repo_id="Wan-AI/Wan-Dancer-14B",
        local_dir="./models/Wan-AI/Wan-Dancer-14B",
        token=os.environ.get("HF_TOKEN")
    )
    print("Pre-download completed successfully.")
except Exception as e:
    print(f"Error pre-downloading models: {e}")
finally:
    if "HF_TOKEN" in os.environ:
        del os.environ["HF_TOKEN"]
        print("Cleared HF_TOKEN from environment variables to prevent ZeroGPU quota override.")

try:
    from local_inference import get_or_init_pipeline
    print("Pre-loading models into CPU memory...")
    get_or_init_pipeline(local_model_path="./models")
    print("Pre-loading models completed successfully.")
except Exception as e:
    print(f"Error pre-loading models: {e}")

DANCE_GENRES = [
    ("Chinese Classical", "chinese_classical"),
    ("K-Pop", "k_pop"),
    ("Street Dance", "street"),
    ("Tap Dance", "tap"),
    ("Latin Dance", "latin"),
]
DANCE_GENRE_LABELS = {v: k for k, v in DANCE_GENRES}

HISTORY_PAGE_SIZE = 20

def _image_to_thumbnail_base64(image_path: str, size: int = 64) -> str | None:
    if not image_path or not os.path.isfile(image_path):
        return None
    try:
        img = Image.open(image_path)
        img = img.convert("RGB")
        img.thumbnail((size, size), Image.Resampling.LANCZOS)
        import io
        buf = io.BytesIO()
        img.save(buf, format="JPEG", quality=75)
        return base64.b64encode(buf.getvalue()).decode("ascii")
    except Exception:
        return None

def _render_history_page(tasks: list, page: int) -> tuple[str, str]:
    total = len(tasks)
    total_pages = max(1, (total + HISTORY_PAGE_SIZE - 1) // HISTORY_PAGE_SIZE)
    page = max(0, min(page, total_pages - 1))
    start = page * HISTORY_PAGE_SIZE
    end = start + HISTORY_PAGE_SIZE
    page_tasks = tasks[start:end]

    if not page_tasks:
        return "No task history.", "Page 0/0, 0 total"

    lines = [
        "| Thumbnail | task_id | Status | Genre | Created | Result Video |",
        "| --- | --- | --- | --- | --- | --- |",
    ]
    for t in page_tasks:
        img_path = t.get("image_local_path") or ""
        thumb_b64 = _image_to_thumbnail_base64(img_path)
        if thumb_b64:
            thumb_cell = f'<img src="data:image/jpeg;base64,{thumb_b64}" width="64" height="64" style="object-fit:cover" />'
        else:
            thumb_cell = "—"
        task_id = (t.get("task_id") or "").replace("|", "\\|")
        status = (t.get("status") or "—").replace("|", "\\|")
        genre_val = t.get("dance_genres") or ""
        genre_label = DANCE_GENRE_LABELS.get(genre_val, genre_val or "—").replace("|", "\\|")
        
        # Format created_at
        created = t.get("created_at") or ""
        if created:
            try:
                from datetime import datetime
                dt = datetime.fromisoformat(created.replace("Z", "+00:00"))
                created = dt.strftime("%Y-%m-%d %H:%M")
            except Exception:
                created = created[:16]
        else:
            created = "—"
            
        video_local = (t.get("video_local_path") or "").strip()
        if video_local and os.path.isfile(video_local):
            video_cell = f"[View Video](/gradio_api/file={video_local})"
        else:
            video_cell = "—"
            
        lines.append(f"| {thumb_cell} | {task_id} | {status} | {genre_label} | {created} | {video_cell} |")

    page_info = f"Page {page + 1}/{total_pages}, {total} total"
    return "\n".join(lines), page_info

def refresh_history(page: int = 0):
    tasks = task_history.load_history()
    md, info = _render_history_page(tasks, page)
    total_pages = max(1, (len(tasks) + HISTORY_PAGE_SIZE - 1) // HISTORY_PAGE_SIZE)
    current_page = max(0, min(page, total_pages - 1))
    return md, info, current_page

def history_prev_page(current_page: int):
    return refresh_history(max(0, current_page - 1))

def history_next_page(current_page: int):
    return refresh_history(current_page + 1)

# ZeroGPU decorated entry handler
@spaces.GPU(duration=90)
def run_wan_dance(
    image_file,
    audio_file,
    dance_genre_value,
    max_audio_duration_value,
    audio_start_sec_value,
    steps_value,
    cfg_value,
    progress=gr.Progress(track_tqdm=True),
    request: gr.Request = None
):
    # Decode crop image if base64 data URL
    if image_file and isinstance(image_file, str) and image_file.startswith("data:image"):
        try:
            header, b64data = image_file.split(",", 1)
            import tempfile
            tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".png", dir=str(TASK_CACHE_DIR))
            tmp.write(base64.b64decode(b64data))
            tmp.close()
            image_file = tmp.name
            logger.info("Decoded crop image: {}", image_file)
        except Exception as e:
            logger.error("Decoding base64 image failed: {}", e)
            return None, f"Image decoding failed: {e}", gr.skip(), gr.skip(), gr.skip()

    if not image_file:
        return None, "Please upload and crop a portrait image (9:16 vertical).", gr.skip(), gr.skip(), gr.skip()
    if not audio_file:
        return None, "Please upload a music file.", gr.skip(), gr.skip(), gr.skip()

    # Truncate audio
    max_dur = int(max_audio_duration_value) if max_audio_duration_value else 10
    start_sec = float(audio_start_sec_value) if audio_start_sec_value else 0.0
    original_audio = audio_file
    audio_file = truncate_audio(audio_file, max_dur, start_sec)
    if audio_file != original_audio:
        logger.info("Audio truncated to {}s (starting at {}s): {}", max_dur, start_sec, audio_file)

    steps = int(steps_value) if steps_value else 24
    cfg = float(cfg_value) if cfg_value else 5.0
    task_id = f"task_{int(time.time())}"

    # Setup local cache copies for task history
    img_ext = Path(image_file).suffix or ".jpg"
    aud_ext = Path(audio_file).suffix or ".mp3"
    
    # Save files to task cache folder
    task_sub_dir = TASK_CACHE_DIR / task_id
    task_sub_dir.mkdir(parents=True, exist_ok=True)
    
    import shutil
    image_local_path = str(task_sub_dir / f"image{img_ext}")
    audio_local_path = str(task_sub_dir / f"audio{aud_ext}")
    shutil.copy2(image_file, image_local_path)
    shutil.copy2(audio_file, audio_local_path)

    try:
        # Run local generation on ZeroGPU
        video_output = generate_dance_video(
            image_path=image_local_path,
            music_path=audio_local_path,
            genre=dance_genre_value,
            output_folder=str(TASK_CACHE_DIR),
            local_model_path="./models",
            seed=0,
            height=1280,
            width=720,
            steps=steps,
            cfg=cfg,
            progress=progress
        )
        
        # Save output in local task sub-directory
        video_local_path = str(task_sub_dir / "video.mp4")
        shutil.move(video_output, video_local_path)

        # Save to history
        task_history.add_task(
            task_id=task_id,
            image_url="",
            audio_url="",
            status="SUCCEEDED",
            result_video_url=None,
            dance_genres=dance_genre_value,
            image_local_path=image_local_path,
            audio_local_path=audio_local_path,
            owner_id=""
        )
        task_history.update_task(task_id, status="SUCCEEDED", video_local_path=video_local_path)

        # Reload history Markdown
        md, info, _ = refresh_history(0)

        return video_local_path, "Generation succeeded!", md, info, 0

    except Exception as e:
        logger.exception("Inference failed")
        task_history.add_task(
            task_id=task_id,
            image_url="",
            audio_url="",
            status="FAILED",
            message=str(e),
            dance_genres=dance_genre_value,
            image_local_path=image_local_path,
            audio_local_path=audio_local_path,
            owner_id=""
        )
        md, info, _ = refresh_history(0)
        return None, f"Generation failed: {e}", md, info, 0

def build_ui():
    with gr.Blocks(
        title="Wan-dancer",
        theme=gr.themes.Soft(
            font=[gr.themes.GoogleFont("Source Sans Pro"), "Arial", "sans-serif"],
        ),
        css="""
        .gradio-container {max-width: none !important;}
        #logo-banner {
            background: #ffffff !important;
            padding: 16px 0 !important;
            border-radius: 8px;
        }
        """,
    ) as demo:
        # Logo Banner
        try:
            _logo_path = Path(__file__).resolve().parent / "logo.png"
            _logo_b64 = base64.b64encode(_logo_path.read_bytes()).decode()
            _logo_html = (
                '<div align="center" style="margin: 12px 0 4px 0;">'
                f'<img src="data:image/png;base64,{_logo_b64}" alt="Wan-dancer" '
                'style="max-width: 480px; width: 100%; height: auto;">'
                '</div>'
            )
        except Exception as _e:
            _logo_html = '<div align="center"><h1>Wan-dancer</h1></div>'
        gr.HTML(_logo_html, elem_id="logo-banner")
        
        gr.Markdown(
            "<div align=\"center\">\n\n"
            "[GitHub](https://github.com/Wan-Video/Wan-Dancer) · "
            "[Hugging Face Model](https://huggingface.co/Wan-AI/Wan-Dancer-14B)\n\n"
            "</div>"
        )
        
        gr.Markdown(
            "**Instructions**\n\n"
            "1. Upload a portrait image (9:16 vertical, full-body/half-body facing front) and a music file (wav/mp3, 10s-30s).\n"
            "2. Select a dance genre and customize generation parameters if desired, then click **Generate**.\n"
            "3. **ZeroGPU Optimization**: To stay within Hugging Face ZeroGPU time limits (120s max), keep **Inference Steps around 12-15** and **Max Audio Duration around 10-15s**."
        )


        with gr.Row(equal_height=False):
            with gr.Column(scale=1):
                gr.Markdown("### Input")
                with gr.Group():
                    image_in = gr.Image(
                        type="filepath",
                        label="Portrait Image (will be auto-cropped to 9:16 vertical)",
                    )
                    audio_in = gr.Audio(
                        type="filepath",
                        label="Music (wav or mp3)",
                        sources=["upload"],
                    )
                    max_audio_duration = gr.Slider(
                        minimum=10,
                        maximum=30,
                        value=10,
                        step=1,
                        label="Max Audio Duration (seconds)",
                        info="Video generation is computed sequentially in 5-second segments. Standard length is 10s.",
                    )
                    audio_start_sec = gr.Slider(
                        minimum=0,
                        maximum=120,
                        value=0,
                        step=1,
                        label="Audio Start Time (seconds)",
                        info="Choose where in the song the dance should start (e.g. when the beat drops).",
                    )
                    dance_genre = gr.Dropdown(
                        choices=[(f"{label} ({v})", v) for label, v in DANCE_GENRES],
                        value="chinese_classical",
                        label="Dance Genre",
                    )
                with gr.Accordion("Parameters", open=False):
                    with gr.Group():
                        with gr.Row():
                            steps = gr.Slider(
                                minimum=10,
                                maximum=50,
                                value=15,
                                step=1,
                                label="Inference Steps",
                                info="Number of denoising steps. Default is 15."
                            )
                            cfg = gr.Slider(
                                minimum=1.0,
                                maximum=15.0,
                                value=5.0,
                                step=0.5,
                                label="Classifier-free Guidance Scale",
                                info="Guidance scale. Default is 5.0."
                            )
                submit_btn = gr.Button("Generate", variant="primary")
            with gr.Column(scale=1):
                gr.Markdown("### Result")
                video_out = gr.Video(label="Generated Video")
                msg_out = gr.Textbox(
                    label="Status / Log",
                    lines=4,
                    interactive=False,
                )

        # ==================== My Tasks ====================
        with gr.Accordion("My Tasks", open=True):
            history_page_state = gr.State(value=0)
            history_markdown = gr.Markdown(value="")
            with gr.Row():
                prev_page_btn = gr.Button("⬅ Prev", scale=1)
                page_info = gr.Textbox(value="", interactive=False, show_label=False, scale=2)
                next_page_btn = gr.Button("Next ➡", scale=1)
                refresh_history_btn = gr.Button("Refresh", scale=1)

        history_outputs = [history_markdown, page_info, history_page_state]
        demo.load(
            fn=refresh_history,
            inputs=[history_page_state],
            outputs=history_outputs,
        )
        refresh_history_btn.click(
            fn=refresh_history,
            inputs=[history_page_state],
            outputs=history_outputs,
        )
        prev_page_btn.click(
            fn=history_prev_page,
            inputs=[history_page_state],
            outputs=history_outputs,
        )
        next_page_btn.click(
            fn=history_next_page,
            inputs=[history_page_state],
            outputs=history_outputs,
        )

        # ==================== Examples (示例) ====================
        _samples_dir = Path(__file__).resolve().parent / "samples"
        gr.Examples(
            examples=[
                [str(_samples_dir / "1/image.png"), str(_samples_dir / "1/audio.mp3"), "street"],
                [str(_samples_dir / "2/image.jpeg"), str(_samples_dir / "2/audio.wav"), "chinese_classical"],
            ],
            inputs=[image_in, audio_in, dance_genre],
            label="Examples"
        )

        submit_btn.click(
            fn=run_wan_dance,
            inputs=[image_in, audio_in, dance_genre, max_audio_duration, audio_start_sec, steps, cfg],
            outputs=[video_out, msg_out, history_markdown, page_info, history_page_state],
        )

    return demo

def main():
    demo = build_ui()
    _samples_dir = str(Path(__file__).resolve().parent / "samples")
    demo.launch(
        allowed_paths=[str(TASK_CACHE_DIR), _samples_dir],
    )

if __name__ == "__main__":
    main()