"""Kaleidoscope: upload an image, fan it out to every configured model
provider concurrently, and preview each provider's generated video.

Every run (input image + per-provider output videos + metadata) is
persisted under /data so past runs can be shown in the history section.

Provider API keys are BYOK: each browser user supplies their own key(s) in
the UI. Keys are persisted in browser localStorage and are never written to
disk by this server.
"""

from __future__ import annotations

import dataclasses
import logging
import os
import time
import warnings
from concurrent.futures import ThreadPoolExecutor, as_completed

os.environ.setdefault("GRADIO_ANALYTICS_ENABLED", "False")

import gradio as gr
from dotenv import load_dotenv

from providers import PROVIDERS, apply_global_settings
from runs import (
    Run,
    RunResult,
    create_run,
    ensure_data_dirs,
    get_data_dir,
    load_runs,
    save_output_video,
    save_run,
)

load_dotenv()

warnings.filterwarnings(
    "ignore",
    message=".*HTTP_422_UNPROCESSABLE_ENTITY.*",
    category=DeprecationWarning,
    module=r"gradio\.routes",
)

TOKEN_ENVS = sorted({provider.api_token_env for provider in PROVIDERS})
EXTRA_CONFIG_ENVS = sorted({name for provider in PROVIDERS for name in provider.extra_config_envs})

logging.basicConfig(
    level=os.environ.get("LOG_LEVEL", "INFO").upper(),
    format="%(asctime)s %(levelname)s %(name)s: %(message)s",
)
logger = logging.getLogger(__name__)

ensure_data_dirs()
logger.info("Data directory resolved to %s", get_data_dir())



def _call_with_timing(provider, image_path, prompt, duration_seconds, resolution):
    start = time.monotonic()
    effective_provider = apply_global_settings(provider, prompt, duration_seconds, resolution)
    logger.info("Provider %s: starting", provider.name)
    try:
        video_bytes = effective_provider.call(effective_provider, image_path)
        duration = time.monotonic() - start
        logger.info("Provider %s: succeeded in %.1fs", provider.name, duration)
        return effective_provider, video_bytes, None, duration
    except Exception as exc:  # noqa: BLE001 - isolate provider failure, report it in the UI
        duration = time.monotonic() - start
        logger.exception("Provider %s: failed after %.1fs", provider.name, duration)
        return effective_provider, None, exc, duration


def run_providers(providers, image_path, prompt, duration_seconds, resolution):
    """Fans out image_path to every provider concurrently.

    Yields (provider, video_bytes, error, duration_seconds) tuples in
    completion order (not submission order) so the caller can update the UI
    incrementally instead of waiting on the slowest provider.
    """
    with ThreadPoolExecutor(max_workers=max(len(providers), 1)) as executor:
        futures = [
            executor.submit(_call_with_timing, provider, image_path, prompt, duration_seconds, resolution)
            for provider in providers
        ]
        for future in as_completed(futures):
            yield future.result()


def format_metadata(provider, status, error=None, duration=None):
    lines = [f"**{provider.name}**", f"- status: {status}"]
    if duration is not None:
        lines.append(f"- duration: {duration:.1f}s")
    if provider.params:
        lines.append(f"- params: {provider.params}")
    if error is not None:
        lines.append(f"- error: {error}")
    return "\n".join(lines)


# Small per-provider status indicator shown next to its checkbox (visible
# without expanding that provider's accordion) - a CSS spinner while a
# request is in flight, then a static icon once it settles. The spinner's
# look is defined by STATUS_SPINNER_CSS, injected once via gr.Blocks(css=...).
_STATUS_ICONS = {
    "idle": "",
    "running": '<span class="ks-spinner" title="Running"></span>',
    "ok": '<span title="Done" style="font-size:1.1em;">\u2705</span>',
    "error": '<span title="Error" style="font-size:1.1em;">\u274c</span>',
    "skipped": '<span title="Skipped" style="font-size:1.1em;">\u23ed\ufe0f</span>',
}


def format_status_icon(status: str) -> str:
    return _STATUS_ICONS.get(status, "")


