#
# Copyright (c) 2024, Daily
#
# SPDX-License-Identifier: BSD 2-Clause License
#

import os
import logging
import asyncio
import dataclasses
from typing import List, Literal, Optional

from pydantic import BaseModel, ValidationError
from apipeline.pipeline.pipeline import Pipeline
from apipeline.processors.frame_processor import FrameDirection, FrameProcessor
from apipeline.frames.sys_frames import SystemFrame, MetricsFrame
from apipeline.frames.control_frames import StartFrame

from achatbot.types.frames.sys_frames import BotInterruptionFrame
from achatbot.types.frames.control_frames import (
    LLMModelUpdateFrame,
    TTSVoiceUpdateFrame,
    UserStartedSpeakingFrame,
    UserStoppedSpeakingFrame,
)
from achatbot.types.frames.data_frames import (
    Frame,
    TextFrame,
    TranscriptionFrame,
    InterimTranscriptionFrame,
    LLMMessagesAppendFrame,
    LLMMessagesUpdateFrame,
    TTSSpeakFrame,
    TransportMessageFrame,
)
from achatbot.types.frames.control_frames import (
    LLMFullResponseEndFrame, LLMFullResponseStartFrame,
)
from achatbot.processors.aggregators.llm_response import (
    LLMAssistantResponseAggregator, LLMUserResponseAggregator,
)
from achatbot.processors.aggregators.openai_llm_context import OpenAILLMContext
from achatbot.transports.base import BaseTransport

from dotenv import load_dotenv
load_dotenv(override=True)


DEFAULT_LLM_SYS_MESSAGES = [
    {
        "role": "system",
        "content": os.getenv("LLM_CHAT_SYSTEM", ""),
    }]

GROQ_LLM_URL = "https://api.groq.com/openai/v1"
GROQ_LLM_MODEL = "llama-3.1-70b-versatile"

# https://docs.together.ai/docs/chat-models
TOGETHER_LLM_URL = "https://api.together.xyz/v1"
TOGETHER_LLM_MODEL = "Qwen/Qwen2-72B-Instruct"

DEFAULT_LLM_URL = os.getenv("LLM_OPENAI_BASE_URL", GROQ_LLM_URL)
DEFAULT_LLM_MODEL = os.getenv("LLM_OPENAI_MODEL", GROQ_LLM_MODEL)

DEFAULT_LLM_LANG = "zh"


class RTVILLMConfig(BaseModel):
    base_url: Optional[str] = DEFAULT_LLM_URL
    model: Optional[str] = DEFAULT_LLM_MODEL
    language: Optional[str] = DEFAULT_LLM_LANG
    messages: Optional[List[dict]] = None
    tag: Optional[str] = "openai_llm_processor"
    args: Optional[dict] = None


class RTVITTSConfig(BaseModel):
    voice: Optional[str] = None
    language: Optional[str] = None
    tag: Optional[str] = None
    args: Optional[dict] = None


class RTVIASRConfig(BaseModel):
    tag: Optional[str] = None
    args: Optional[dict] = None


class RTVIConfig(BaseModel):
    asr: Optional[RTVIASRConfig] = None
    llm: Optional[RTVILLMConfig] = None
    tts: Optional[RTVITTSConfig] = None


class RTVISetup(BaseModel):
    config: Optional[RTVIConfig] = None


class RTVILLMMessageData(BaseModel):
    messages: List[dict]


class RTVITTSMessageData(BaseModel):
    text: str
    interrupt: Optional[bool] = False


class RTVIMessageData(BaseModel):
    setup: Optional[RTVISetup] = None
    config: Optional[RTVIConfig] = None
    llm: Optional[RTVILLMMessageData] = None
    tts: Optional[RTVITTSMessageData] = None


class RTVIMessage(BaseModel):
    label: Literal["rtvi"] = "rtvi"
    type: str
    id: str
    data: Optional[RTVIMessageData] = None


class RTVIResponseData(BaseModel):
    success: bool
    error: Optional[str] = None


class RTVIResponse(BaseModel):
    label: Literal["rtvi"] = "rtvi"
    type: Literal["response"] = "response"
    id: str
    data: RTVIResponseData


class RTVIErrorData(BaseModel):
    message: str


class RTVIError(BaseModel):
    label: Literal["rtvi"] = "rtvi"
    type: Literal["error"] = "error"
    data: RTVIErrorData


class RTVILLMContextMessageData(BaseModel):
    messages: List[dict]


class RTVILLMContextMessage(BaseModel):
    label: Literal["rtvi"] = "rtvi"
    type: Literal["llm-context"] = "llm-context"
    data: RTVILLMContextMessageData


class RTVITTSTextMessageData(BaseModel):
    text: str


