# Copyright (c) 2026 DeepBeepMeep - part of WanGP, WanGP Community License 2.0 (see LICENSE.txt)
import os
from importlib.util import find_spec

import torch

from shared.utils import files_locator as fl
from shared.utils.hf import build_hf_url

from .stable_audio3_prompt_enhancers import get_stable_audio3_prompt_enhancer


STABLE_AUDIO3_REPO_ID = "DeepBeepMeep/TTS"
STABLE_AUDIO3_SMALL = "stable_audio3_small"
STABLE_AUDIO3_MEDIUM = "stable_audio3_medium"
STABLE_AUDIO3_TEXT_ENCODER_FOLDER = "t5gemma-b-b-ul2"
STABLE_AUDIO3_TEXT_ENCODER_BF16 = "t5gemma-b-b-ul2_bf16.safetensors"
STABLE_AUDIO3_SMALL_CONFIG = "stable_audio3_small_config.json"
STABLE_AUDIO3_MEDIUM_CONFIG = "stable_audio3_medium_config.json"
STABLE_AUDIO3_SAME_S_WEIGHTS = "stable_audio3_same_s_bf16.safetensors"
STABLE_AUDIO3_SAME_L_WEIGHTS = "stable_audio3_same_l_bf16.safetensors"
STABLE_AUDIO3_TOKENIZER_FILES = [
    "config.json",
    "special_tokens_map.json",
    "tokenizer.json",
    "tokenizer.model",
    "tokenizer_config.json",
]
STABLE_AUDIO3_AUDIO_MODE_CHOICES = ["", "AE", "AI", "AC"]
STABLE_AUDIO3_AUDIO_MODE_LABELS = {
    "": "Text to audio",
    "AE": "Audio to audio edit",
    "AI": "Inpaint source audio",
    "AC": "Continue source audio",
}
STABLE_AUDIO3_SAMPLE_SOLVERS = [
    ("PingPong", "pingpong"),
    ("Euler", "euler"),
    ("DPM++", "dpmpp"),
    ("RK4", "rk4"),
]
STABLE_AUDIO3_CUSTOM_SETTINGS = [
    {"id": "inpaint_start_seconds", "label": "Inpaint Start Seconds", "name": "Inpaint Start Seconds", "type": "float", "default": 0.0, "min": 0.0, "inc": 0.1},
    {"id": "inpaint_end_seconds", "label": "Inpaint End Seconds", "name": "Inpaint End Seconds", "type": "float", "default": 10.0, "min": 0.0, "inc": 0.1},
]


def _root_checkpoint_path(filename):
    return fl.locate_file(filename, error_if_none=False) or filename


def _config_path(base_model_type):
    return os.path.join(os.path.dirname(__file__), "stable_audio3", "configs", _config_name(base_model_type))


def _asset_weights(base_model_type):
    return STABLE_AUDIO3_SAME_L_WEIGHTS if base_model_type == STABLE_AUDIO3_MEDIUM else STABLE_AUDIO3_SAME_S_WEIGHTS


def _config_name(base_model_type):
    return STABLE_AUDIO3_MEDIUM_CONFIG if base_model_type == STABLE_AUDIO3_MEDIUM else STABLE_AUDIO3_SMALL_CONFIG


def _max_duration(base_model_type):
    return 380 if base_model_type == STABLE_AUDIO3_MEDIUM else 120


def _duration_slider(base_model_type):
    return {"label": "Duration (seconds)", "name": "Target Duration", "min": 1, "max": _max_duration(base_model_type), "increment": 1, "default": 30}


def _custom_settings(base_model_type):
    return [dict(one, max=float(_max_duration(base_model_type))) for one in STABLE_AUDIO3_CUSTOM_SETTINGS]


def _mode_from_audio_prompt_type(audio_prompt_type):
    audio_prompt_type = str(audio_prompt_type or "").upper()
    if "A" not in audio_prompt_type:
        return "text"
    if "E" in audio_prompt_type:
        return "audio_to_audio"
    if "I" in audio_prompt_type:
        return "inpaint"
    if "C" in audio_prompt_type:
        return "continue"
    return "text"


