from __future__ import annotations

import asyncio
from array import array
import base64
import os
from typing import Optional
from urllib.parse import urlparse, urlunparse

try:
    from dotenv import load_dotenv
except Exception:
    load_dotenv = None

if load_dotenv:
    _dotenv_path = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", ".env"))
    load_dotenv(dotenv_path=_dotenv_path)

from azure.ai.voicelive.aio import connect
from azure.ai.voicelive.models import (
    AudioFormat,
    AzureStandardVoice,
    Modality,
    RequestSession,
    ServerEventType,
    ServerVad,
)
from azure.core.credentials import AzureKeyCredential

from app.observability.logging import get_logger


def _voicelive_settings() -> tuple[str, str, str, str, str, Optional[str]]:
    endpoint = (
        os.getenv("AZURE_VOICELIVE_ENDPOINT")
        or os.getenv("VOICELIVE_ENDPOINT")
        or os.getenv("SPEECH_ENDPOINT")
    )
    region = os.getenv("AZURE_VOICELIVE_REGION") or os.getenv("SPEECH_REGION")
    api_key = os.getenv("AZURE_VOICELIVE_API_KEY") or os.getenv("SPEECH_KEY")
    api_version = os.getenv("AZURE_VOICELIVE_API_VERSION") or os.getenv("VOICELIVE_API_VERSION") or "2025-05-01-preview"
    model = os.getenv("VOICELIVE_MODEL", "gpt-4o-realtime-preview")
    voice = os.getenv("VOICELIVE_VOICE", os.getenv("SPEECH_VOICE", "en-US-JennyNeural"))
    rate = (
        os.getenv("VOICELIVE_RATE")
        or os.getenv("VOICELIVE_VOICE_RATE")
        or os.getenv("SPEECH_RATE")
    )
    if not api_key:
        raise RuntimeError("Missing AZURE_VOICELIVE_API_KEY (or SPEECH_KEY) for VoiceLive.")
    endpoint = _normalize_voicelive_endpoint(endpoint, region=region)
    if rate is not None:
        rate = rate.strip() or None
    return endpoint, api_key, model, voice, api_version, rate


def _normalize_voicelive_endpoint(raw: str | None, *, region: str | None) -> str:
    if not raw:
        if region:
            raw = region
        else:
            raise RuntimeError("Missing AZURE_VOICELIVE_ENDPOINT (or AZURE_VOICELIVE_REGION).")
    endpoint = raw.strip()
    if not endpoint:
        raise RuntimeError("Empty AZURE_VOICELIVE_ENDPOINT provided.")
    if "://" not in endpoint:
        if "." not in endpoint and "/" not in endpoint:
            endpoint = f"https://{endpoint}.api.cognitive.microsoft.com"
        else:
            endpoint = f"https://{endpoint}"
    parsed = urlparse(endpoint)
    scheme = parsed.scheme.lower()
    if scheme == "wss":
        scheme = "https"
    elif scheme == "ws":
        scheme = "http"

    path = parsed.path or ""
    lower_path = path.lower().rstrip("/")
    for suffix in ("/voice-agent/realtime", "/voice-live/realtime", "/voice-agent", "/voice-live"):
        if lower_path.endswith(suffix):
            path = path[: -len(suffix)]
            break
    path = path.rstrip("/")
    return urlunparse((scheme, parsed.netloc, path, "", parsed.query, ""))


def _int_env(name: str, default: int) -> int:
    value = os.getenv(name)
    if not value:
        return default
    try:
        parsed = int(value)
    except Exception:
        return default
    return parsed if parsed > 0 else default


