from __future__ import annotations

import asyncio
import json
from typing import Any

from aiortc.mediastreams import MediaStreamError
from av.audio.resampler import AudioResampler
import numpy as np

from dictation.audio import AudioEventKind, AudioSegmenter
from dictation.config import DictationConfig
from dictation.transcriber import Transcriber


class DictationSession:
    def __init__(
        self,
        session_id: str,
        peer: Any,
        transcriber: Transcriber,
        config: DictationConfig,
    ) -> None:
        self.session_id = session_id
        self.peer = peer
        self.transcriber = transcriber
        self.segmenter = AudioSegmenter(
            speech_threshold=config.speech_threshold,
            silence_ms=config.silence_ms,
            partial_interval_ms=config.partial_interval_ms,
            max_utterance_seconds=config.max_utterance_seconds,
        )
        self.channel = None
        self.audio_task: asyncio.Task | None = None
        self.partial_task: asyncio.Task | None = None
        self.final_tasks: set[asyncio.Task] = set()
        self.closed = False
        self.stopping = False
        self._resampler = AudioResampler(
            format="s16",
            layout="mono",
            rate=16000,
        )

    def attach_channel(self, channel: Any) -> None:
        if channel.label != "transcripts":
            channel.close()
            return
        self.channel = channel

        @channel.on("open")
        def on_open() -> None:
            self.send("ready", sessionId=self.session_id)

        @channel.on("message")
        def on_message(message: Any) -> None:
            if message == "stop":
                self.stopping = True
                event = self.segmenter.flush()
                if event and event.samples is not None:
                    self._start_final(event.samples)
                else:
                    self.send("speech.ended")

        if channel.readyState == "open":
            self.send("ready", sessionId=self.session_id)

    def attach_track(self, track: Any) -> None:
        if track.kind != "audio" or self.audio_task is not None:
            return
        self.audio_task = asyncio.create_task(self._consume_audio(track))

    async def _consume_audio(self, track: Any) -> None:
        try:
            while not self.closed:
                frame = await track.recv()
                for converted in self._resampler.resample(frame):
                    if self.stopping:
                        continue
                    samples = converted.to_ndarray().reshape(-1)
                    normalized = samples.astype(np.float32) / 32768.0
                    for event in self.segmenter.feed(normalized):
                        if event.kind == AudioEventKind.SPEECH_STARTED:
                            self.send("speech.started")
                        elif (
                            event.kind == AudioEventKind.PARTIAL_READY
                            and event.samples is not None
                        ):
                            self._start_partial(event.samples)
                        elif (
                            event.kind == AudioEventKind.FINAL_READY
                            and event.samples is not None
                        ):
                            self._start_final(event.samples)
        except MediaStreamError:
            if not self.stopping:
                event = self.segmenter.flush()
                if event and event.samples is not None:
                    self._start_final(event.samples)
                else:
                    self.send("speech.ended")
        except asyncio.CancelledError:
            raise
        except Exception as error:
            if not self.closed:
                self.send(
                    "error",
                    code="audio_track_failed",
                    message=str(error),
                )

    def _start_partial(self, samples: np.ndarray) -> None:
        if self.partial_task and not self.partial_task.done():
            return
        self.partial_task = asyncio.create_task(
            self._transcribe(samples, final=False)
        )

    def _start_final(self, samples: np.ndarray) -> None:
        task = asyncio.create_task(self._transcribe(samples, final=True))
        self.final_tasks.add(task)
        task.add_done_callback(self.final_tasks.discard)

    async def _transcribe(self, samples: np.ndarray, *, final: bool) -> None:
        try:
            transcript = await self.transcriber.transcribe(
                samples,
                final=final,
            )
            if transcript.text:
                self.send(
                    "transcript.final" if final else "transcript.partial",
                    text=transcript.text,
                    language=transcript.language,
                    probability=transcript.probability,
                )
            if final:
                self.send("speech.ended")
        except asyncio.CancelledError:
            raise
        except Exception as error:
            self.send(
                "error",
                code="transcription_failed",
                message=str(error),
            )

    def send(self, event_type: str, **payload: Any) -> None:
        if not self.channel or self.channel.readyState != "open":
            return
        self.channel.send(
            json.dumps(
                {"type": event_type, **payload},
                separators=(",", ":"),
            )
        )

    async def close(self) -> None:
        if self.closed:
            return
        self.closed = True
        tasks = [
            task
            for task in [
                self.audio_task,
                self.partial_task,
                *self.final_tasks,
            ]
            if task is not None and not task.done()
        ]
        for task in tasks:
            task.cancel()
        if tasks:
            await asyncio.gather(*tasks, return_exceptions=True)
        if self.channel and self.channel.readyState != "closed":
            self.channel.close()
        await self.peer.close()