from __future__ import annotations

import asyncio
import base64
import os
from dataclasses import dataclass, field
from typing import Any, Dict, Optional

from app.observability.logging import get_logger
from app.providers.null_stt import NullSTTProvider
from app.providers.null_tts import NullTTSProvider
from app.providers.stt_base import STTProvider, STTStream, TranscriptEvent
from app.providers.tts_base import TTSProvider
from app.sales_brain.brain import SalesBrainEngine, merge_memory
from app.sales_brain.chunker import chunk_text_for_speech
from app.voice.audio import pcm16_bytes_to_seconds
from app.voice.ssml import build_ssml


@dataclass
class UsageMetrics:
    audio_bytes: int = 0
    stt_seconds: float = 0.0
    tts_seconds: float = 0.0
    turns: int = 0
    llm_tokens: int = 0


class VoiceSession:
    def __init__(
        self,
        websocket,
        principal: Dict[str, Any],
        sample_rate: int = 16000,
        stt_provider: Optional[STTProvider] = None,
        tts_provider: Optional[TTSProvider] = None,
    ) -> None:
        self.websocket = websocket
        self.principal = principal
        self.sample_rate = sample_rate
        self.stt_provider = stt_provider or NullSTTProvider()
        self.tts_provider = tts_provider or NullTTSProvider()
        self.brain = SalesBrainEngine()
        self.state = "listening"
        self.sales_state = "GREETING"
        self.memory: Dict[str, Any] = {}
        self.metrics = UsageMetrics()
        self.session_id = principal.get("sid")
        self.allow_barge_in = False
        self._listening_paused = False
        self._stt_stream: Optional[STTStream] = None
        self._stt_task: Optional[asyncio.Task] = None
        self._llm_task: Optional[asyncio.Task] = None
        self._tts_task: Optional[asyncio.Task] = None
        self._logger = get_logger("voice").bind(sid=self.session_id, user=principal.get("sub"))

    async def start(self) -> None:
        self._stt_stream = await self.stt_provider.start_stream(self.sample_rate)
        self._stt_task = asyncio.create_task(self._consume_transcripts())
        await self._send_state("listening")

    async def update_sample_rate(self, sample_rate: int) -> None:
        if sample_rate <= 0 or sample_rate == self.sample_rate:
            return
        self.sample_rate = sample_rate
        await self._restart_listening()

    async def end_utterance(self) -> None:
        if self._listening_paused:
            return
        await self._restart_listening()

    async def _restart_listening(self) -> None:
        if self._stt_stream:
            await self._stt_stream.end()
            self._stt_stream = None
        if self._stt_task:
            await self._stt_task
            self._stt_task = None
        if not self._listening_paused:
            self._stt_stream = await self.stt_provider.start_stream(self.sample_rate)
            self._stt_task = asyncio.create_task(self._consume_transcripts())

    async def close(self) -> None:
        await self.interrupt("session_end")
        if self._stt_stream:
            await self._stt_stream.end()
        if self._stt_task:
            await self._stt_task
        self._logger.info(
            "session_metrics",
            audio_bytes=self.metrics.audio_bytes,
            stt_seconds=round(self.metrics.stt_seconds, 2),
            tts_seconds=round(self.metrics.tts_seconds, 2),
            turns=self.metrics.turns,
        )

    async def on_audio(self, data: bytes) -> None:
        if not data:
            return
        self.metrics.audio_bytes += len(data)
        self.metrics.stt_seconds += pcm16_bytes_to_seconds(len(data), self.sample_rate)
        if self.state == "speaking":
            if not self.allow_barge_in:
                return
            await self.interrupt("barge_in")
        if self._stt_stream:
            await self._stt_stream.send_audio(data)

    async def interrupt(self, reason: str) -> None:
        self._logger.info("interrupt", reason=reason)
        for task in (self._llm_task, self._tts_task):
            if task and not task.done():
                task.cancel()
        await self._send_state("listening")

    async def _consume_transcripts(self) -> None:
        if not self._stt_stream:
            return
        async for event in self._stt_stream.results():
            if self.state == "speaking" and not self.allow_barge_in:
                continue
            if not event.text:
                continue
            if event.is_final:
                await self._handle_user_text(event.text)
            else:
                await self._send_json({"type": "stt.partial", "text": event.text})

    async def _handle_user_text(self, text: str) -> None:
        await self.interrupt("new_input")
        self.metrics.turns += 1
        await self._send_json({"type": "stt.final", "text": text})

        async def _run():
            try:
                await self._send_state("thinking")
                output = await asyncio.to_thread(self.brain.run_turn, self.sales_state, text, self.memory)
                self.sales_state = output.next_state
                merge_memory(self.memory, output)
                await self._send_json({"type": "agent.text.final", "text": output.reply})
                await self._speak(output.reply)
            except Exception as exc:
                self._logger.error("llm_error", error=str(exc))
                await self._send_json({"type": "agent.error", "detail": "LLM connection error"})
                await self._send_state("listening")

        self._llm_task = asyncio.create_task(_run())

    async def _speak(self, text: str) -> None:
        if not self.allow_barge_in:
            await self._pause_listening()
        await self._send_state("speaking")

        async def _run_tts() -> None:
            seq = 0
            sent_chunks = 0
            meta = self.memory.get("_meta") or {}
            decision_pressure = str(meta.get("decision_pressure") or "medium")
            rate, pitch, pause_ms = _voice_style(self.sales_state, decision_pressure)
            for chunk in chunk_text_for_speech(text, max_chars=180, max_sentences=2):
                await self._send_json({"type": "agent.text.delta", "text": chunk})
                ssml = build_ssml(chunk, rate=rate, pitch=pitch, pause_ms=pause_ms)
                async for audio_chunk in self.tts_provider.synthesize_stream(ssml):
                    if not audio_chunk:
                        continue
                    seq += 1
                    sent_chunks += 1
                    self.metrics.tts_seconds += pcm16_bytes_to_seconds(len(audio_chunk), self.sample_rate)
                    b64 = base64.b64encode(audio_chunk).decode("ascii")
                    await self._send_json({"type": "tts.audio.chunk", "seq": seq, "audio": b64})
            if sent_chunks == 0:
                await self._send_json({"type": "tts.error", "detail": "No audio returned from TTS provider"})
            await self._send_state("listening")
            if not self.allow_barge_in:
                await self._resume_listening()

        self._tts_task = asyncio.create_task(_run_tts())

    async def _send_state(self, state: str) -> None:
        self.state = state
        await self._send_json({"type": "agent.state", "state": state})

    async def _send_json(self, payload: Dict[str, Any]) -> None:
        try:
            await self.websocket.send_json(payload)
        except Exception as exc:
            self._logger.warning("send_failed", error=str(exc))

    async def _pause_listening(self) -> None:
        if self._listening_paused:
            return
        self._listening_paused = True
        if self._stt_stream:
            await self._stt_stream.end()
            self._stt_stream = None
        if self._stt_task:
            await self._stt_task
            self._stt_task = None

    async def _resume_listening(self) -> None:
        if not self._listening_paused:
            return
        self._listening_paused = False
        self._stt_stream = await self.stt_provider.start_stream(self.sample_rate)
        self._stt_task = asyncio.create_task(self._consume_transcripts())


def _voice_style(sales_state: str, decision_pressure: str) -> tuple[str, str, int]:
    base_rate = float(os.getenv("SPEECH_RATE", "1.3"))
    base_pct = int(round((base_rate - 1.0) * 100))
    base_pct = max(-50, min(100, base_pct))
    pressure_delta = {"low": -10, "medium": 0, "high": 10}
    rate_pct = base_pct + pressure_delta.get(decision_pressure, 0)
    rate_pct = max(-50, min(100, rate_pct))
    rate = f"{rate_pct}%"
    pitch = {"low": "-2st", "medium": "0st", "high": "+1st"}.get(decision_pressure, "0st")
    pause_ms = {"low": 380, "medium": 250, "high": 150}.get(decision_pressure, 250)
    if sales_state in ("DISCOVERY", "QUALIFICATION"):
        pause_ms = max(pause_ms, 300)
    if sales_state in ("OBJECTIONS", "VALUE"):
        pause_ms = max(pause_ms, 220)
    if sales_state in ("NEXT_STEP", "HANDOFF"):
        pause_ms = min(pause_ms, 140)
    return rate, pitch, pause_ms
