import logging
import uuid

from dotenv import load_dotenv
from fastapi import WebSocket
from apipeline.pipeline.pipeline import Pipeline
from apipeline.pipeline.task import PipelineParams, PipelineTask
from apipeline.pipeline.runner import PipelineRunner
from apipeline.processors.logger import FrameLogger
from apipeline.frames import AudioRawFrame, TextFrame

from achatbot.common.session import Session, SessionCtx
from achatbot.serializers.transcription_protobuf import TranscriptionFrameSerializer
from achatbot.processors.speech.vad_audio_save_processor import VADAudioSaveProcessor
from achatbot.modules.speech.asr_live import ASRLiveEnvInit
from achatbot.cmd.bots.base_fastapi_websocket_server import AIFastapiWebsocketBot
from achatbot.modules.speech.vad_analyzer import VADAnalyzerEnvInit
from achatbot.cmd.bots import register_ai_fastapi_ws_bots
from achatbot.types.network.fastapi_websocket import FastapiWebsocketServerParams
from achatbot.transports.fastapi_websocket_server import FastapiWebsocketTransport
from achatbot.processors.speech.asr.asr_live_processor import ASRLiveProcessor
from achatbot.processors.punctuation_processor import PunctuationProcessor
from achatbot.modules.punctuation import PuncEnvInit

load_dotenv(override=True)


@register_ai_fastapi_ws_bots.register
class FastapiWebsocketStreamingASRBot(AIFastapiWebsocketBot):
    """
    fastapi websocket input(audio)/output(text) server bot with vad,asr,punc,nt
    """

    def __init__(self, websocket: WebSocket | None = None, **args) -> None:
        super().__init__(websocket=websocket, **args)
        self.init_bot_config()

        self.vad_analyzer = None
        self.asr_engine = None
        self.punc_engine = None

    def load(self):
        # load vad analyer
        if self._bot_config.vad:
            tag = self._bot_config.vad.tag
            args = self._bot_config.vad.args or {}
            self.vad_analyzer = VADAnalyzerEnvInit.initVADAnalyzerEngine(tag, args)
        else:
            self.vad_analyzer = VADAnalyzerEnvInit.initVADAnalyzerEngine()

        # load asr live engine
        if self._bot_config.asr:
            tag = self._bot_config.asr.tag
            args = self._bot_config.asr.args or {}
            self.asr_engine = ASRLiveEnvInit.getEngine(tag, **args)
        else:
            logging.info("use defualt asr live engine")
            self.asr_engine = ASRLiveEnvInit.getEngine()

        # load punctuation engine
        if self.asr_engine.get_args_dict().get("textnorm", False) is False:
            if self._bot_config.punctuation:
                tag = self._bot_config.punctuation.tag
                args = self._bot_config.punctuation.args or {}
                self.punc_engine = PuncEnvInit.initEngine(tag, **args)

    async def arun(self):
        if self._websocket is None:
            return

        serializer = TranscriptionFrameSerializer()
        self.params = FastapiWebsocketServerParams(
            audio_in_enabled=True,
            audio_out_enabled=False,
            add_wav_header=True,
            vad_enabled=True,
            vad_analyzer=self.vad_analyzer,
            vad_audio_passthrough=True,
            serializer=serializer,
        )
        transport = FastapiWebsocketTransport(
            websocket=self._websocket,
            params=self.params,
        )

        session = Session(**SessionCtx(str(uuid.uuid4())).__dict__)
        asr_live_processor = ASRLiveProcessor(asr=self.asr_engine, session=session)

        processors = [
            transport.input_processor(),
            # FrameLogger(include_frame_types=[AudioRawFrame]),
            # VADAudioSaveProcessor(prefix_name="streaming_asr_vad", pass_raw_audio=True),
            asr_live_processor,
            FrameLogger(include_frame_types=[TextFrame]),
        ]
        if self.punc_engine:
            punc_processor = PunctuationProcessor(engine=self.punc_engine, session=session)
            processors.append(punc_processor)

        processors.append(FrameLogger(include_frame_types=[TextFrame]))
        processors.append(transport.output_processor())

        self.task = PipelineTask(
            Pipeline(processors=processors),
            params=PipelineParams(
                allow_interruptions=False,
                enable_metrics=True,
                send_initial_empty_metrics=False,
            ),
        )

        transport.add_event_handler("on_client_connected", self.on_client_connected)
        transport.add_event_handler("on_client_disconnected", self.on_client_disconnected)

        await PipelineRunner(handle_sigint=self._handle_sigint).run(self.task)
