"""Inter-2 Audio playground: a Gradio UI and API around the Inter-2 Audio model.

Task presets (tasks.py) fill in the prompts for emotion, or for engagement, social
signals, emotion and vocal delivery together; each can optionally be constrained to
its JSON schema. Prompts and decoding are editable.
"""

from __future__ import annotations

import html
import json
import logging
import os
import sys
import time
from collections.abc import Callable
from pathlib import Path
from typing import Any, TypeVar

import gradio as gr
import spaces
import torch
from huggingface_hub import snapshot_download
from inference import (
    DEFAULT_MAX_NEW_TOKENS,
    DEFAULT_REPETITION_PENALTY,
    ENCODER_WINDOW_S,
    TARGET_SR,
    AudioInputError,
    build_messages,
    format_summary,
    load_audio,
    parse_json,
)
from tasks import DEFAULT_TASK, TASKS
from theme import CSS, FOOTER, HERO, THEME

MODEL_ID = os.environ["MODEL_ID"]  # a Space variable: the model repo (or a local copy of it)
# cuda on a GPU or ZeroGPU Space (ZeroGPU has no GPU at import time), else cpu.
_HAS_GPU = torch.cuda.is_available() or bool(os.getenv("SPACES_ZERO_GPU"))
MODEL_DEVICE = os.getenv("MODEL_DEVICE") or ("cuda" if _HAS_GPU else "cpu")
SPACES_GPU_DURATION = int(os.getenv("SPACES_GPU_DURATION", "90"))  # rationale-heavy tasks write more
logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO)


def _load() -> tuple[Any, Any]:
    """(model, processor), via the loader that ships with the model repo."""
    model_dir = MODEL_ID if Path(MODEL_ID).is_dir() else snapshot_download(MODEL_ID)
    sys.path.insert(0, model_dir)
    from inter2_loader import load

    loaded = load(model_dir, MODEL_DEVICE)
    logger.info("Model loaded")
    return loaded


def _compile_grammars() -> dict[str, Any]:
    """Each task's answer schema for xgrammar, in compact JSON layout."""
    try:
        import xgrammar as xgr
    except ImportError:
        logger.warning("xgrammar is not installed; the schema toggle will be ignored.")
        return {}
    vocab_size = model.get_output_embeddings().weight.shape[0]
    compiler = xgr.GrammarCompiler(xgr.TokenizerInfo.from_huggingface(processor.tokenizer, vocab_size=vocab_size))
    return {
        name: compiler.compile_json_schema(
            json.dumps(task.schema), any_whitespace=False, indent=None, separators=(",", ":")
        )
        for name, task in TASKS.items()
    }


model, processor = _load()
grammars = _compile_grammars()

F = TypeVar("F", bound=Callable[..., Any])


def _spaces_gpu(func: F) -> F:
    setting = os.getenv("ENABLE_SPACES_GPU", "auto").lower()
    enabled = setting in {"1", "true", "yes"} or (setting == "auto" and bool(os.getenv("SPACE_ID")))
    return spaces.GPU(duration=SPACES_GPU_DURATION)(func) if enabled else func


