from __future__ import annotations

import base64
import io
import pathlib
import shutil
import struct
import sys
import wave
import aiofiles
import aiohttp
from typing import Optional
import asyncio

async def synthesize_and_play_audio(
    input_text: str = "Hello!",
    reference_audio_path: str = "~/Downloads/jmsample.wav",
    model: str = "voxtral-mini-tts-260213",
    api_key: str = "MISTRAL_API_KEY",
    output_path: str = "/tmp/voxtral.wav",
    timeout: float = 180.0,
    raw_sample_rate: int = 24000,
    raw_channels: int = 1,
    no_play: bool = False,
    url: str = "https://api.mistral.ai/v1/audio/text-to-speech",
    reference_format: str = "raw-base64",
) -> int:
    """
    Asynchronously synthesize speech from input text using a reference audio file and play the output.

    Args:
        input_text: Text to synthesize.
        reference_audio_path: Path to the reference WAV file.
        model: Model name sent in the request.
        api_key: API key for authentication.
        output_path: Output audio file path.
        timeout: HTTP timeout in seconds.
        raw_sample_rate: Sample rate used when response audio is raw f32le (non-WAV).
        raw_channels: Channel count used when response audio is raw f32le (non-WAV).
        no_play: If True, only save audio, do not launch a player.
        url: TTS endpoint URL.
        reference_format: How to serialize reference_audio ("data-uri" or "raw-base64").

    Returns:
        int: 0 on success, non-zero on failure.
    """
    print(f"Synthesizing: {input_text!r}")

    try:
        # Use async file operations
        reference_path = pathlib.Path(reference_audio_path).expanduser().resolve()
        if not reference_path.is_file():
            print(f"Reference audio not found: {reference_path}", file=sys.stderr)
            return 2

        # Read reference audio asynchronously
        async with aiofiles.open(reference_path, 'rb') as f:
            reference_bytes = await f.read()

        reference_b64 = base64.b64encode(reference_bytes).decode("ascii")
        if reference_format == "data-uri":
            reference_audio = f"data:audio/wav;base64,{reference_b64}"
        else:
            reference_audio = reference_b64

        payload = {
            "model": model,
            "input": input_text,
            "reference_audio": reference_audio,
            "response_format": "wav",
        }
        headers = {
            "content-type": "application/json",
            "x-api-key": api_key,
        }

        # Use async HTTP client
        async with aiohttp.ClientSession() as session:
            try:
                async with session.post(url, headers=headers, json=payload, timeout=timeout) as response:
                    if response.status >= 400:
                        error_text = await response.text()
                        print(f"Request failed: {response.status}", file=sys.stderr)
                        print(error_text, file=sys.stderr)
                        return 1

                    content_type = response.headers.get("content-type", "")
                    audio_bytes = await response.read()
                    
                    if "application/json" in content_type.lower():
                        body = await response.json()
                        audio_field = body.get("audio") if isinstance(body, dict) else None
                        if not isinstance(audio_field, str) or not audio_field:
                            print("JSON response does not contain an 'audio' field.", file=sys.stderr)
                            print(body, file=sys.stderr)
                            return 1

                        if "," in audio_field and audio_field.startswith("data:"):
                            audio_field = audio_field.split(",", 1)[1]
                        try:
                            audio_bytes = base64.b64decode(audio_field, validate=False)
                        except Exception as exc:
                            print(f"Failed to decode JSON audio field as base64: {exc}", file=sys.stderr)
                            return 1

                    if not _looks_like_wav(audio_bytes):
                        converted = _convert_f32le_to_wav(
                            audio_bytes,
                            sample_rate=raw_sample_rate,
                            channels=raw_channels,
                        )
                        if converted is not None:
                            audio_bytes = converted
                            print(
                                "Response was non-WAV raw audio; converted f32le stream to WAV.",
                                file=sys.stderr,
                            )
                        else:
                            print(
                                "Response bytes are non-WAV and could not be auto-converted.",
                                file=sys.stderr,
                            )

                    # Write output file asynchronously
                    output_path_obj = pathlib.Path(output_path).expanduser()
                    async with aiofiles.open(output_path_obj, 'wb') as f:
                        await f.write(audio_bytes)
                    print(f"Wrote {len(audio_bytes)} bytes to {output_path_obj}")

                    if not no_play:
                        await maybe_play_audio_async(output_path_obj)

                    return 0
                    
            except asyncio.TimeoutError:
                print(f"Request timed out after {timeout} seconds", file=sys.stderr)
                return 1
            except aiohttp.ClientError as e:
                print(f"HTTP client error: {e}", file=sys.stderr)
                return 1
                
    except Exception as e:
        print(f"Unexpected error: {e}", file=sys.stderr)
        return 1

def _looks_like_wav(audio_bytes: bytes) -> bool:
    return len(audio_bytes) >= 12 and audio_bytes[:4] == b"RIFF" and audio_bytes[8:12] == b"WAVE"

def _convert_f32le_to_wav(audio_bytes: bytes, sample_rate: int, channels: int) -> Optional[bytes]:
    if len(audio_bytes) % 4 != 0:
        return None
    if channels <= 0 or sample_rate <= 0:
        return None

    pcm16 = bytearray()
    for (sample,) in struct.iter_unpack("<f", audio_bytes):
        clipped = max(-1.0, min(1.0, sample))
        pcm16.extend(struct.pack("<h", int(clipped * 32767)))

    buffer = io.BytesIO()
    with wave.open(buffer, "wb") as wav_file:
        wav_file.setnchannels(channels)
        wav_file.setsampwidth(2)
        wav_file.setframerate(sample_rate)
        wav_file.writeframes(bytes(pcm16))
    return buffer.getvalue()

async def maybe_play_audio_async(path: pathlib.Path) -> None:
    """Asynchronously play audio using available players."""
    players = (
        ("afplay", [str(path)]),
        ("ffplay", ["-nodisp", "-autoexit", str(path)]),
        ("mpv", [str(path)]),
    )

    for command, extra_args in players:
        if shutil.which(command) is None:
            continue
        try:
            proc = await asyncio.create_subprocess_exec(
                command, *extra_args,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE
            )
            await proc.wait()
            if proc.returncode == 0:
                return
            print(
                f"Player '{command}' failed with exit code {proc.returncode}, trying next player.",
                file=sys.stderr,
            )
        except Exception as e:
            print(f"Error with player '{command}': {e}", file=sys.stderr)
            continue

    print("No local audio player found (afplay/ffplay/mpv).", file=sys.stderr)