class RTVITTSTextMessage(BaseModel):
    label: Literal["rtvi"] = "rtvi"
    type: Literal["tts-text"] = "tts-text"
    data: RTVITTSTextMessageData


class RTVIBotReady(BaseModel):
    label: Literal["rtvi"] = "rtvi"
    type: Literal["bot-ready"] = "bot-ready"


class RTVITranscriptionMessageData(BaseModel):
    text: str
    user_id: str
    timestamp: str
    final: bool


class RTVITranscriptionMessage(BaseModel):
    label: Literal["rtvi"] = "rtvi"
    type: Literal["user-transcription"] = "user-transcription"
    data: RTVITranscriptionMessageData


class RTVIUserStartedSpeakingMessage(BaseModel):
    label: Literal["rtvi"] = "rtvi"
    type: Literal["user-started-speaking"] = "user-started-speaking"


class RTVIUserStoppedSpeakingMessage(BaseModel):
    label: Literal["rtvi"] = "rtvi"
    type: Literal["user-stopped-speaking"] = "user-stopped-speaking"


class RTVIJSONCompletion(BaseModel):
    label: Literal["rtvi"] = "rtvi"
    type: Literal["json-completion"] = "json-completion"
    data: str


class RTVITTSTextProcessor(FrameProcessor):

    def __init__(self):
        super().__init__()

    async def process_frame(self, frame: Frame, direction: FrameDirection):
        await super().process_frame(frame, direction)

        await self.push_frame(frame, direction)

        if isinstance(frame, TextFrame):
            message = RTVITTSTextMessage(data=RTVITTSTextMessageData(text=frame.text))
            await self.push_frame(TransportMessageFrame(message=message.model_dump(exclude_none=True)))


class FunctionCallProcessor(FrameProcessor):

    def __init__(self, context):
        super().__init__()
        self._checking = False
        self._aggregating = False
        self._emitted_start = False
        self._aggregation = ""
        self._context = context

        self._callbacks = {}
        self._start_callbacks = {}

    def register_function(self, function_name: str, callback, start_callback=None):
        self._callbacks[function_name] = callback
        if start_callback:
            self._start_callbacks[function_name] = start_callback

    def unregister_function(self, function_name: str):
        del self._callbacks[function_name]
        if self._start_callbacks[function_name]:
            del self._start_callbacks[function_name]

    def has_function(self, function_name: str):
        return function_name in self._callbacks.keys()

    async def call_function(self, function_name: str, args):
        if function_name in self._callbacks.keys():
            return await self._callbacks[function_name](self, args)
        return None

    async def call_start_function(self, function_name: str):
        if function_name in self._start_callbacks.keys():
            await self._start_callbacks[function_name](self)

    async def process_frame(self, frame: Frame, direction: FrameDirection):
        await super().process_frame(frame, direction)

        if isinstance(frame, LLMFullResponseStartFrame):
            self._checking = True
            await self.push_frame(frame, direction)
        elif isinstance(frame, TextFrame) and self._checking:
            # TODO-CB: should we expand this to any non-text character to start the completion?
            if frame.text.strip().startswith("{") or frame.text.strip().startswith("```"):
                self._emitted_start = False
                self._checking = False
                self._aggregation = frame.text
                self._aggregating = True
            else:
                self._checking = False
                self._aggregating = False
                self._aggregation = ""
                self._emitted_start = False
                await self.push_frame(frame, direction)
        elif isinstance(frame, TextFrame) and self._aggregating:
            self._aggregation += frame.text
            # TODO-CB: We can probably ignore function start I think
            # if not self._emitted_start:
            #     fn = re.search(r'{"function_name":\s*"(.*)",', self._aggregation)
            #     if fn and fn.group(1):
            #         await self.call_start_function(fn.group(1))
            #         self._emitted_start = True
        elif isinstance(frame, LLMFullResponseEndFrame) and self._aggregating:
            try:
                self._aggregation = self._aggregation.replace("```json", "").replace("```", "")
                self._context.add_message({"role": "assistant", "content": self._aggregation})
                message = RTVIJSONCompletion(data=self._aggregation)
                msg = message.model_dump(exclude_none=True)
                await self.push_frame(TransportMessageFrame(message=msg))

            except Exception as e:
                print(f"Error parsing function call json: {e}")
                print(f"aggregation was: {self._aggregation}")

            self._aggregating = False
            self._aggregation = ""
            self._emitted_start = False
        elif isinstance(frame, LLMFullResponseEndFrame):
            await self.push_frame(frame, direction)
        else:
            await self.push_frame(frame, direction)