@_spaces_gpu
def _generate(inputs: dict[str, Any], gen_kwargs: dict[str, Any], schema_task: str | None) -> tuple[str, int]:
    inputs = {k: v.to(model.device) for k, v in inputs.items()}
    if "input_features" in inputs:
        inputs["input_features"] = inputs["input_features"].to(model.dtype)
    if schema_task is not None:
        from xgrammar.contrib.hf import LogitsProcessor  # fresh per call: it keeps per-sequence state

        gen_kwargs = {**gen_kwargs, "logits_processor": [LogitsProcessor(grammars[schema_task])]}
    prompt_len = inputs["input_ids"].shape[-1]
    with torch.inference_mode():
        output_ids = model.generate(**inputs, **gen_kwargs)
    new_ids = output_ids[:, prompt_len:]
    text = processor.batch_decode(new_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
    return text.strip(), int(new_ids.shape[-1])


def predict(
    audio_path: str | None,
    system_prompt: str,
    user_prompt: str,
    audio_first: bool,
    constrain_json: bool,
    max_new_tokens: int,
    temperature: float,
    top_p: float,
    repetition_penalty: float,
    task: str = DEFAULT_TASK,
) -> tuple[str, str, str]:
    """Run one clip. Returns (summary HTML, parsed JSON, raw text). `task` picks the schema used
    when constrain_json is on; the prompts are whatever was passed in."""
    if not audio_path:
        return _message("Please upload or record an audio clip."), "", ""
    if not user_prompt.strip():
        return _message("The user prompt is empty."), "", ""
    try:
        wav, duration = load_audio(audio_path)
        chat = processor.apply_chat_template(
            build_messages(user_prompt, system_prompt.strip(), audio_first), tokenize=False, add_generation_prompt=True
        )
        inputs = dict(processor(text=[chat], audio=[wav], sampling_rate=TARGET_SR, return_tensors="pt", padding=True))
        gen_kwargs: dict[str, Any] = {
            "max_new_tokens": int(max_new_tokens),
            "repetition_penalty": float(repetition_penalty),
            "do_sample": temperature > 0,
        }
        if temperature > 0:
            gen_kwargs.update(temperature=float(temperature), top_p=float(top_p))
        start = time.time()
        schema_task = task if constrain_json and task in grammars else None
        raw, n_tokens = _generate(inputs, gen_kwargs, schema_task)
        logger.info("%s: %d tokens in %.1fs (constrained=%s)", task, n_tokens, time.time() - start, bool(schema_task))
    except Exception as exc:
        logger.exception("predict failed")
        if isinstance(exc, AudioInputError):  # our own input checks; other errors may carry internals
            return _message(str(exc)), "", ""
        return _message("Something went wrong while analysing this clip. Please try again or use another file."), "", ""
    payload = parse_json(raw)
    pretty = json.dumps(payload, indent=2, ensure_ascii=False) if payload is not None else ""
    summary = format_summary(payload, duration if duration > ENCODER_WINDOW_S else None, n_tokens >= max_new_tokens)
    return summary, pretty, raw


def _message(text: str) -> str:
    return f'<div class="ih-card"><p class="ih-note">{html.escape(text)}</p></div>'


def _preset(task: str) -> tuple[str, str, bool]:
    preset = TASKS[task]
    return preset.system_prompt, preset.user_prompt, preset.constrain


with gr.Blocks(title="Inter-2-Audio Playground") as demo:
    gr.HTML(HERO)
    with gr.Row(equal_height=False):
        with gr.Column(scale=5):
            audio_in = gr.Audio(label="Audio clip", sources=["upload", "microphone"], type="filepath")
            task_in = gr.Dropdown(list(TASKS), value=DEFAULT_TASK, label="Task")
            user_prompt_in = gr.Textbox(label="Prompt", value=TASKS[DEFAULT_TASK].user_prompt, lines=8, max_lines=12)
            system_prompt_in = gr.Textbox(
                label="System prompt (empty = the model's default)",
                value=TASKS[DEFAULT_TASK].system_prompt,
                lines=3,
                max_lines=8,
            )
            constrain_in = gr.Checkbox(
                value=TASKS[DEFAULT_TASK].constrain, label="Constrain the answer to the task's JSON schema"
            )
            with gr.Accordion("Advanced", open=False):
                audio_first_in = gr.Checkbox(value=True, label="Audio before text")
                max_tokens_in = gr.Slider(16, 1024, value=DEFAULT_MAX_NEW_TOKENS, step=16, label="Max new tokens")
                temperature_in = gr.Slider(0.0, 1.5, value=0.0, step=0.05, label="Temperature (0 = greedy)")
                top_p_in = gr.Slider(0.05, 1.0, value=1.0, step=0.05, label="Top-p (sampling only)")
                rep_in = gr.Slider(1.0, 1.5, value=DEFAULT_REPETITION_PENALTY, step=0.05, label="Repetition penalty")
            with gr.Row():
                submit_btn = gr.Button("Analyze", variant="primary", scale=2)
                reset_btn = gr.Button("Reset", variant="secondary", scale=1)
        with gr.Column(scale=4):
            summary_out = gr.HTML(_message("Upload a clip and press Analyze."))
            with gr.Accordion("JSON", open=False):
                json_out = gr.Code(show_label=False, language="json", wrap_lines=True)
            with gr.Accordion("Raw model output", open=False):
                raw_out = gr.Textbox(show_label=False, lines=6)
    gr.HTML(FOOTER)
    inputs = [
        audio_in,
        system_prompt_in,
        user_prompt_in,
        audio_first_in,
        constrain_in,
        max_tokens_in,
        temperature_in,
        top_p_in,
        rep_in,
        task_in,
    ]
    submit_btn.click(predict, inputs=inputs, outputs=[summary_out, json_out, raw_out], api_name="generate")
    task_in.change(_preset, inputs=task_in, outputs=[system_prompt_in, user_prompt_in, constrain_in], api_name=False)
    reset_btn.click(
        lambda: (
            *_preset(DEFAULT_TASK),
            True,
            DEFAULT_MAX_NEW_TOKENS,
            0.0,
            1.0,
            DEFAULT_REPETITION_PENALTY,
            DEFAULT_TASK,
        ),
        outputs=[system_prompt_in, user_prompt_in, constrain_in, audio_first_in, max_tokens_in]
        + [temperature_in, top_p_in, rep_in, task_in],
        api_name=False,
    )


def _gradio_auth() -> tuple[str, str] | None:
    raw_auth = os.getenv("GRADIO_AUTH", "").strip()
    if not raw_auth:
        return None
    if ":" not in raw_auth:
        raise ValueError("GRADIO_AUTH must be formatted as username:password")
    user, password = raw_auth.split(":", 1)
    return user, password


if __name__ == "__main__":
    demo.queue(default_concurrency_limit=1).launch(
        server_name=os.getenv("GRADIO_SERVER_NAME") or None,
        server_port=int(os.getenv("GRADIO_SERVER_PORT", "7860")),
        auth=_gradio_auth(),
        theme=THEME,
        css=CSS,
    )