import logging
from typing import AsyncGenerator
import uuid

from apipeline.frames.data_frames import Frame
from apipeline.frames.sys_frames import ErrorFrame

from src.common.factory import EngineClass
from src.common.session import Session, SessionCtx
from src.common.utils.time import time_now_iso8601
from src.common.interface import IAsr
from src.processors.speech.asr.base import SegmentedASRProcessor, SegmentedVolumeASRProcessor
from src.types.frames.data_frames import TranscriptionFrame
from src.types.speech.language import Language


# class ASRProcessor(SegmentedVolumeASRProcessor):
class ASRProcessor(SegmentedASRProcessor):
    def __init__(
        self,
        *,
        min_volume: float = 0.6,
        max_silence_secs: float = 0.3,
        max_buffer_secs: float = 1.5,
        sample_rate: int = 16000,
        num_channels: int = 1,
        asr: IAsr | EngineClass | None = None,
        session: Session | None = None,
        **kwargs,
    ):
        super().__init__(
            min_volume=min_volume,
            max_silence_secs=max_silence_secs,
            max_buffer_secs=max_buffer_secs,
            sample_rate=sample_rate,
            num_channels=num_channels,
            **kwargs,
        )
        self._asr = asr
        self._session = session or Session(**SessionCtx(str(uuid.uuid4())).__dict__)

    def set_asr(self, asr: IAsr):
        self._asr = asr

    async def set_asr_args(self, **args):
        self._asr.set_args(**args)

    async def run_asr(self, audio: bytes, **kwargs) -> AsyncGenerator[Frame, None]:
        if self._asr is None:
            logging.error(f"{self} error: ASR engine not available")
            yield ErrorFrame("ASR engine not available")
            return

        await self.start_processing_metrics()
        await self.start_ttfb_metrics()

        self._asr.set_audio_data(audio)
        text: str = ""
        i = 0
        async for segment in self._asr.transcribe_stream(self._session):
            if i == 0:
                await self.stop_ttfb_metrics()
            text += f"{segment}"
            i += 1

        await self.stop_processing_metrics()

        language = None
        args = self._asr.get_args_dict()
        if "language" in args:
            language = Language(args["language"])

        if text:
            logging.info(f"{self._asr.SELECTED_TAG} Transcription: [{text}] language: {language}")
            yield TranscriptionFrame(
                text=text,
                user_id=self._session.ctx.client_id,
                timestamp=time_now_iso8601(),
                language=language,
                speech_id=kwargs.get("speech_id", ""),
                start_at_s=kwargs.get("start_at_s", 0.0),
                end_at_s=kwargs.get("end_at_s", 0.0),
            )
