import asyncio
from collections import deque
import os
import threading
import time
import traceback
import av
import numpy as np
import streamlit as st
from streamlit_webrtc import WebRtcMode, webrtc_streamer
import pydub
import torch
# import av
# import cv2
from sample_utils.turn import get_ice_servers
import json
from typing import List

from vosk import SetLogLevel, Model, KaldiRecognizer
SetLogLevel(-1) # mutes vosk verbosity

from dotenv import load_dotenv
load_dotenv()

webrtc_ctx = None

# Initialize Ray
import ray
if not ray.is_initialized():
    # Try to connect to a running Ray cluster
    ray_address = os.getenv('RAY_ADDRESS')
    if ray_address:
        ray.init(ray_address, namespace="project_charles")
    else:
        ray.init(namespace="project_charles")



async def main():

    system_one_audio_status = st.empty()

    playing = st.checkbox("Playing", value=True)

    system_one_audio_status.write("Initializing streaming")
    system_one_audio_output = st.empty()

    system_one_video_output = st.empty()

    system_one_audio_history = []
    system_one_audio_history_output = st.empty()

    # Initialize resources if not already done
    system_one_audio_status.write("Initializing streaming")
    if "streamlit_av_queue" not in st.session_state:
        from streamlit_av_queue import StreamlitAVQueue
        st.session_state.streamlit_av_queue = StreamlitAVQueue()

    system_one_audio_status.write("resources referecned")



    system_one_audio_status.write("Initializing webrtc_streamer")
    webrtc_ctx = webrtc_streamer(
        key="charles",
        desired_playing_state=playing,
        queued_audio_frames_callback=st.session_state.streamlit_av_queue.queued_audio_frames_callback,
        queued_video_frames_callback=st.session_state.streamlit_av_queue.queued_video_frames_callback,
        mode=WebRtcMode.SENDRECV,
        rtc_configuration={"iceServers": get_ice_servers()},
        async_processing=True,
    )

    if not webrtc_ctx.state.playing:
        exit

    system_one_audio_status.write("Initializing speech")
    
    from charles_actor import CharlesActor
    charles_actor = None


    try:
        while True:
            if not webrtc_ctx.state.playing:
                system_one_audio_status.write("Stopped.")
                await asyncio.sleep(0.1)
                continue
            if charles_actor is None:
                try:
                    charles_actor = ray.get_actor("CharlesActor")
                    system_one_audio_status.write("Charles is here.")
                except ValueError as e:
                    system_one_audio_status.write("Charles is sleeping.")
                    pass
            if charles_actor is not None: 
                audio_history = ray.get(charles_actor.get_system_one_audio_history_output.remote())
                system_one_audio_history_output.markdown(audio_history)
            await asyncio.sleep(0.1)

    except Exception as e:
        print(f"An error occurred: {e}")
        traceback.print_exc()
        raise e


if __name__ == "__main__":
    try:
        asyncio.run(main())
    except Exception as e:
        if  webrtc_ctx is not None:
            del webrtc_ctx
            webrtc_ctx = None
        if "streamlit_av_queue" in st.session_state:
            del st.session_state.streamlit_av_queue    
    finally:
        pass