import os
import logging
import asyncio
import uuid

from apipeline.frames.control_frames import EndFrame
from apipeline.pipeline.task import PipelineTask

from achatbot.processors.image.base import ImageGenProcessor
from achatbot.processors.image import get_image_gen_processor
from achatbot.modules.vision.ocr import VisionOCREnvInit
from achatbot.modules.vision.detector import VisionDetectorEnvInit
from achatbot.processors.ai_processor import AIProcessor
from achatbot.processors.vision.vision_processor import MockVisionProcessor
from achatbot.processors.speech.asr.base import ASRProcessorBase
from achatbot.processors.llm.base import LLMProcessor
from achatbot.processors.speech.tts.base import TTSProcessorBase
from achatbot.modules.speech.vad_analyzer import VADAnalyzerEnvInit
from achatbot.modules.speech.asr import ASREnvInit
from achatbot.core.llm import LLMEnvInit
from achatbot.modules.speech.tts import TTSEnvInit
from achatbot.types.ai_conf import ASRConfig, LLMConfig, TTSConfig, AIConfig
from achatbot.common import interface
from achatbot.common.factory import EngineClass
from achatbot.common.types import BotRunArgs
from achatbot.common.interface import IBot, IVisionDetector
from achatbot.common.session import Session
from achatbot.common.types import SessionCtx

from dotenv import load_dotenv

load_dotenv(override=True)


