import aiohttp
import asyncio
import logging
import os
import random
from typing import AsyncGenerator
from PIL import Image
from dailyai.pipeline.pipeline import Pipeline

from dailyai.transports.daily_transport import DailyTransport
from dailyai.services.open_ai_services import OpenAILLMService
from dailyai.services.elevenlabs_ai_service import ElevenLabsTTSService
from dailyai.pipeline.aggregators import (
    LLMUserContextAggregator,
    LLMAssistantContextAggregator,
)
from dailyai.pipeline.frames import (
    Frame,
    TextFrame,
    ImageFrame,
    SpriteFrame,
    TranscriptionFrame,
)
from dailyai.services.ai_services import AIService

from runner import configure

from dotenv import load_dotenv
load_dotenv(override=True)

logging.basicConfig(format=f"%(levelno)s %(asctime)s %(message)s")
logger = logging.getLogger("dailyai")
logger.setLevel(logging.DEBUG)

sprites = {}
image_files = [
    "sc-default.png",
    "sc-talk.png",
    "sc-listen-1.png",
    "sc-think-1.png",
    "sc-think-2.png",
    "sc-think-3.png",
    "sc-think-4.png",
]

script_dir = os.path.dirname(__file__)

for file in image_files:
    # Build the full path to the image file
    full_path = os.path.join(script_dir, "assets", file)
    # Get the filename without the extension to use as the dictionary key
    filename = os.path.splitext(os.path.basename(full_path))[0]
    # Open the image and convert it to bytes
    with Image.open(full_path) as img:
        sprites[file] = img.tobytes()

# When the bot isn't talking, show a static image of the cat listening
quiet_frame = ImageFrame(sprites["sc-listen-1.png"], (720, 1280))
# When the bot is talking, build an animation from two sprites
talking_list = [sprites["sc-default.png"], sprites["sc-talk.png"]]
talking = [random.choice(talking_list) for x in range(30)]
talking_frame = SpriteFrame(images=talking)

# TODO: Support "thinking" as soon as we get a valid transcript, while LLM
# is processing
thinking_list = [
    sprites["sc-think-1.png"],
    sprites["sc-think-2.png"],
    sprites["sc-think-3.png"],
    sprites["sc-think-4.png"],
]
thinking_frame = SpriteFrame(images=thinking_list)


class TranscriptFilter(AIService):
    def __init__(self, bot_participant_id=None):
        self.bot_participant_id = bot_participant_id

    async def process_frame(self, frame: Frame) -> AsyncGenerator[Frame, None]:
        if isinstance(frame, TranscriptionFrame):
            if frame.participantId != self.bot_participant_id:
                yield frame


class NameCheckFilter(AIService):
    def __init__(self, names: list[str]):
        self.names = names
        self.sentence = ""

    async def process_frame(self, frame: Frame) -> AsyncGenerator[Frame, None]:
        content: str = ""

        # TODO: split up transcription by participant
        if isinstance(frame, TextFrame):
            content = frame.text

        self.sentence += content
        if self.sentence.endswith((".", "?", "!")):
            if any(name in self.sentence for name in self.names):
                out = self.sentence
                self.sentence = ""
                yield TextFrame(out)
            else:
                out = self.sentence
                self.sentence = ""


class ImageSyncAggregator(AIService):
    def __init__(self):
        pass

    async def process_frame(self, frame: Frame) -> AsyncGenerator[Frame, None]:
        yield talking_frame
        yield frame
        yield quiet_frame


async def main(room_url: str, token):
    async with aiohttp.ClientSession() as session:
        transport = DailyTransport(
            room_url,
            token,
            "Santa Cat",
            duration_minutes=3,
            start_transcription=True,
            mic_enabled=True,
            mic_sample_rate=16000,
            camera_enabled=True,
            camera_width=720,
            camera_height=1280,
        )

        llm = OpenAILLMService(
            api_key=os.getenv("OPENAI_API_KEY"),
            model="gpt-4-turbo-preview")

        tts = ElevenLabsTTSService(
            aiohttp_session=session,
            api_key=os.getenv("ELEVENLABS_API_KEY"),
            voice_id="jBpfuIE2acCO8z3wKNLl",
        )
        isa = ImageSyncAggregator()

        messages = [
            {
                "role": "system",
                "content": "You are Santa Cat, a cat that lives in Santa's workshop at the North Pole. You should be clever, and a bit sarcastic. You should also tell jokes every once in a while.  Your responses should only be a few sentences long.",
            },
        ]

        tma_in = LLMUserContextAggregator(
            messages, transport._my_participant_id)
        tma_out = LLMAssistantContextAggregator(
            messages, transport._my_participant_id
        )
        tf = TranscriptFilter(transport._my_participant_id)
        ncf = NameCheckFilter(["Santa Cat", "Santa"])

        pipeline = Pipeline([isa, tf, ncf, tma_in, llm, tma_out, tts])

        @transport.event_handler("on_first_other_participant_joined")
        async def on_first_other_participant_joined(transport, participant):
            await transport.say(
                "Hi! If you want to talk to me, just say 'hey Santa Cat'.",
                tts,
            )

        async def starting_image():
            await transport.send_queue.put(quiet_frame)

        await asyncio.gather(transport.run(pipeline), starting_image())


if __name__ == "__main__":
    (url, token) = configure()
    asyncio.run(main(url, token))