def _flash_attention2_available():
    if find_spec("flash_attn") is None:
        return False
    try:
        import flash_attn
    except Exception:
        return False
    return callable(getattr(flash_attn, "flash_attn_varlen_func", None))


def _medium_flash_attention_error():
    return "Stable Audio 3 Medium requires Flash Attention 2 for SAME-L sliding-window attention. Install flash-attn or use Stable Audio 3 Small Music/SFX."


def _require_medium_flash_attention(base_model_type):
    if base_model_type == STABLE_AUDIO3_MEDIUM and not _flash_attention2_available():
        raise RuntimeError(_medium_flash_attention_error())


def _model_id(base_model_type, model_def):
    default_model_id = "medium" if base_model_type == STABLE_AUDIO3_MEDIUM else "small-music"
    return str((model_def or {}).get("stable_audio3_model_id") or default_model_id)


class family_handler:
    @staticmethod
    def query_supported_types():
        return [STABLE_AUDIO3_SMALL, STABLE_AUDIO3_MEDIUM]

    @staticmethod
    def query_family_maps():
        return {}, {}

    @staticmethod
    def query_model_family():
        return "tts"

    @staticmethod
    def query_family_infos():
        return {"music": (2195, "Music"), "tts": (2200, "TTS")}

    @staticmethod
    def get_lora_dir(base_model_type):
        if base_model_type == STABLE_AUDIO3_MEDIUM:
            return STABLE_AUDIO3_MEDIUM
        return STABLE_AUDIO3_SMALL

    @staticmethod
    def query_model_def(base_model_type, model_def):
        model_id = _model_id(base_model_type, model_def)
        prompt_enhancer_instructions, prompt_enhancer_max_tokens, prompt_enhancer_button_label = get_stable_audio3_prompt_enhancer(model_id)
        speciality, example = {
            "small-music": ("instrumental music, loops and grooves", "A relaxed 90 BPM instrumental with warm electric piano, rounded bass and brushed drums; a gentle introduction builds into a steady groove, then resolves softly."),
            "small-sfx": ("sound effects, foley and ambience", "A heavy wooden door creaks open close to the listener, followed by footsteps across a stone floor and a short reverberant door slam."),
            "medium": ("longer music, evolving ambience and layered sound design", "A slow cinematic instrumental with low strings and soft piano, gradually joined by distant thunder and wind, building to a restrained climax before fading."),
        }[model_id]
        input_help = (
            f"Stable Audio 3 is suited to {speciality}. It produces stereo audio at 44.1 kHz. "
            f"**Duration** sets the target output length, from 1 to {_max_duration(base_model_type)} seconds; audible sound may finish earlier.\n\n"
            "### Choose a Workflow\n\n"
            "- **Text to Audio:** describe a new sound or piece of music.\n"
            "- **Edit Audio:** supply source audio and describe the desired change. Audio strength 0 retains more source detail; 1 changes it more.\n"
            "- **Inpaint Audio:** supply source audio and select the start/end seconds of the interval to regenerate.\n"
            "- **Continue Audio:** supply source audio and choose a total duration longer than the source. The duration includes the original section.\n\n"
            "Inpainting and continuation may also change the sound slightly outside the regenerated region; preservation is not sample-exact.\n\n"
            "Start with 8 inference steps, guidance 1 and the PingPong solver. Prompt enhancement is optional. "
            "This model is not designed for speech or voice generation, or transcription. Inspect the result for silence and level before combining it with video."
        )
        if base_model_type == STABLE_AUDIO3_MEDIUM:
            input_help += " Medium requires Flash Attention 2 for its audio codec. WanGP allows up to 380 seconds, but the upstream Medium checkpoint is documented for about 4.75 minutes; prefer that range."
        prompt_help = (
            f"Describe the sound in English, choosing details relevant to {speciality}. The model was trained on English descriptions; other languages may underperform. "
            "For music, specify genre, tempo or BPM, instruments, mood, production texture and a simple progression. "
            "For effects, specify the source, action, material, distance, environment and how the sound develops. "
            "For source editing, describe the desired resulting sound and the intended change. "
            "Use the duration control for timing; prose alone does not set the output length. "
            "Exact lyrics, spoken wording, seamless loops and precisely timed events are not guaranteed.\n\n"
            f"Example: {example}"
        )
        return {
            "group": "music",
            "infos": "## Stable Audio 3\n\n" + input_help,
            "prompt_infos": "## Describe the Sound\n\n" + prompt_help,
            "specialities": [{"name": name} for name in ("audio editing", "audio inpainting", "audio continuation", *(("sound effects",) if model_id != "small-music" else ()), *(("music generation",) if model_id != "small-sfx" else ()))],
            "deepy_infos": f"Stereo 44.1 kHz; {speciality}; target 1–{_max_duration(base_model_type)}s. Modes audio_prompt_type: empty=text; AE=edit audio_guide (audio_scale 0..1, higher changes more); AI=inpaint audio_guide using custom_settings.inpaint_start_seconds/end_seconds; AC=continue audio_guide, duration_seconds is total length and must exceed the source. Codec reconstruction is not sample-exact. Start 8 steps, guidance 1, pingpong. Not designed for speech/voice generation or transcription; inspect silence/levels. " + ("Medium needs Flash Attention 2. WanGP allows 380s, but upstream documents Medium for about 285s; prefer that range." if base_model_type == STABLE_AUDIO3_MEDIUM else ""),
            "deepy_prompt_infos": f"English description of {speciality}: audible sources/instruments, rhythm, mood, texture, space and development. English-trained; other languages may underperform. Describe the target change when editing. Set timing with duration_seconds. Exact words, event timing and seamless loops are not guaranteed. Example: {example}",
            "device_explicit": True,
            "audio_only": True,
            "image_outputs": False,
            "sliding_window": False,
            "guidance_max_phases": 1,
            "image_prompt_types_allowed": "",
            "supports_early_stop": True,
            "profiles_dir": [base_model_type],
            "text_encoder_URLs": [build_hf_url(STABLE_AUDIO3_REPO_ID, STABLE_AUDIO3_TEXT_ENCODER_FOLDER, STABLE_AUDIO3_TEXT_ENCODER_BF16)],
            "text_encoder_folder": STABLE_AUDIO3_TEXT_ENCODER_FOLDER,
            "text_prompt_enhancer_instructions": prompt_enhancer_instructions,
            "text_prompt_enhancer_max_tokens": prompt_enhancer_max_tokens,
            "prompt_enhancer_button_label": prompt_enhancer_button_label,
            "prompt_enhancer_choices_allowed": ["T"],
            "inference_steps": True,
            "sample_solvers": STABLE_AUDIO3_SAMPLE_SOLVERS,
            "temperature": False,
            "any_audio_prompt": True,
            "audio_prompt_choices": True,
            "enabled_audio_lora": True,
            "audio_guide_label": "Source audio",
            "audio_scale_name": "Edit Noise Level",
            "audio_prompt_type_sources": {
                "selection": STABLE_AUDIO3_AUDIO_MODE_CHOICES,
                "labels": STABLE_AUDIO3_AUDIO_MODE_LABELS,
                "default": "",
                "label": "Source Audio",
                "letters_filter": "AEIC",
            },
            "duration_slider": _duration_slider(base_model_type),
            "custom_settings": _custom_settings(base_model_type),
            "compile": False,
        }

    @staticmethod
    def query_model_files(computeList, base_model_type, model_def=None):
        return {
            "repoId": STABLE_AUDIO3_REPO_ID,
            "sourceFolderList": ["", STABLE_AUDIO3_TEXT_ENCODER_FOLDER],
            "fileList": [[_asset_weights(base_model_type)], STABLE_AUDIO3_TOKENIZER_FILES],
        }

    @staticmethod
    def load_model(
        model_filename,
        model_type,
        base_model_type,
        model_def,
        quantizeTransformer=False,
        text_encoder_quantization=None,
        dtype=None,
        VAE_dtype=None,
        mixed_precision_transformer=False,
        save_quantized=False,
        submodel_no_list=None,
        text_encoder_filename=None,
        profile=0,
        **kwargs,
    ):
        _require_medium_flash_attention(base_model_type)

        from .stable_audio3.pipeline import StableAudio3Pipeline

        transformer_weights = model_filename[0] if isinstance(model_filename, (list, tuple)) else model_filename
        config_path = _config_path(base_model_type)
        autoencoder_weights = _root_checkpoint_path(_asset_weights(base_model_type))
        tokenizer_dir = fl.locate_folder(STABLE_AUDIO3_TEXT_ENCODER_FOLDER)
        pipeline = StableAudio3Pipeline(
            transformer_weights,
            config_path,
            autoencoder_weights,
            text_encoder_filename,
            tokenizer_dir,
            model_id=_model_id(base_model_type, model_def),
            max_duration=_max_duration(base_model_type),
            dtype=dtype or torch.bfloat16,
        )

        prompt_conditioner = pipeline.model.conditioner.conditioners["prompt"] if "prompt" in pipeline.model.conditioner.conditioners else None
        pipe = {
            "transformer": pipeline.main_model,
            "codec": pipeline.model.pretransform,
        }
        if prompt_conditioner is not None and hasattr(prompt_conditioner, "model"):
            pipe["text_encoder"] = prompt_conditioner.model

        if save_quantized and transformer_weights:
            from wgp import save_quantized_model

            save_quantized_model(pipeline.main_model, model_type, transformer_weights, dtype or torch.bfloat16, config_path)

        return pipeline, {"pipe": pipe, "coTenantsMap": {}}

    @staticmethod
    def fix_settings(base_model_type, settings_version, model_def, ui_defaults):
        ui_defaults.setdefault("audio_prompt_type", "")
        ui_defaults.setdefault("sample_solver", "pingpong")

    @staticmethod
    def update_default_settings(base_model_type, model_def, ui_defaults):
        ui_defaults.update(
            {
                "audio_prompt_type": "",
                "prompt": "An anthemic pop rock instrumental with bright guitars, punchy drums, and a nostalgic festival chorus.",
                "duration_seconds": 30,
                "repeat_generation": 1,
                "video_length": 0,
                "num_inference_steps": 8,
                "guidance_scale": 1.0,
                "negative_prompt": "poor quality, distorted, noisy",
                "sample_solver": "pingpong",
                "audio_scale": 0.9,
                "multi_prompts_gen_type": "FG",
            }
        )

    @staticmethod
    def validate_generative_prompt(base_model_type, model_def, inputs, one_prompt):
        if one_prompt is None or len(str(one_prompt).strip()) == 0:
            return "Prompt text cannot be empty for Stable Audio 3."
        return None

    @staticmethod
    def validate_generative_settings(base_model_type, model_def, inputs):
        if base_model_type == STABLE_AUDIO3_MEDIUM and not _flash_attention2_available():
            return _medium_flash_attention_error()
        audio_prompt_type = inputs.get("audio_prompt_type", "") or ""
        mode = _mode_from_audio_prompt_type(audio_prompt_type)
        if mode in ("audio_to_audio", "inpaint", "continue") and inputs.get("audio_guide") is None:
            return "Stable Audio 3 source-audio modes require a source audio file."
        if mode == "audio_to_audio" and inputs.get("audio_scale") is not None:
            try:
                inputs["audio_scale"] = float(inputs["audio_scale"])
            except (TypeError, ValueError):
                return "Stable Audio 3 Edit Noise Level must be a number."
            if not 0 <= inputs["audio_scale"] <= 1:
                return "Stable Audio 3 Edit Noise Level must be between 0 and 1."
        custom_settings = inputs.get("custom_settings", None)
        if isinstance(custom_settings, dict):
            for key in ("inpaint_start_seconds", "inpaint_end_seconds"):
                value = custom_settings.get(key, None)
                if value is None or value == "":
                    continue
                try:
                    custom_settings[key] = float(value)
                except (TypeError, ValueError):
                    return f"Stable Audio 3 custom setting '{key}' must be a number."
        return None