def _resample_pcm16_mono(audio: bytes, src_rate: int, dst_rate: int) -> bytes:
    if not audio or src_rate <= 0 or dst_rate <= 0 or src_rate == dst_rate:
        return audio
    if len(audio) < 4:
        return audio
    if len(audio) % 2 == 1:
        audio = audio[:-1]

    samples = array("h")
    samples.frombytes(audio)
    n_in = len(samples)
    if n_in < 2:
        return audio

    n_out = max(1, int(n_in * (dst_rate / src_rate)))
    out = array("h", [0]) * n_out

    step = src_rate / dst_rate
    pos = 0.0
    for i in range(n_out):
        j = int(pos)
        if j >= n_in - 1:
            out[i] = samples[-1]
            break
        frac = pos - j
        s0 = samples[j]
        s1 = samples[j + 1]
        v = int(round(s0 + (s1 - s0) * frac))
        if v > 32767:
            v = 32767
        elif v < -32768:
            v = -32768
        out[i] = v
        pos += step

    return out.tobytes()


def _maybe_resample_voicelive_tts(audio: bytes) -> bytes:
    # The WebUI playback is hard-coded to 16kHz PCM16. If the VoiceLive service emits
    # PCM16 at a higher sample rate (commonly 24kHz), playing it at 16kHz sounds deep/slow.
    disable = (os.getenv("VOICELIVE_DISABLE_RESAMPLE") or "").strip().lower()
    if disable in ("1", "true", "yes", "on"):
        return audio

    dst_rate = _int_env("VOICELIVE_CLIENT_SAMPLE_RATE", 16000)
    src_rate = _int_env("VOICELIVE_TTS_SOURCE_SAMPLE_RATE", 24000)
    if src_rate == dst_rate:
        return audio
    return _resample_pcm16_mono(audio, src_rate, dst_rate)


async def run_voicelive_session(websocket, principal) -> None:
    logger = get_logger("voicelive").bind(sid=principal.get("sid"), user=principal.get("sub"))
    endpoint, api_key, model, voice_name, api_version, voice_rate = _voicelive_settings()
    credential = AzureKeyCredential(api_key)

    try:
        async with connect(
            endpoint=endpoint,
            credential=credential,
            model=model,
            api_version=api_version,
            connection_options={
                "max_msg_size": 10 * 1024 * 1024,
                "heartbeat": 20,
                "timeout": 20,
            },
        ) as conn:
            await _configure_session(conn, voice_name, voice_rate)
            await websocket.send_json({"type": "session.ready", "engine": "voicelive"})
            
            # Send welcome message immediately after session is ready
            await _send_welcome_message(conn, websocket)

            async def client_loop() -> None:
                while True:
                    message = await websocket.receive()
                    if message["type"] == "websocket.disconnect":
                        break
                    if "bytes" in message and message["bytes"] is not None:
                        audio_b64 = base64.b64encode(message["bytes"]).decode("ascii")
                        await conn.input_audio_buffer.append(audio=audio_b64)
                        continue
                    if "text" in message and message["text"] is not None:
                        payload = await _safe_json(message["text"])
                        if payload is None:
                            continue
                        event_type = payload.get("type")
                        if event_type == "session.end":
                            break
                        if event_type == "control.interrupt":
                            try:
                                await conn.response.cancel()
                            except Exception:
                                pass
                            continue
                        if event_type == "audio.frame":
                            audio_b64 = payload.get("audio")
                            if audio_b64:
                                await conn.input_audio_buffer.append(audio=audio_b64)
                            continue
                try:
                    await websocket.close()
                except Exception:
                    return

            async def server_loop() -> None:
                speaking = False
                async for event in conn:
                    etype = getattr(event, "type", None)
                    if etype == ServerEventType.SESSION_UPDATED:
                        continue
                    if etype == ServerEventType.INPUT_AUDIO_BUFFER_SPEECH_STARTED:
                        await websocket.send_json({"type": "agent.state", "state": "listening"})
                        continue
                    if etype == ServerEventType.RESPONSE_CREATED:
                        await websocket.send_json({"type": "agent.state", "state": "thinking"})
                        continue
                    if etype == ServerEventType.RESPONSE_AUDIO_DELTA:
                        if not speaking:
                            speaking = True
                            await websocket.send_json({"type": "agent.state", "state": "speaking"})
                        audio = getattr(event, "delta", None)
                        if audio:
                            audio = _maybe_resample_voicelive_tts(audio)
                            b64 = base64.b64encode(audio).decode("ascii")
                            await websocket.send_json({"type": "tts.audio.chunk", "audio": b64})
                        continue
                    if etype == ServerEventType.RESPONSE_AUDIO_DONE:
                        if speaking:
                            speaking = False
                            await websocket.send_json({"type": "agent.state", "state": "listening"})
                        continue
                    if etype == ServerEventType.RESPONSE_TEXT_DELTA:
                        text_delta = getattr(event, "delta", None)
                        if text_delta:
                            await websocket.send_json({"type": "agent.text.delta", "text": text_delta})
                        continue
                    if etype == ServerEventType.RESPONSE_DONE:
                        text = getattr(getattr(event, "response", None), "text", None)
                        if text:
                            await websocket.send_json({"type": "agent.text.final", "text": text})
                        continue
                    if etype == ServerEventType.ERROR:
                        msg = getattr(getattr(event, "error", None), "message", "VoiceLive error")
                        logger.error("voicelive_error", error=msg)
                        await websocket.send_json({"type": "agent.error", "detail": msg})

            await asyncio.gather(client_loop(), server_loop())
    except Exception as exc:
        logger.error("voicelive_connect_failed", error=str(exc), endpoint=endpoint)
        try:
            await websocket.send_json(
                {"type": "agent.error", "detail": "VoiceLive connection failed. Check AZURE_VOICELIVE_ENDPOINT."}
            )
            await websocket.close()
        except Exception:
            return


