import base64
import io
import json
import os

import numpy as np
import scipy.io.wavfile as wavfile
from fastmcp import FastMCP

from pocket_tts import TTSModel
from pocket_tts.default_parameters import get_default_voice_for_language
from pocket_tts.utils.utils import _ORIGINS_OF_PREDEFINED_VOICES

DEFAULT_LANGUAGE = os.environ.get("POCKET_TTS_MCP_LANGUAGE", "spanish_24l")
DEFAULT_TEMP = float(os.environ.get("POCKET_TTS_MCP_TEMP", "0.5"))
DEFAULT_LSD_DECODE_STEPS = int(os.environ.get("POCKET_TTS_MCP_LSD_STEPS", "5"))
DEFAULT_EOS_THRESHOLD = float(os.environ.get("POCKET_TTS_MCP_EOS_THRESHOLD", "-3.0"))
DEFAULT_HOST = os.environ.get("POCKET_TTS_MCP_HOST", "0.0.0.0")
DEFAULT_PORT = int(os.environ.get("POCKET_TTS_MCP_PORT", "8008"))

_model: TTSModel | None = None

mcp = FastMCP(
    name="Pocket TTS MCP",
    instructions=(
        "Generate WAV speech audio with Pocket TTS and inspect the built-in voice presets."
    ),
)


def get_model() -> TTSModel:
    global _model
    if _model is None:
        print("Loading Pocket TTS model for FastMCP server...")
        _model = TTSModel.load_model(
            language=DEFAULT_LANGUAGE,
            temp=DEFAULT_TEMP,
            lsd_decode_steps=DEFAULT_LSD_DECODE_STEPS,
            eos_threshold=DEFAULT_EOS_THRESHOLD,
        )
    return _model


def list_predefined_voices() -> list[dict[str, str | bool]]:
    default_voice = get_default_voice_for_language(DEFAULT_LANGUAGE)
    return [
        {
            "name": name,
            "source": source,
            "is_default": name == default_voice,
        }
        for name, source in sorted(_ORIGINS_OF_PREDEFINED_VOICES.items())
    ]


def encode_wav_base64(audio_tensor, sample_rate: int) -> tuple[str, int]:
    if hasattr(audio_tensor, "detach"):
        audio_numpy = audio_tensor.detach().cpu().numpy()
    else:
        audio_numpy = np.array(audio_tensor)

    if audio_numpy.ndim > 1:
        audio_numpy = audio_numpy.squeeze()

    audio_numpy = np.asarray(audio_numpy, dtype=np.float32)

    buffer = io.BytesIO()
    wavfile.write(buffer, sample_rate, audio_numpy)
    wav_bytes = buffer.getvalue()
    return base64.b64encode(wav_bytes).decode("ascii"), len(wav_bytes)


@mcp.tool()
def list_voices() -> dict:
    """List the built-in Pocket TTS voice presets available to the MCP client."""
    return {
        "language": DEFAULT_LANGUAGE,
        "default_voice": get_default_voice_for_language(DEFAULT_LANGUAGE),
        "voices": list_predefined_voices(),
    }


@mcp.resource("voice://predefined")
def predefined_voices_resource() -> str:
    """Return the predefined voice catalog as JSON."""
    return json.dumps(list_voices(), indent=2)


@mcp.tool()
def generate_speech(
    text: str,
    voice: str | None = None,
    frames_after_eos: int | None = None,
    copy_state: bool = True,
) -> dict:
    """Generate speech audio and return a base64-encoded WAV payload."""
    if not text.strip():
        raise ValueError("text must not be empty")

    model = get_model()
    selected_voice = voice or get_default_voice_for_language(DEFAULT_LANGUAGE)
    voice_state = model.get_state_for_audio_prompt(selected_voice)
    audio_tensor = model.generate_audio(
        model_state=voice_state,
        text_to_generate=text,
        frames_after_eos=frames_after_eos,
        copy_state=copy_state,
    )
    audio_base64, wav_size_bytes = encode_wav_base64(audio_tensor, model.sample_rate)

    return {
        "text": text,
        "voice": selected_voice,
        "language": DEFAULT_LANGUAGE,
        "sample_rate": model.sample_rate,
        "mime_type": "audio/wav",
        "wav_size_bytes": wav_size_bytes,
        "audio_base64": audio_base64,
    }


def main() -> None:
    mcp.run(transport="sse", host=DEFAULT_HOST, port=DEFAULT_PORT)


if __name__ == "__main__":
    main()