def format_result_metadata(result: RunResult) -> str:
    lines = [
        f"**{result.provider_name}**",
        f"- status: {result.status}",
        f"- duration: {result.duration_seconds:.1f}s",
        f"- params: {result.params_used}",
    ]
    if result.error:
        lines.append(f"- error: {result.error}")
    return "\n".join(lines)


def on_submit(image_path, prompt, duration_seconds, resolution, history, *dynamic_inputs):
    if not image_path:
        raise gr.Error("Please upload an image first.")

    duration_seconds = int(duration_seconds) if duration_seconds is not None else None

    token_values = list(dynamic_inputs[: len(TOKEN_ENVS)])
    extra_config_values = list(dynamic_inputs[len(TOKEN_ENVS) : len(TOKEN_ENVS) + len(EXTRA_CONFIG_ENVS)])
    enabled_values = list(dynamic_inputs[len(TOKEN_ENVS) + len(EXTRA_CONFIG_ENVS) :])

    token_map = {
        token_env: (token_value.strip() if isinstance(token_value, str) else "")
        for token_env, token_value in zip(TOKEN_ENVS, token_values)
    }
    extra_config_map = {
        config_env: (config_value.strip() if isinstance(config_value, str) else "")
        for config_env, config_value in zip(EXTRA_CONFIG_ENVS, extra_config_values)
    }

    selected_providers = [provider for provider, enabled in zip(PROVIDERS, enabled_values) if enabled]
    if not selected_providers:
        raise gr.Error("Select at least one model to run.")

    missing_token_envs = sorted(
        {
            provider.api_token_env
            for provider in selected_providers
            if not token_map.get(provider.api_token_env)
        }
    )
    missing_extra_config_envs = sorted(
        {
            config_env
            for provider in selected_providers
            for config_env in provider.extra_config_envs
            if not extra_config_map.get(config_env)
        }
    )
    if missing_token_envs or missing_extra_config_envs:
        missing_label = ", ".join(missing_token_envs + missing_extra_config_envs)
        raise gr.Error(f"Missing required configuration: {missing_label}")

    selected_runtime_providers = [
        dataclasses.replace(
            provider,
            api_token_value=token_map.get(provider.api_token_env, ""),
            extra_config_values={
                config_env: extra_config_map.get(config_env, "") for config_env in provider.extra_config_envs
            },
        )
        for provider in selected_providers
    ]

    run_id, run_dir, input_path = create_run(image_path)
    logger.info(
        "Run %s: created (selected_providers=%d total_providers=%d, prompt=%r)",
        run_id,
        len(selected_runtime_providers),
        len(PROVIDERS),
        bool(prompt),
    )
    provider_index = {provider.name: i for i, provider in enumerate(PROVIDERS)}
    results_by_name: dict[str, RunResult] = {}
    selected_names = {provider.name for provider in selected_runtime_providers}

    metadata_values = []
    status_values = []
    for provider in PROVIDERS:
        if provider.name in selected_names:
            metadata_values.append(
                format_metadata(apply_global_settings(provider, prompt, duration_seconds, resolution), "running")
            )
            status_values.append(format_status_icon("running"))
        else:
            metadata_values.append(
                format_metadata(apply_global_settings(provider, prompt, duration_seconds, resolution), "skipped")
            )
            status_values.append(format_status_icon("skipped"))
            results_by_name[provider.name] = RunResult(
                provider_name=provider.name,
                output_path=None,
                status="skipped",
                error=None,
                duration_seconds=0.0,
                params_used=provider.params,
            )
    video_values = [None] * len(PROVIDERS)
    yield metadata_values + status_values + video_values + [history]

    for provider, video_bytes, error, duration in run_providers(
        selected_runtime_providers, image_path, prompt, duration_seconds, resolution
    ):
        index = provider_index[provider.name]

        if error is None:
            output_path = save_output_video(run_dir, provider.name, video_bytes)
            results_by_name[provider.name] = RunResult(
                provider_name=provider.name,
                output_path=output_path,
                status="ok",
                error=None,
                duration_seconds=duration,
                params_used=provider.params,
            )
            metadata_values[index] = format_metadata(provider, "ok", duration=duration)
            status_values[index] = format_status_icon("ok")
            video_values[index] = output_path
        else:
            results_by_name[provider.name] = RunResult(
                provider_name=provider.name,
                output_path=None,
                status="error",
                error=str(error),
                duration_seconds=duration,
                params_used=provider.params,
            )
            metadata_values[index] = format_metadata(provider, "error", error=str(error), duration=duration)
            status_values[index] = format_status_icon("error")
            video_values[index] = None

        yield metadata_values + status_values + video_values + [history]

    run = Run(
        run_id=run_id,
        timestamp=time.time(),
        input_path=input_path,
        prompt=prompt or None,
        results=[results_by_name[provider.name] for provider in PROVIDERS],
    )
    save_run(run)
    logger.info("Run %s: saved", run_id)

    yield metadata_values + status_values + video_values + [history + [run]]