async def _configure_session(conn, voice_name: str, voice_rate: Optional[str]) -> None:
    voice_config: AzureStandardVoice | str
    if "-" in voice_name:
        voice_kwargs: dict[str, object] = {"name": voice_name, "type": "azure-standard"}
        if voice_rate:
            voice_kwargs["rate"] = voice_rate
        # Use standard AriaNeural without style modifications that can cause voice distortion
        voice_config = AzureStandardVoice(**voice_kwargs)
    else:
        voice_config = voice_name
    session_config = RequestSession(
        modalities=[Modality.TEXT, Modality.AUDIO],
        instructions=_instructions(),
        voice=voice_config,
        input_audio_format=AudioFormat.PCM16,
        output_audio_format=AudioFormat.PCM16,
        turn_detection=ServerVad(threshold=0.5, prefix_padding_ms=200, silence_duration_ms=350),
    )
    await conn.session.update(session=session_config)


async def _send_welcome_message(conn, websocket) -> None:
    """Send an initial welcome message when the session starts."""
    try:
        welcome_text = _welcome_message()
        # Create a conversation item with the welcome message
        await conn.conversation.item.create(
            item={
                "type": "message", 
                "role": "assistant",
                "content": [{"type": "text", "text": welcome_text}]
            }
        )
        # Create a response to speak the welcome message
        await conn.response.create()
    except Exception as exc:
        # If welcome message fails, just log and continue
        logger = get_logger("voicelive")
        logger.warning("welcome_message_failed", error=str(exc))


def _instructions() -> str:
    return os.getenv(
        "VOICELIVE_INSTRUCTIONS",
        "You are an elite sales professional with a warm, trustworthy voice. Be concise, confident, and customer-focused. "
        "Use a friendly yet professional tone that builds rapport. Ask at most one question and keep responses under 2 sentences unless the user asks otherwise. "
        "Speak with enthusiasm about solutions and show genuine interest in helping the customer succeed.",
    )


def _welcome_message() -> str:
    return os.getenv(
        "VOICELIVE_WELCOME_MESSAGE",
        "Hello! I'm your AI sales assistant. I'm here to help you with any questions about our products and services. "
        "How can I assist you today?",
    )


async def _safe_json(text: str) -> Optional[dict]:
    try:
        import json

        return json.loads(text)
    except Exception:
        return None