class AIBot(IBot):
    r"""
    use ai bot config
    !TODONE: need config processor with bot config (redefine api params) @weedge
    bot config: Dict[str, Dict[str,Any]]
    e.g. {"llm":{"key":val,"tag":TAG,"args":{}}, "tts":{"key":val,"tag":TAG,"args":{}}}
    !TIPS: RTVI config options can transfer to ai bot config

    !NOTE:
    use multiprocessing pipe to run bot with unix socket, bot __init__ to new a bot obj must be serializable (pickle); or wraper a func, don't use bot obj methed.
    """

    def __init__(self, **args) -> None:
        self.args = BotRunArgs(**args)
        if self.args.bot_name is None or len(self.args.bot_name) == 0:
            self.args.bot_name = self.__class__.__name__

        self.task: PipelineTask | None = None
        self.session = Session(**SessionCtx(uuid.uuid4()).__dict__)

        self._bot_config_list = self.args.bot_config_list
        self._bot_config = self.args.bot_config
        self._handle_sigint = self.args.handle_sigint

    def init_bot_config(self):
        try:
            logging.debug(f'args.bot_config: {self.args.bot_config}')
            self._bot_config: AIConfig = AIConfig(**self.args.bot_config)
            if self._bot_config.llm is None:
                self._bot_config.llm = LLMConfig()
        except Exception as e:
            raise Exception(f"Failed to parse bot configuration: {e}")
        logging.info(f'ai bot_config: {self._bot_config}')

    def bot_config(self):
        return self._bot_config

    def run(self):
        asyncio.run(self.try_run())

    async def try_run(self):
        try:
            await self.arun()
        except asyncio.CancelledError:
            logging.info("CancelledError, Exiting!")
            if self.task is not None:
                await self.task.queue_frame(EndFrame())
        except KeyboardInterrupt:
            logging.warning("Ctrl-C detected. Exiting!")
            if self.task is not None:
                await self.task.queue_frame(EndFrame())
        except Exception as e:
            logging.error(f"run error: {e}", exc_info=True)

    async def arun(self):
        pass

    def get_vad_analyzer(self) -> interface.IVADAnalyzer | EngineClass:
        vad_analyzer: interface.IVADAnalyzer | EngineClass = None
        if self._bot_config.vad and self._bot_config.vad.tag \
                and len(self._bot_config.vad.tag) > 0  \
                and self._bot_config.vad.args:
            vad_analyzer = VADAnalyzerEnvInit.getEngine(
                self._bot_config.vad.tag, **self._bot_config.vad.args)
        else:
            vad_analyzer = VADAnalyzerEnvInit.initVADAnalyzerEngine()
        return vad_analyzer

    def get_asr_processor(self) -> ASRProcessorBase:
        asr_processor: ASRProcessorBase | None = None
        if self._bot_config.asr and self._bot_config.asr.tag \
                and self._bot_config.asr.tag == "deepgram_asr_processor" \
                and self._bot_config.asr.args:
            from achatbot.processors.speech.asr.deepgram_asr_processor import DeepgramAsrProcessor
            asr_processor = DeepgramAsrProcessor(
                api_key=os.getenv("DEEPGRAM_API_KEY"),
                **self._bot_config.asr.args)
        else:
            # use asr engine processor
            from achatbot.processors.speech.asr.asr_processor import ASRProcessor
            asr: interface.IAsr | EngineClass | None = None
            if self._bot_config.asr and self._bot_config.asr.tag \
                    and self._bot_config.asr.args:
                asr = ASREnvInit.getEngine(
                    self._bot_config.asr.tag, **self._bot_config.asr.args)
            else:
                logging.info(f"use default asr engine processor")
                asr = ASREnvInit.initASREngine()
                self._bot_config.asr = ASRConfig(tag=asr.SELECTED_TAG, args=asr.get_args_dict())
            asr_processor = ASRProcessor(
                asr=asr,
                session=self.session
            )
        return asr_processor

    def get_vision_llm_processor(self, llm_config: LLMConfig | None = None) -> LLMProcessor:
        """
        get local vision llm
        """
        from achatbot.processors.vision.vision_processor import VisionProcessor
        if not llm_config:
            llm_config = self._bot_config.vision_llm
        if "mock" in llm_config.tag:
            llm_processor = MockVisionProcessor()
        else:
            logging.debug(f"init engine llm processor tag: {llm_config.tag}")
            llm_engine = LLMEnvInit.initLLMEngine(llm_config.tag, llm_config.args)
            llm_processor = VisionProcessor(llm_engine, self.session)
        return llm_processor

    def get_remote_llm_processor(self, llm: LLMConfig | None = None) -> LLMProcessor:
        """
        get remote vision llm
        """
        if not llm:
            llm = self._bot_config.llm
        if llm and llm.tag \
                and "google" in llm.tag:
            llm_processor = self.get_google_llm_processor(llm)
        elif llm and llm.tag and "litellm" in llm.tag:
            llm_processor = self.get_litellm_processor(llm)
        else:
            llm_processor = self.get_openai_llm_processor(llm)
        return llm_processor

    def get_openai_llm_processor(self, llm: LLMConfig | None = None) -> LLMProcessor:
        from achatbot.processors.llm.openai_llm_processor import OpenAILLMProcessor, OpenAIGroqLLMProcessor
        if not llm:
            llm = self._bot_config.llm
        # default use openai llm processor
        api_key = os.environ.get("OPENAI_API_KEY")
        if "groq" in llm.base_url:
            # https://console.groq.com/docs/models
            api_key = os.environ.get("GROQ_API_KEY")
            llm_processor = OpenAIGroqLLMProcessor(
                model=llm.model,
                base_url=llm.base_url,
                api_key=api_key,
            )
            return llm_processor
        elif "together" in llm.base_url:
            # https://docs.together.ai/docs/chat-models
            api_key = os.environ.get("TOGETHER_API_KEY")
        llm_processor = OpenAILLMProcessor(
            model=llm.model,
            base_url=llm.base_url,
            api_key=api_key,
        )
        return llm_processor

    def get_google_llm_processor(self, llm: LLMConfig) -> LLMProcessor:
        from achatbot.processors.llm.google_llm_processor import GoogleAILLMProcessor
        llm_config = llm
        if llm_config and llm_config.args:
            llm_processor = GoogleAILLMProcessor(**llm_config.args)
        else:
            api_key = os.environ.get("GOOGLE_API_KEY")
            model = os.environ.get("MODEL", "gemini-1.5-flash-latest")
            llm_processor = GoogleAILLMProcessor(
                api_key=api_key,
                model=model,
            )

        return llm_processor

    def get_litellm_processor(self, llm: LLMConfig) -> LLMProcessor:
        from achatbot.processors.llm.litellm_processor import LiteLLMProcessor
        llm_processor = LiteLLMProcessor(model=llm.model, set_verbose=False)
        return llm_processor

    def get_llm_processor(self, llm: LLMConfig | None = None) -> LLMProcessor:
        if not llm:
            llm = self._bot_config.llm
        if llm and llm.tag and "vision" in llm.tag:
            # engine llm processor(just support vision model, other TODO):
            # (llm_llamacpp, llm_personalai_proxy, llm_transformers etc..)
            llm_processor = self.get_vision_llm_processor(llm)
        else:
            llm_processor = self.get_remote_llm_processor(llm)
        return llm_processor

    def get_vision_annotate_processor(self) -> AIProcessor:
        from achatbot.processors.vision.annotate_processor import AnnotateProcessor
        detector: IVisionDetector | EngineClass = VisionDetectorEnvInit.initVisionDetectorEngine(
            self._bot_config.vision_detector.tag,
            self._bot_config.vision_detector.args)
        processor = AnnotateProcessor(detector, self.session)
        return processor

    def get_vision_detect_processor(self) -> AIProcessor:
        from achatbot.processors.vision.detect_processor import DetectProcessor
        detector = VisionDetectorEnvInit.initVisionDetectorEngine(
            self._bot_config.vision_detector.tag,
            self._bot_config.vision_detector.args)

        desc, out_desc = "", ""
        if "desc" in self._bot_config.vision_detector.args:
            desc = self._bot_config.vision_detector.args["desc"]
        if "out_desc" in self._bot_config.vision_detector.args:
            out_desc = self._bot_config.vision_detector.args["out_desc"]
        if desc and out_desc:
            processor = DetectProcessor(detected_text=desc,
                                        out_detected_text=out_desc,
                                        detector=detector, session=self.session)
        elif desc:
            processor = DetectProcessor(detected_text=desc,
                                        detector=detector, session=self.session)
        elif out_desc:
            processor = DetectProcessor(out_detected_text=out_desc,
                                        detector=detector, session=self.session)
        else:
            processor = DetectProcessor(detector=detector, session=self.session)

        return processor

    def get_vision_ocr_processor(self) -> AIProcessor:
        from achatbot.processors.vision.ocr_processor import OCRProcessor
        ocr = VisionOCREnvInit.initVisionOCREngine(
            self._bot_config.vision_ocr.tag,
            self._bot_config.vision_ocr.args)
        processor = OCRProcessor(ocr=ocr, session=self.session)
        return processor

    def get_voice_processor(self, llm: LLMConfig | None = None) -> LLMProcessor:
        if not llm:
            llm = self._bot_config.voice_llm
        if llm and llm.tag and "moshi" in llm.tag:
            from achatbot.processors.voice.moshi_voice_processor import MoshiVoiceOpusStreamProcessor, MoshiVoiceProcessor, VoiceOpusStreamEchoProcessor
            if "moshi_opus" in llm.tag:
                if llm.args:
                    llm_processor = MoshiVoiceOpusStreamProcessor(**llm.args)
                else:
                    llm_processor = MoshiVoiceOpusStreamProcessor()
            else:
                if llm.args:
                    llm_processor = MoshiVoiceProcessor(**llm.args)
                else:
                    llm_processor = MoshiVoiceProcessor()
        else:
            # default VoiceOpusStreamEchoProcessor
            llm_processor = VoiceOpusStreamEchoProcessor()
        return llm_processor

    def get_tts_processor(self) -> TTSProcessorBase:
        tts_processor: TTSProcessorBase | None = None
        if self._bot_config.tts and self._bot_config.tts.tag and self._bot_config.tts.args:
            if self._bot_config.tts.tag == "elevenlabs_tts_processor":
                from achatbot.processors.speech.tts.elevenlabs_tts_processor import ElevenLabsTTSProcessor
                tts_processor = ElevenLabsTTSProcessor(**self._bot_config.tts.args)
            elif self._bot_config.tts.tag == "cartesia_tts_processor":
                from achatbot.processors.speech.tts.cartesia_tts_processor import CartesiaTTSProcessor
                tts_processor = CartesiaTTSProcessor(**self._bot_config.tts.args)
            else:
                # use tts engine processor
                from achatbot.processors.speech.tts.tts_processor import TTSProcessor
                tts = TTSEnvInit.getEngine(
                    self._bot_config.tts.tag, **self._bot_config.tts.args)
                self._bot_config.tts.tag = tts.SELECTED_TAG,
                self._bot_config.tts.args = tts.get_args_dict()
                tts_processor = TTSProcessor(tts=tts, session=self.session)
        else:
            # default tts engine processor
            from achatbot.processors.speech.tts.tts_processor import TTSProcessor
            logging.info(f"use default tts engine processor")
            tag = None
            if self._bot_config.tts and self._bot_config.tts.tag:
                tag = self._bot_config.tts.tag
            tts = TTSEnvInit.initTTSEngine(tag)
            self._bot_config.tts = TTSConfig(tag=tts.SELECTED_TAG, args=tts.get_args_dict())
            tts_processor = TTSProcessor(tts=tts, session=self.session)

        return tts_processor

    def get_image_gen_processor(self) -> ImageGenProcessor:
        if not self._bot_config.img_gen \
                or not self._bot_config.img_gen.args:
            raise Exception("need img_gen args params")
        return get_image_gen_processor(
            self._bot_config.img_gen.tag,
            **self._bot_config.img_gen.args)


class AIRoomBot(AIBot):
    def __init__(self, **args) -> None:
        super().__init__(**args)