# CSS for the small per-provider running spinner (see format_status_icon) -
# a plain rotating-border circle so no extra asset/dependency is needed.
STATUS_SPINNER_CSS = """
.ks-spinner {
    display: inline-block;
    width: 14px;
    height: 14px;
    border: 2px solid var(--border-color-primary, #999);
    border-top-color: var(--color-accent, #555);
    border-radius: 50%;
    animation: ks-spin 0.8s linear infinite;
    vertical-align: middle;
}
@keyframes ks-spin {
    to { transform: rotate(360deg); }
}
"""

with gr.Blocks(title="Kaleidoscope") as demo:
    gr.Markdown(
        "# Kaleidoscope\n"
        "Upload an image to generate a short video with each configured provider."
    )

    with gr.Accordion(label="Settings", open=False):
        gr.Markdown(
            "Applied to every model that supports it (translated into each "
            "model's own params as needed); models without a matching field "
            "ignore it."
        )
        duration_input = gr.Number(
            value=6,
            precision=0,
            label="Duration (seconds)",
        )
        resolution_input = gr.Dropdown(
            choices=["720p", "1080p"],
            value="720p",
            label="Resolution",
        )

    with gr.Accordion(label="API Keys (BYOK)", open=False):
        gr.Markdown("Keys are saved in your browser local storage and never persisted by this server.")
        token_inputs = []
        for token_env in TOKEN_ENVS:
            token_inputs.append(
                gr.Textbox(
                    label=token_env,
                    placeholder=f"Enter {token_env}",
                    type="password",
                )
            )

        extra_config_inputs = []
        if EXTRA_CONFIG_ENVS:
            gr.Markdown(
                "Some providers need more than just a key (e.g. a "
                "per-resource endpoint URL) - configure those here too."
            )
            # Per-field placeholder overrides for config values whose
            # expected format isn't obvious from the label alone (falls back
            # to a generic placeholder for anything not listed here).
            extra_config_placeholders = {
                "AZURE_SORA_ENDPOINT": "https://<resource>.openai.azure.com/openai/v1",
            }
            for config_env in EXTRA_CONFIG_ENVS:
                extra_config_inputs.append(
                    gr.Textbox(
                        label=config_env,
                        placeholder=extra_config_placeholders.get(config_env, f"Enter {config_env}"),
                    )
                )

    with gr.Accordion(label="Input image", open=True) as image_accordion:
        image_input = gr.Image(type="filepath", label="Input image")
    prompt_input = gr.Textbox(
        label="Prompt (optional)",
        placeholder="Used by providers that support a prompt; ignored by the rest.",
    )
    submit_btn = gr.Button("Submit", variant="primary")

    gr.Markdown("## Results")
    check_all_btn = gr.Button("Check all models")
    enabled_components = []
    status_components = []
    metadata_components = []
    video_components = []
    for provider in PROVIDERS:
        with gr.Row():
            with gr.Column(scale=0, min_width=90):
                enabled_components.append(gr.Checkbox(label="Use", value=True))
                status_components.append(gr.HTML(value=format_status_icon("idle")))
            with gr.Accordion(label=provider.name, open=False):
                with gr.Row():
                    with gr.Column(scale=1):
                        metadata_components.append(gr.Markdown(format_metadata(provider, "idle")))
                    with gr.Column(scale=2):
                        video_components.append(gr.Video(label=provider.name, interactive=False))

    gr.Markdown("## Past runs")
    # Start empty and populate via demo.load() below, so every new page
    # load/session re-reads runs from disk instead of reusing a snapshot
    # taken once when the server process started.
    history_state = gr.State([])
    demo.load(load_runs, outputs=history_state)

    persisted_inputs = token_inputs + extra_config_inputs
    storage_keys = [f"kaleidoscope.byok.{token_env.lower()}" for token_env in TOKEN_ENVS] + [
        f"kaleidoscope.byok.{config_env.lower()}" for config_env in EXTRA_CONFIG_ENVS
    ]
    if persisted_inputs:
        # Gradio's JS-return convention mirrors Python fn returns: with exactly
        # one output, return the bare value (not wrapped in an array); with
        # multiple outputs, return an array of values in output order.
        # Returning a 1-element array for a single output corrupts that
        # component's internal block registry (it gets replaced by the raw
        # list), crashing later interactions with
        # "'list' object has no attribute 'stateful'" - so the single- and
        # multi-output cases must be built differently below.
        if len(persisted_inputs) == 1:
            load_js = f"() => localStorage.getItem('{storage_keys[0]}') || ''"
        else:
            load_js = (
                "() => ["
                + ", ".join(f"localStorage.getItem('{key}') || ''" for key in storage_keys)
                + "]"
            )
        demo.load(fn=None, inputs=None, outputs=persisted_inputs, js=load_js)

    for storage_key, persisted_input in zip(storage_keys, persisted_inputs):
        persisted_input.change(
            fn=None,
            inputs=[persisted_input],
            outputs=[],
            js=f"(value) => localStorage.setItem('{storage_key}', value || '')",
        )

    check_all_btn.click(
        # A bare `True` (not a list) when there's exactly one output, else a
        # list matching the output count - see the load_js comment above for
        # why this distinction matters.
        fn=(lambda: True) if len(enabled_components) == 1 else (lambda: [True] * len(enabled_components)),
        inputs=None,
        outputs=enabled_components,
    )

    @gr.render(inputs=history_state)
    def render_history(history):
        if not history:
            gr.Markdown("_No past runs yet._")
            return
        for run in reversed(history):
            run_label = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(run.timestamp))
            with gr.Accordion(label=run_label, open=False):
                gr.Image(value=run.input_path, label="Input")
                if run.prompt:
                    gr.Markdown(f"**Prompt:** {run.prompt}")
                for result in run.results:
                    with gr.Row():
                        with gr.Column(scale=1):
                            gr.Markdown(format_result_metadata(result))
                        with gr.Column(scale=2):
                            gr.Video(value=result.output_path, interactive=False)

    submit_btn.click(
        fn=lambda: (gr.update(interactive=False), gr.update(open=False)),
        inputs=None,
        outputs=[submit_btn, image_accordion],
    ).then(
        fn=on_submit,
        inputs=[image_input, prompt_input, duration_input, resolution_input, history_state]
        + token_inputs
        + extra_config_inputs
        + enabled_components,
        outputs=metadata_components + status_components + video_components + [history_state],
    ).then(
        fn=lambda: gr.update(interactive=True),
        inputs=None,
        outputs=submit_btn,
    )

demo.queue()

if __name__ == "__main__":
    logger.info("Starting Kaleidoscope with %d provider(s): %s", len(PROVIDERS), [p.name for p in PROVIDERS])
    # /data (on a Hugging Face Space) lives outside the cwd and the system
    # temp dir, so Gradio refuses to serve run files from it unless the
    # directory is explicitly allow-listed here.
    demo.launch(allowed_paths=[get_data_dir()], css=STATUS_SPINNER_CSS)