from __future__ import annotations

import os
import time
import urllib.request

import gradio as gr

from pyharp import *
from gradio_client import Client, handle_file


_BACKEND_SPACE = "ACE-Step/ACE-Step"
_BACKEND_API_NAME = "/__call__"
_BACKEND_TOKEN_ENV = "HF_TOKEN"
_ACCEPT_USER_TOKEN = False
# How many times to wake+retry a sleeping backend, and how long to wait for
# it to boot (a free Space cold start can take a few minutes).
_CALL_RETRIES = int(os.environ.get("BACKEND_CALL_RETRIES", "4"))
_WAKE_TIMEOUT = float(os.environ.get("BACKEND_WAKE_TIMEOUT", "420"))
_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 _reset_client():
    # Drop the cached connection so the next attempt reconnects to a Space
    # that has since finished waking.
    global _client
    _client = None


def _make_conn(tok):
    tok = (tok or '').strip()
    if tok:
        return Client(_BACKEND_SPACE, hf_token=tok)
    return _backend_client()


def _space_url(space):
    slug = space.strip().lower().replace('/', '-').replace('_', '-')
    return f'https://{slug}.hf.space/'


def _is_cold_start(message):
    # Errors that mean 'the backend was asleep/booting', worth waking+retrying
    # (vs. a real application error, which we surface immediately).
    _low = (message or '').lower()
    return any(s in _low for s in (
        'read operation timed out', 'timed out', 'timeout', 'starting',
        'building', 'not ready', 'no application', 'connection', '503', '502',
    ))


def _wake_backend():
    # A sleeping Space boots when its URL is hit; poll until it answers (or
    # the budget expires) so the retried call lands on a running backend.
    _url = _space_url(_BACKEND_SPACE)
    _deadline = time.time() + _WAKE_TIMEOUT
    _delay = 5.0
    while time.time() < _deadline:
        try:
            _req = urllib.request.Request(_url, headers={'User-Agent': 'harp-frontend'})
            with urllib.request.urlopen(_req, timeout=30) as _resp:
                if getattr(_resp, 'status', 200) < 500:
                    return True
        except Exception:
            pass
        time.sleep(_delay)
        _delay = min(_delay * 1.5, 30.0)
    return False


def _quota_hint(message):
    # Turn a backend 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."
        )
    # Opaque backend failure: the Space raised an exception it refuses to
    # expose (it runs with show_error=False), so all we get is a generic
    # 'Internal Gradio error'. The frontend can't fix a server-side crash --
    # point the user at where the real cause lives.
    if (
        "internal gradio error" in _low
        or "internal server error" in _low
        or "apperror" in _low
        or _low.strip() in ("", "none")
    ):
        _hint = (
            "The backend Space raised an error it did not expose (it runs with "
            "show_error disabled), so the real cause is only in the backend "
            "Space's Logs tab (" + _BACKEND_SPACE + ")."
        )
        if _ACCEPT_USER_TOKEN:
            _hint += (
                " If it is a ZeroGPU Space, an anonymous call can fail this way -- "
                "paste a Hugging Face token in the token field and retry."
            )
        return _hint
    return message or "Backend call failed."


model_card = ModelCard(
    name="Ace Step",
    description="TODO: describe this model.",
    author="ACE-Step",
    tags=[],
)


