from __future__ import annotations

import asyncio
import time
from typing import Awaitable, Callable, Optional
import numpy as np

from core.state_machine import Arbitrator
from services.asr import ASRService
from services.llm import LLMService
from services.tts import TTSService


AudioChunkCallback = Callable[[np.ndarray, int, str, int, int], Awaitable[None]]


class ChatPipeline:
    def __init__(
        self,
        arbitrator: Arbitrator,
        asr: ASRService,
        llm: LLMService,
        tts: TTSService,
    ) -> None:
        self.arbitrator = arbitrator
        self.asr = asr
        self.llm = llm
        self.tts = tts

    async def _finish_avatar_after_playback(self, playback_s: float) -> None:
        t_end = time.perf_counter() + playback_s
        while time.perf_counter() < t_end:
            if self.arbitrator.cancel_avatar_event.is_set():
                break
            await asyncio.sleep(0.05)
        await self.arbitrator.on_avatar_done()

    @staticmethod
    def _split_sentences(text: str, max_len: int = 24) -> list[str]:
        # Two-stage split:
        # 1) hard sentence boundaries
        # 2) soft commas / max length chunks for lower first-chunk latency
        hard_seps = "。！？!?；;\n"
        soft_seps = "，、, "

        stage1 = []
        buf = []
        for ch in text:
            buf.append(ch)
            if ch in hard_seps:
                seg = "".join(buf).strip()
                if seg:
                    stage1.append(seg)
                buf = []
        tail = "".join(buf).strip()
        if tail:
            stage1.append(tail)
        if not stage1:
            stage1 = [text]

        out = []
        for seg in stage1:
            cur = []
            for ch in seg:
                cur.append(ch)
                cur_len = len(cur)
                if ch in soft_seps and cur_len >= 8:
                    out.append("".join(cur).strip())
                    cur = []
                elif cur_len >= max_len:
                    out.append("".join(cur).strip())
                    cur = []
            if cur:
                out.append("".join(cur).strip())

        return [s for s in out if s]

    async def _synthesize_sentence_first(
        self,
        reply_text: str,
        on_audio_chunk: Optional[AudioChunkCallback] = None,
        on_text_segment: Optional[Callable[[str, int, int], Awaitable[None]]] = None,
        should_interrupt: Optional[Callable[[], bool]] = None,
    ) -> tuple[np.ndarray, int, int, int]:
        segments = self._split_sentences(reply_text)
        audio_chunks = []
        sr = 24000
        t0 = time.perf_counter()
        first_chunk_ms = 0
        for i, seg in enumerate(segments):
            if should_interrupt and should_interrupt():
                break
            if on_text_segment is not None:
                await on_text_segment(seg, i, len(segments))
            seg_audio, seg_sr = await self.tts.synthesize(seg)
            sr = seg_sr
            if i == 0:
                first_chunk_ms = int((time.perf_counter() - t0) * 1000)
            if on_audio_chunk is not None:
                await on_audio_chunk(seg_audio, seg_sr, seg, i, len(segments))
            audio_chunks.append(seg_audio)
        tts_total_ms = int((time.perf_counter() - t0) * 1000)
        audio = np.concatenate(audio_chunks) if audio_chunks else np.zeros(1, dtype=np.float32)
        return audio, sr, first_chunk_ms, tts_total_ms

    async def process_turn(
        self,
        audio_16k: np.ndarray,
        on_audio_chunk: Optional[AudioChunkCallback] = None,
        on_user_text: Optional[Callable[[str], Awaitable[None]]] = None,
        on_reply_text: Optional[Callable[[str], Awaitable[None]]] = None,
        on_reply_segment: Optional[Callable[[str, int, int], Awaitable[None]]] = None,
    ) -> tuple[dict, np.ndarray, int]:
        await self.arbitrator.on_speech_end()
        asr_t0 = time.perf_counter()
        user_text = await self.asr.transcribe(audio_16k)
        asr_latency_ms = int((time.perf_counter() - asr_t0) * 1000)
        if on_user_text is not None:
            await on_user_text(user_text)
        llm_t0 = time.perf_counter()
        reply_text, llm_source = await self.llm.reply_with_meta(user_text)
        if on_reply_text is not None:
            await on_reply_text(reply_text)
        llm_latency_ms = int((time.perf_counter() - llm_t0) * 1000)
        audio, sr, tts_first_chunk_ms, tts_latency_ms = await self._synthesize_sentence_first(
            reply_text,
            on_audio_chunk=on_audio_chunk,
            on_text_segment=on_reply_segment,
            should_interrupt=lambda: self.arbitrator.cancel_avatar_event.is_set(),
        )
        await self.arbitrator.on_avatar_start()
        playback_s = max(0.2, min(5.0, float(audio.shape[0]) / float(sr)))
        # Do not block request completion on playback duration.
        asyncio.create_task(self._finish_avatar_after_playback(playback_s))
        payload = {
            "user_text": user_text,
            "reply_text": reply_text,
            "audio_samples": int(audio.shape[0]),
            "sample_rate": sr,
            "asr_latency_ms": asr_latency_ms,
            "llm_latency_ms": llm_latency_ms,
            "llm_source": llm_source,
            "tts_latency_ms": tts_latency_ms,
            "tts_first_chunk_ms": tts_first_chunk_ms,
        }
        return payload, audio, sr

    async def process_text_turn(
        self,
        user_text: str,
        on_audio_chunk: Optional[AudioChunkCallback] = None,
        on_reply_text: Optional[Callable[[str], Awaitable[None]]] = None,
        on_reply_segment: Optional[Callable[[str, int, int], Awaitable[None]]] = None,
    ) -> tuple[dict, np.ndarray, int]:
        llm_t0 = time.perf_counter()
        reply_text, llm_source = await self.llm.reply_with_meta(user_text)
        if on_reply_text is not None:
            await on_reply_text(reply_text)
        llm_latency_ms = int((time.perf_counter() - llm_t0) * 1000)
        audio, sr, tts_first_chunk_ms, tts_latency_ms = await self._synthesize_sentence_first(
            reply_text,
            on_audio_chunk=on_audio_chunk,
            on_text_segment=on_reply_segment,
            should_interrupt=lambda: self.arbitrator.cancel_avatar_event.is_set(),
        )
        await self.arbitrator.on_avatar_start()
        playback_s = max(0.2, min(5.0, float(audio.shape[0]) / float(sr)))
        asyncio.create_task(self._finish_avatar_after_playback(playback_s))
        payload = {
            "user_text": user_text,
            "reply_text": reply_text,
            "audio_samples": int(audio.shape[0]),
            "sample_rate": sr,
            "llm_latency_ms": llm_latency_ms,
            "llm_source": llm_source,
            "tts_latency_ms": tts_latency_ms,
            "tts_first_chunk_ms": tts_first_chunk_ms,
        }
        return payload, audio, sr

    async def run_once(self, audio_16k: np.ndarray) -> dict:
        payload, _, _ = await self.process_turn(audio_16k)
        return payload
