from __future__ import annotations

import os

import gradio as gr

from pyharp import *
from gradio_client import Client, handle_file


_BACKEND_SPACE = "facebook/MelodyFlow"
_BACKEND_API_NAME = "/predict"
_BACKEND_TOKEN_ENV = "HF_TOKEN"
_ACCEPT_USER_TOKEN = True
_client = None


def _backend_client():
    # Lazily create and cache one warm connection using this Space's own
    # token (from the HF_TOKEN secret) or anonymous if none is set. User
    # tokens are NOT cached here -- they get a fresh per-call connection.
    global _client
    if _client is None:
        _token = os.environ.get(_BACKEND_TOKEN_ENV) or None
        _client = Client(_BACKEND_SPACE, hf_token=_token)
    return _client


def _quota_hint(message):
    # Turn a backend ZeroGPU quota error into an actionable message.
    # NOTE: 'message' is the backend's error text; it never contains our token.
    _low = (message or "").lower()
    if "quota" in _low or "zerogpu" in _low:
        if _ACCEPT_USER_TOKEN:
            return (
                "The backend's ZeroGPU quota is exhausted for the identity making "
                "this call. Paste your own Hugging Face token in the token field "
                "(read scope) so usage is attributed to your account."
            )
        return (
            "The backend's ZeroGPU quota is exhausted. This Space's calls are "
            "anonymous unless an HF_TOKEN secret is set (Settings -> Variables "
            "and secrets); use a token from a PRO account or a ZeroGPU-enabled org."
        )
    return message or "Backend call failed."


model_card = ModelCard(
    name="Melodyflow",
    description="TODO: describe this model.",
    author="facebook",
    tags=[],
)


def process_fn(text, steps, target_flowstep, regularize, regularization_strength, duration, melody, _hf_user_token=''):
    _tok = (_hf_user_token or '').strip()
    if _tok:
        _conn = Client(_BACKEND_SPACE, hf_token=_tok)
    else:
        _conn = _backend_client()
    try:
        _raw = _conn.predict(
            'facebook/melodyflow-t24-30secs',
            text,
            'midpoint',
            steps,
            target_flowstep,
            regularize,
            regularization_strength,
            duration,
            handle_file(melody),
            api_name="/predict",
        )
    except Exception as _exc:  # surface a token-aware hint, never the token
        raise gr.Error(_quota_hint(str(_exc)))
    _values = list(_raw) if isinstance(_raw, (list, tuple)) else [_raw]
    _detail = " | ".join(str(_v) for _v in _values if isinstance(_v, str) and _v.strip())
    _out_generated_audio_variation_1 = _values[0] if len(_values) > 0 else None
    if not _out_generated_audio_variation_1:
        raise gr.Error(_detail or "The backend Space returned no 'generated_audio_variation_1' output. Check the backend Space's logs; if it uses ZeroGPU it may need a moment to warm up.")
    _out_generated_audio_variation_2 = _values[1] if len(_values) > 1 else None
    if not _out_generated_audio_variation_2:
        raise gr.Error(_detail or "The backend Space returned no 'generated_audio_variation_2' output. Check the backend Space's logs; if it uses ZeroGPU it may need a moment to warm up.")
    _out_generated_audio_variation_3 = _values[2] if len(_values) > 2 else None
    if not _out_generated_audio_variation_3:
        raise gr.Error(_detail or "The backend Space returned no 'generated_audio_variation_3' output. Check the backend Space's logs; if it uses ZeroGPU it may need a moment to warm up.")
    return _out_generated_audio_variation_1, _out_generated_audio_variation_2, _out_generated_audio_variation_3


with gr.Blocks() as demo:
    input_components = [
        gr.Textbox(label="Input Text"),
        gr.Slider(minimum=0.0, maximum=1.0, step=0.1, value=128.0, label="Inference steps"),
        gr.Slider(minimum=0.0, maximum=1.0, step=0.1, value=0.0, label="Target Flow step"),
        gr.Checkbox(value=False, label="Regularize"),
        gr.Slider(minimum=0.0, maximum=1.0, step=0.1, value=0.2, label="Regularization Strength"),
        gr.Slider(minimum=0.0, maximum=1.0, step=0.1, value=30.0, label="Duration"),
        gr.Audio(type="filepath", label="File or Microphone"),
        gr.Textbox(label="Hugging Face token (optional)", type="password", info="Optional. Paste a Hugging Face token (Settings -> Access Tokens, read scope) so ZeroGPU usage on the backend is charged to YOUR account. Used only for this call; not stored. Leave blank to use this Space's own token."),
    ]
    output_components = [
        gr.Audio(type="filepath", label="Generated Audio - variation 1"),
        gr.Audio(type="filepath", label="Generated Audio - variation 2"),
        gr.Audio(type="filepath", label="Generated Audio - variation 3"),
    ]
    build_endpoint(
        model_card=model_card,
        input_components=input_components,
        output_components=output_components,
        process_fn=process_fn,
    )

demo.queue().launch(share=True, show_error=False, pwa=True)