class RTVIProcessor(FrameProcessor):
    def __init__(
            self,
            *,
            transport: BaseTransport,
            setup: RTVISetup | None = None,
            llm_processor: FrameProcessor | None = None,
            tts_processor: FrameProcessor | None = None):
        if llm_processor is None or tts_processor is None:
            raise ValueError("llm_processor and tts_processor must be provided")

        super().__init__()
        self._transport = transport
        self._setup = setup
        self._llm_processor = llm_processor
        self._tts_processor = tts_processor

        self._start_frame: Frame | None = None
        self._pipeline: FrameProcessor | None = None
        self._first_participant_joined: bool = False

        # Register transport event so we can send a `bot-ready` event (and maybe
        # others) when the participant joins.
        transport.add_event_handler(
            "on_first_participant_joined",
            self._on_first_participant_joined)

        self._frame_handler_task = self.get_event_loop().create_task(self._frame_handler())
        self._frame_queue = asyncio.Queue()

    async def process_frame(self, frame: Frame, direction: FrameDirection):
        await super().process_frame(frame, direction)

        if isinstance(frame, SystemFrame):
            await self.push_frame(frame, direction)
        else:
            await self._frame_queue.put((frame, direction))

        if isinstance(frame, StartFrame):
            self._start_frame = frame
            try:
                await self._handle_setup(self._setup)
            except Exception as e:
                logging.error(f"unable to setup RTVI, exception: {e}")
                await self._send_error(f"unable to setup RTVI: {e}")

    async def cleanup(self):
        self._frame_handler_task.cancel()
        await self._frame_handler_task

    async def _frame_handler(self):
        while True:
            try:
                (frame, direction) = await self._frame_queue.get()
                await self._handle_frame(frame, direction)
                self._frame_queue.task_done()
            except asyncio.CancelledError:
                break

    async def _handle_frame(self, frame: Frame, direction: FrameDirection):
        if isinstance(frame, TransportMessageFrame):
            await self._handle_message(frame)
        else:
            await self.push_frame(frame, direction)

        if isinstance(frame, TranscriptionFrame) or isinstance(frame, InterimTranscriptionFrame):
            await self._handle_transcriptions(frame)
        elif isinstance(frame, UserStartedSpeakingFrame) or isinstance(frame, UserStoppedSpeakingFrame):
            await self._handle_interruptions(frame)

    async def _handle_transcriptions(self, frame: Frame):
        # TODO(aleix): Once we add support for using custom piplines, the STTs will
        # be in the pipeline after this processor. This means the STT will have to
        # push transcriptions upstream as well.

        message = None
        if isinstance(frame, TranscriptionFrame):
            message = RTVITranscriptionMessage(
                data=RTVITranscriptionMessageData(
                    text=frame.text,
                    user_id=frame.user_id,
                    timestamp=frame.timestamp,
                    final=True))
        elif isinstance(frame, InterimTranscriptionFrame):
            message = RTVITranscriptionMessage(
                data=RTVITranscriptionMessageData(
                    text=frame.text,
                    user_id=frame.user_id,
                    timestamp=frame.timestamp,
                    final=False))

        if message:
            frame = TransportMessageFrame(message=message.model_dump(exclude_none=True))
            await self.push_frame(frame)

    async def _handle_interruptions(self, frame: Frame):
        message = None
        if isinstance(frame, UserStartedSpeakingFrame):
            message = RTVIUserStartedSpeakingMessage()
        elif isinstance(frame, UserStoppedSpeakingFrame):
            message = RTVIUserStoppedSpeakingMessage()

        if message:
            frame = TransportMessageFrame(message=message.model_dump(exclude_none=True))
            await self.push_frame(frame)

    async def _handle_message(self, frame: TransportMessageFrame):
        try:
            message = RTVIMessage.model_validate(frame.message)
        except ValidationError as e:
            await self._send_error(f"invalid message: {e}")
            return
        logging.info(f"handle message: {message}")

        # https://github.com/rtvi-ai#docs-events-and-data-structures-and-extensions
        # js sdk:
        # https://github.com/rtvi-ai/rtvi-client-web/blob/main/rtvi-client-js/src/messages.ts
        try:
            success = True
            error = None
            match message.type:
                case "setup":
                    setup = None
                    if message.data:
                        setup = message.data.setup
                    await self._handle_setup(message.id, setup)
                case "config-update":
                    await self._handle_config_update(message.data.config)
                case "llm-get-context":
                    await self._handle_llm_get_context()
                case "llm-append-context":
                    await self._handle_llm_append_context(message.data.llm)
                case "llm-update-context":
                    await self._handle_llm_update_context(message.data.llm)
                case "tts-speak":
                    await self._handle_tts_speak(message.data.tts)
                case "tts-interrupt":
                    await self._handle_tts_interrupt()
                case _:
                    success = False
                    error = f"unsupported type {message.type}"

            await self._send_response(message.id, success, error)
        except ValidationError as e:
            await self._send_response(message.id, False, f"invalid message: {e}")
        except Exception as e:
            await self._send_response(message.id, False, f"{e}")

    async def _handle_setup(self, setup: RTVISetup | None):
        self._tma_in = LLMUserResponseAggregator()
        self._tma_out = LLMAssistantResponseAggregator()

        # TODO-CB: Eventually we'll need to switch the context aggregators to use the
        # OpenAI context frames instead of message frames
        context = OpenAILLMContext()
        self._fc = FunctionCallProcessor(context)

        self._tts_text = RTVITTSTextProcessor()

        pipeline = Pipeline([
            self._tma_in,
            self._llm_processor,
            self._fc,
            self._tts_processor,
            self._tts_text,
            self._tma_out,
            self._transport.output_processor(),
        ])

        parent_pipeline: Pipeline = self.get_parent_pipeline()
        if parent_pipeline and self._start_frame:
            parent_pipeline.link(pipeline)

            # We need to initialize the new pipeline with the same settings
            # as the initial one.
            start_frame = dataclasses.replace(self._start_frame)
            await self.push_frame(start_frame)

            # Send new initial metrics with the new processors
            processors = parent_pipeline.processors_with_metrics()
            processors.extend(pipeline.processors_with_metrics())
            ttfb = [{"processor": p.name, "value": 0.0} for p in processors]
            processing = [{"processor": p.name, "value": 0.0} for p in processors]
            await self.push_frame(MetricsFrame(ttfb=ttfb, processing=processing))

        self._pipeline = pipeline

        await self._maybe_send_bot_ready()

    async def _handle_config_update(self, config: RTVIConfig):
        # Change voice before LLM updates, so we can hear the new vocie.
        if config.tts and config.tts.voice:
            frame = TTSVoiceUpdateFrame(config.tts.voice)
            await self.push_frame(frame)
        if config.llm and config.llm.model:
            frame = LLMModelUpdateFrame(config.llm.model)
            await self.push_frame(frame)
        if config.llm and config.llm.messages:
            frame = LLMMessagesUpdateFrame(config.llm.messages)
            await self.push_frame(frame)

    async def _handle_llm_get_context(self):
        data = RTVILLMContextMessageData(messages=self._tma_in.messages)
        message = RTVILLMContextMessage(data=data)
        frame = TransportMessageFrame(message=message.model_dump(exclude_none=True))
        await self.push_frame(frame)

    async def _handle_llm_append_context(self, data: RTVILLMMessageData):
        if data and data.messages:
            frame = LLMMessagesAppendFrame(data.messages)
            await self.push_frame(frame)

    async def _handle_llm_update_context(self, data: RTVILLMMessageData):
        if data and data.messages:
            frame = LLMMessagesUpdateFrame(data.messages)
            await self.push_frame(frame)

    async def _handle_tts_speak(self, data: RTVITTSMessageData):
        if data and data.text:
            if data.interrupt:
                await self._handle_tts_interrupt()
            frame = TTSSpeakFrame(text=data.text)
            await self.push_frame(frame)

    async def _handle_tts_interrupt(self):
        await self.push_frame(BotInterruptionFrame(), FrameDirection.UPSTREAM)

    async def _on_first_participant_joined(self, transport, participant):
        self._first_participant_joined = True
        await self._maybe_send_bot_ready()

    async def _maybe_send_bot_ready(self):
        if self._pipeline and self._first_participant_joined:
            message = RTVIBotReady()
            frame = TransportMessageFrame(message=message.model_dump(exclude_none=True))
            await self.push_frame(frame)

    async def _send_error(self, error: str):
        message = RTVIError(data=RTVIErrorData(message=error))
        frame = TransportMessageFrame(message=message.model_dump(exclude_none=True))
        await self.push_frame(frame)

    async def _send_response(self, id: str, success: bool, error: str | None = None):
        # TODO(aleix): This is a bit hacky, but we might get invalid
        # configuration or something might going wrong during setup and we would
        # like to send the error to the client. However, if the pipeline is not
        # setup yet we don't have an output transport and therefore we can't
        # send any messages. So, we setup a super basic pipeline with just the
        # output transport so we can send messages.
        if not self._pipeline:
            pipeline = Pipeline([self._transport.output_processor()])
            self._pipeline = pipeline

            parent_pipeline = self.get_parent_pipeline()
            if parent_pipeline and self._start_frame:
                parent_pipeline.link(pipeline)

        message = RTVIResponse(id=id, data=RTVIResponseData(success=success, error=error))
        frame = TransportMessageFrame(message=message.model_dump(exclude_none=True))
        await self.push_frame(frame)