def process_fn(audio_duration, prompt, lyrics, infer_step, guidance_scale, omega_scale, manual_seeds, guidance_interval, guidance_interval_decay, min_guidance_scale, use_erg_tag, use_erg_lyric, use_erg_diffusion, oss_steps, guidance_scale_text, guidance_scale_lyric, audio2audio_enable, ref_audio_strength, ref_audio_input, lora_name_or_path):
    _tok = ''
    # Call the backend, waking it and retrying if it was asleep (a cold
    # start otherwise fails the first hit with 'read operation timed out').
    _raw = None
    for _attempt in range(_CALL_RETRIES + 1):
        try:
            _conn = _make_conn(_tok)
            _raw = _conn.predict(
                audio_duration,
                prompt,
                lyrics,
                infer_step,
                guidance_scale,
                'euler',
                'apg',
                omega_scale,
                manual_seeds,
                guidance_interval,
                guidance_interval_decay,
                min_guidance_scale,
                use_erg_tag,
                use_erg_lyric,
                use_erg_diffusion,
                oss_steps,
                guidance_scale_text,
                guidance_scale_lyric,
                audio2audio_enable,
                ref_audio_strength,
                (handle_file(ref_audio_input) if ref_audio_input else None),
                lora_name_or_path,
                api_name="/__call__",
            )
            break
        except Exception as _exc:  # never surfaces the token
            if _attempt < _CALL_RETRIES and _is_cold_start(str(_exc)):
                _reset_client()
                _wake_backend()
                continue
            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_text2music_generated_audio_1 = _values[0] if len(_values) > 0 else None
    if not _out_text2music_generated_audio_1:
        raise gr.Error(_detail or "The backend Space returned no 'text2music_generated_audio_1' output. Check the backend Space's logs; if it uses ZeroGPU it may need a moment to warm up.")
    return _out_text2music_generated_audio_1


with gr.Blocks() as demo:
    input_components = [
        gr.Slider(minimum=-1, maximum=240.0, step=1e-05, value=-1, label="Audio Duration"),
        gr.Textbox(label="Tags", value="funk, pop, soul, rock, melodic, guitar, drums, bass, keyboard, percussion, 105 BPM, energetic, upbeat, groovy, vibrant, dynamic"),
        gr.Textbox(label="Lyrics", value="[verse]\nNeon lights they flicker bright\nCity hums in dead of night\nRhythms pulse through concrete veins\nLost in echoes of refrains\n\n[verse]\nBassline groovin' in my chest\nHeartbeats match the city's zest\nElectric whispers fill the air\nSynthesized dreams everywhere\n\n[chorus]\nTurn it up and let it flow\nFeel the fire let it grow\nIn this rhythm we belong\nHear the night sing out our song\n\n[verse]\nGuitar strings they start to weep\nWake the soul from silent sleep\nEvery note a story told\nIn this night we\u2019re bold and gold\n\n[bridge]\nVoices blend in harmony\nLost in pure cacophony\nTimeless echoes timeless cries\nSoulful shouts beneath the skies\n\n[verse]\nKeyboard dances on the keys\nMelodies on evening breeze\nCatch the tune and hold it tight\nIn this moment we take flight\n"),
        gr.Slider(minimum=1, maximum=200, step=1, value=60, label="Infer Steps"),
        gr.Slider(minimum=0.0, maximum=30.0, step=0.1, value=15.0, label="Guidance Scale"),
        gr.Slider(minimum=-100.0, maximum=100.0, step=0.1, value=10.0, label="Granularity Scale"),
        gr.Textbox(label="manual seeds (default None)"),
        gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=0.5, label="Guidance Interval"),
        gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=0.0, label="Guidance Interval Decay"),
        gr.Slider(minimum=0.0, maximum=200.0, step=0.1, value=3.0, label="Min Guidance Scale"),
        gr.Checkbox(value=True, label="use ERG for tag"),
        gr.Checkbox(value=False, label="use ERG for lyric"),
        gr.Checkbox(value=True, label="use ERG for diffusion"),
        gr.Textbox(label="OSS Steps"),
        gr.Slider(minimum=0.0, maximum=10.0, step=0.1, value=0.0, label="Guidance Scale Text"),
        gr.Slider(minimum=0.0, maximum=10.0, step=0.1, value=0.0, label="Guidance Scale Lyric"),
        gr.Checkbox(value=False, label="Enable Audio2Audio"),
        gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=0.5, label="Refer audio strength"),
        gr.Audio(type="filepath", label="Reference Audio (for Audio2Audio)"),
        gr.Dropdown(choices=["ACE-Step/ACE-Step-v1-chinese-rap-LoRA", "none"], value="none", label="Lora Name or Path"),
    ]
    output_components = [
        gr.Audio(type="filepath", label="Text2Music Generated Audio 1"),
    ]
    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)