Update src/streamlit_app.py

src/streamlit_app.py

@@ -6,42 +6,47 @@ import base64
 import io
 import threading
 import traceback
-import time # For potential delays or timestamps if needed in future
-import atexit # For cleanup actions on exit
-
+import time
+import logging
 from dotenv import load_dotenv
 
-# --- Import main libraries ---
-import cv2
-import pyaudio
+import cv2 # For image processing
+import pyaudio # For audio PLAYBACK
 import PIL.Image
-import mss # For screen capture
 
 from google import genai
 from google.genai import types
 
+# streamlit-webrtc components
+from streamlit_webrtc import (
+    webrtc_streamer,
+    WebRtcMode,
+    AudioProcessorBase,
+    VideoProcessorBase,
+    ClientSettings
+)
+from aiortc import RTCIceServer, RTCConfiguration # For STUN server configuration
+
 # --- Configuration ---
-load_dotenv() # Load environment variables from .env file
+load_dotenv()
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(threadName)s - %(levelname)s - %(message)s')
 
 # Audio configuration
-FORMAT = pyaudio.paInt16
-CHANNELS = 1
-SEND_SAMPLE_RATE = 16000
-RECEIVE_SAMPLE_RATE = 24000  # Gemini documentation recommendation
-CHUNK_SIZE = 1024
-AUDIO_PLAYBACK_QUEUE_MAXSIZE = 50  # Buffer for audio chunks from Gemini for playback
-MEDIA_OUT_QUEUE_MAXSIZE = 20 # Buffer for audio/video frames to send to Gemini
+PYAUDIO_FORMAT = pyaudio.paInt16 # For PyAudio playback
+PYAUDIO_CHANNELS = 1 # For PyAudio playback
+WEBRTC_REQUESTED_AUDIO_CHANNELS = 1 # Request mono audio from WebRTC
+WEBRTC_REQUESTED_SEND_SAMPLE_RATE = 16000 # Target sample rate for audio sent to Gemini
+GEMINI_AUDIO_RECEIVE_SAMPLE_RATE = 24000  # Gemini documentation recommendation for its TTS
+PYAUDIO_PLAYBACK_CHUNK_SIZE = 1024 # For PyAudio playback
+AUDIO_PLAYBACK_QUEUE_MAXSIZE = 50
+MEDIA_TO_GEMINI_QUEUE_MAXSIZE = 30
 
 # Video configuration
-VIDEO_FPS_LIMIT = 1  # Send 1 frame per second to the API
-VIDEO_PREVIEW_RESIZE = (640, 480)  # Size for Streamlit preview display
-VIDEO_API_RESIZE = (1024, 1024)  # Max size for images sent to API (aspect ratio preserved)
+VIDEO_FPS_TO_GEMINI = 2  # Target FPS to send to Gemini (increased slightly)
+VIDEO_API_RESIZE = (1024, 1024) # Max size for images sent to API
 
-# Gemini model configuration
-MODEL_NAME = "models/gemini-2.0-flash-live-001" # VERIFY THIS MODEL NAME IS CORRECT FOR LIVE API
-DEFAULT_VIDEO_MODE = "camera" # Default video input mode ("camera", "screen", "none")
+MODEL_NAME = "models/gemini-2.0-flash-live-001"
 
-# System Prompt for the Medical Assistant
 MEDICAL_ASSISTANT_SYSTEM_PROMPT = """You are an AI Medical Assistant. Your primary function is to analyze visual information from the user's camera or screen and respond via voice.
 
 Your responsibilities are:
@@ -58,467 +63,370 @@ Your responsibilities are:
 Example of a disclaimer you might use: "As an AI assistant, I can describe what I see, but I can't provide medical advice or diagnoses. For any health concerns, it's always best to speak with a doctor or other healthcare professional."
 """
 
-# Initialize PyAudio instance globally
+# --- PyAudio Global Instance and Cleanup ---
 pya = pyaudio.PyAudio()
-
-# Ensure PyAudio is terminated on exit
 def cleanup_pyaudio():
-    print("Terminating PyAudio instance.")
-    pya.terminate()
+    logging.info("Terminating PyAudio instance.")
+    if pya: # Check if pya is not None
+        pya.terminate()
 atexit.register(cleanup_pyaudio)
 
+# --- Global Queues for WebRTC to Backend Communication ---
+video_frames_to_gemini_q = asyncio.Queue(maxsize=MEDIA_TO_GEMINI_QUEUE_MAXSIZE)
+audio_chunks_to_gemini_q = asyncio.Queue(maxsize=MEDIA_TO_GEMINI_QUEUE_MAXSIZE)
+audio_from_gemini_playback_q = asyncio.Queue(maxsize=AUDIO_PLAYBACK_QUEUE_MAXSIZE)
 
-# Initialize Streamlit session state variables
-def init_session_state():
-    defaults = {
-        'app_initialized': True, # To ensure this runs only once per session
-        'session_active': False,
-        'audio_loop_instance': None,
-        'chat_messages': [],
-        'current_frame_preview': None,
-        'video_mode_selection': DEFAULT_VIDEO_MODE,
-    }
-    for key, value in defaults.items():
-        if key not in st.session_state:
-            st.session_state[key] = value
-
-init_session_state()
-
-
-# Configure Streamlit page
-st.set_page_config(page_title="Live Medical Assistant", layout="wide")
-
-# Initialize Gemini client
+# --- Gemini Client Setup ---
 GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY")
-if not GEMINI_API_KEY:
-    st.error("GEMINI_API_KEY not found. Please set it in your environment variables or a .env file.")
+client = None
+if GEMINI_API_KEY:
+    try:
+        client = genai.Client(http_options={"api_version": "v1beta"}, api_key=GEMINI_API_KEY)
+    except Exception as e:
+        # This error will be shown in Streamlit UI if it happens at startup
+        st.error(f"Failed to initialize Gemini client: {e}")
+        logging.critical(f"Gemini client initialization failed: {e}", exc_info=True)
+        st.stop() # Stop Streamlit app if client fails
+else:
+    st.error("GEMINI_API_KEY not found in environment variables. Please set it for the application to run.")
+    logging.critical("GEMINI_API_KEY not found.")
     st.stop()
 
-# Use 'client' consistently for the genai.Client instance
-client = None # Define client in the global scope
-try:
-    client = genai.Client(
-        http_options={"api_version": "v1beta"}, # Required for live
-        api_key=GEMINI_API_KEY,
-    )
-except Exception as e:
-    st.error(f"Failed to initialize Gemini client: {e}")
-    st.stop()
-
-# Gemini LiveConnectConfig
 LIVE_CONNECT_CONFIG = types.LiveConnectConfig(
-    response_modalities=["audio", "text"], # Expect both audio and text responses
+    response_modalities=["audio", "text"],
     speech_config=types.SpeechConfig(
-        voice_config=types.VoiceConfig(
-            prebuilt_voice_config=types.PrebuiltVoiceConfig(voice_name="Zephyr") # Changed voice to Zephyr
-        )
+        voice_config=types.VoiceConfig(prebuilt_voice_config=types.PrebuiltVoiceConfig(voice_name="Zephyr"))
     ),
 )
 
-class AudioVisualLoop:
-    def __init__(self, video_mode_setting=DEFAULT_VIDEO_MODE):
-        self.video_mode = video_mode_setting
+# --- Backend Gemini Interaction Loop ---
+class GeminiInteractionLoop:
+    def __init__(self):
         self.gemini_session = None
-        self.async_event_loop = None # To store the loop for thread-safe calls
-
-        self.audio_playback_queue = None # asyncio.Queue for audio from Gemini
-        self.media_to_gemini_queue = None # asyncio.Queue for audio/video to Gemini
-
-        self.is_running = True # Flag to control all async tasks
-        self.mic_stream = None # PyAudio input stream
-        self.playback_stream = None # PyAudio output stream
-        self.camera_capture = None # OpenCV VideoCapture object
+        self.async_event_loop = None
+        self.is_running = True
+        self.playback_stream = None
 
     async def send_text_input_to_gemini(self, user_text):
         if not user_text or not self.gemini_session or not self.is_running:
-            print("Warning: Cannot send text. Session not active or no text.")
+            logging.warning("Cannot send text. Session not active, no text, or not running.")
             return
         try:
-            print(f"Sending text to Gemini: {user_text}")
+            logging.info(f"Sending text to Gemini: '{user_text[:50]}...'")
             await self.gemini_session.send(input=user_text, end_of_turn=True)
         except Exception as e:
-            st.error(f"Error sending text message to Gemini: {e}")
-            print(f"Traceback for send_text_input_to_gemini: {traceback.format_exc()}")
-
-    def _process_camera_frame(self):
-        if not self.camera_capture or not self.camera_capture.isOpened():
-            print("Camera not available or not open.")
-            return None
-        ret, frame = self.camera_capture.read()
-        if not ret:
-            print("Failed to read frame from camera.")
-            return None
-        
-        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-        img = PIL.Image.fromarray(frame_rgb)
-        
-        preview_img = img.copy()
-        preview_img.thumbnail(VIDEO_PREVIEW_RESIZE)
-        
-        api_img = img.copy()
-        api_img.thumbnail(VIDEO_API_RESIZE) # Preserves aspect ratio
-        
-        image_io = io.BytesIO()
-        api_img.save(image_io, format="jpeg")
-        image_bytes = image_io.getvalue()
-        
-        return {
-            "preview": preview_img,
-            "api_data": {"mime_type": "image/jpeg", "data": base64.b64encode(image_bytes).decode()}
-        }
-
-    async def stream_camera_frames(self):
-        try:
-            self.camera_capture = await asyncio.to_thread(cv2.VideoCapture, 0) # 0 for default camera
-            if not self.camera_capture.isOpened():
-                st.error("Could not open camera. Please check permissions and availability.")
-                self.is_running = False
-                return
-
-            while self.is_running:
-                frame_data = await asyncio.to_thread(self._process_camera_frame)
-                if frame_data:
-                    st.session_state['current_frame_preview'] = frame_data["preview"]
-                    if self.media_to_gemini_queue.full():
-                        await self.media_to_gemini_queue.get() # Make space
-                    await self.media_to_gemini_queue.put(frame_data["api_data"])
-                await asyncio.sleep(1.0 / VIDEO_FPS_LIMIT)
-        except Exception as e:
-            # st.error(f"Camera streaming error: {e}") # Avoid st.error from non-main thread if problematic
-            print(f"Camera streaming error: {e}\nTraceback: {traceback.format_exc()}")
-            self.is_running = False
-        finally:
-            if self.camera_capture:
-                await asyncio.to_thread(self.camera_capture.release)
-                self.camera_capture = None
-            print("Camera streaming task finished.")
-
-    def _process_screen_frame(self):
-        with mss.mss() as sct:
-            monitor_index = 1 
-            if len(sct.monitors) <= monitor_index: 
-                monitor_index = 0 
-            monitor = sct.monitors[monitor_index]
-            
-            sct_img = sct.grab(monitor)
-            img = PIL.Image.frombytes("RGB", (sct_img.width, sct_img.height), sct_img.rgb)
-            
-            preview_img = img.copy()
-            preview_img.thumbnail(VIDEO_PREVIEW_RESIZE)
-            
-            api_img = img.copy()
-            api_img.thumbnail(VIDEO_API_RESIZE)
-            
-            image_io = io.BytesIO()
-            api_img.save(image_io, format="jpeg")
-            image_bytes = image_io.getvalue()
-            
-            return {
-                "preview": preview_img,
-                "api_data": {"mime_type": "image/jpeg", "data": base64.b64encode(image_bytes).decode()}
-            }
-
-    async def stream_screen_frames(self):
-        try:
-            while self.is_running:
-                frame_data = await asyncio.to_thread(self._process_screen_frame)
-                if frame_data:
-                    st.session_state['current_frame_preview'] = frame_data["preview"]
-                    if self.media_to_gemini_queue.full():
-                        await self.media_to_gemini_queue.get()
-                    await self.media_to_gemini_queue.put(frame_data["api_data"])
-                await asyncio.sleep(1.0 / VIDEO_FPS_LIMIT)
-        except Exception as e:
-            # st.error(f"Screen capture error: {e}")
-            print(f"Screen capture error: {e}\nTraceback: {traceback.format_exc()}")
-            self.is_running = False
-        finally:
-            print("Screen streaming task finished.")
+            logging.error(f"Error sending text message to Gemini: {e}", exc_info=True)
 
     async def stream_media_to_gemini(self):
+        logging.info("Task started: Stream media from WebRTC queues to Gemini.")
+        async def get_media_from_queues():
+            try:
+                video_frame = await asyncio.wait_for(video_frames_to_gemini_q.get(), timeout=0.02)
+                video_frames_to_gemini_q.task_done()
+                return video_frame
+            except asyncio.TimeoutError: pass
+            except Exception as e: logging.error(f"Error getting video from queue: {e}", exc_info=True)
+
+            try:
+                audio_chunk = await asyncio.wait_for(audio_chunks_to_gemini_q.get(), timeout=0.02)
+                audio_chunks_to_gemini_q.task_done()
+                return audio_chunk
+            except asyncio.TimeoutError: return None
+            except Exception as e:
+                logging.error(f"Error getting audio from queue: {e}", exc_info=True)
+                return None
         try:
             while self.is_running:
-                if not self.gemini_session: 
-                    await asyncio.sleep(0.1) 
-                    continue
-                try:
-                    media_chunk = await asyncio.wait_for(self.media_to_gemini_queue.get(), timeout=1.0)
-                    if media_chunk and self.gemini_session and self.is_running:
-                         await self.gemini_session.send(input=media_chunk) 
-                    if media_chunk: # Avoid task_done on None if used as sentinel
-                        self.media_to_gemini_queue.task_done()
-                except asyncio.TimeoutError:
-                    continue 
-                except Exception as e:
-                    if self.is_running:
-                        print(f"Error in stream_media_to_gemini: {e}")
-                        await asyncio.sleep(0.1)
-        except asyncio.CancelledError:
-            print("stream_media_to_gemini task cancelled.")
-        finally:
-            print("Media streaming to Gemini task finished.")
-
-    async def capture_microphone_audio(self):
-        try:
-            mic_info = await asyncio.to_thread(pya.get_default_input_device_info)
-            self.mic_stream = await asyncio.to_thread(
-                pya.open,
-                format=FORMAT, channels=CHANNELS, rate=SEND_SAMPLE_RATE,
-                input=True, input_device_index=mic_info["index"],
-                frames_per_buffer=CHUNK_SIZE,
-            )
-            print("Microphone stream opened.")
-            while self.is_running:
-                try:
-                    audio_data = await asyncio.to_thread(self.mic_stream.read, CHUNK_SIZE, exception_on_overflow=False)
-                    if self.media_to_gemini_queue.full():
-                        await self.media_to_gemini_queue.get()
-                    await self.media_to_gemini_queue.put({"data": audio_data, "mime_type": "audio/pcm"})
-                except IOError as e:
-                    if e.errno == pyaudio.paInputOverflowed: # type: ignore
-                        print("Microphone Input overflowed. Skipping.")
-                    else:
-                        print(f"Microphone read IOError: {e}")
-                        self.is_running = False; break 
-                except Exception as e:
-                    print(f"Error in capture_microphone_audio: {e}")
-                    await asyncio.sleep(0.01)
-        except Exception as e:
-            # st.error(f"Failed to open microphone: {e}. Please check permissions.")
-            print(f"Failed to open microphone: {e}\nTraceback: {traceback.format_exc()}")
-            self.is_running = False
-        finally:
-            if self.mic_stream:
-                await asyncio.to_thread(self.mic_stream.stop_stream)
-                await asyncio.to_thread(self.mic_stream.close)
-                self.mic_stream = None
-            print("Microphone capture task finished.")
+                if not self.gemini_session:
+                    await asyncio.sleep(0.1); continue
+                media_data = await get_media_from_queues()
+                if media_data and self.gemini_session and self.is_running:
+                    try:
+                        await self.gemini_session.send(input=media_data)
+                    except Exception as e: logging.error(f"Error sending media chunk to Gemini: {e}", exc_info=True)
+                elif not media_data: await asyncio.sleep(0.05)
+        except asyncio.CancelledError: logging.info("Task cancelled: stream_media_to_gemini.")
+        finally: logging.info("Task finished: stream_media_to_gemini.")
 
     async def process_gemini_responses(self):
+        logging.info("Task started: Process responses from Gemini.")
         try:
             while self.is_running:
-                if not self.gemini_session: 
+                if not self.gemini_session:
                     await asyncio.sleep(0.1); continue
                 try:
-                    turn_response = self.gemini_session.receive() 
+                    turn_response = self.gemini_session.receive()
                     async for chunk in turn_response:
                         if not self.is_running: break
-                        if audio_data := chunk.data: 
-                            if not self.audio_playback_queue.full():
-                                self.audio_playback_queue.put_nowait(audio_data)
-                            else:
-                                print("Audio playback queue full, discarding data.")
-                        if text_response := chunk.text: 
-                            # Schedule Streamlit update from the main thread if possible,
-                            # or use st.session_state and rely on rerun.
-                            st.session_state['chat_messages'] = st.session_state['chat_messages'] + [{"role": "assistant", "content": text_response}]
-                            # If immediate UI update is needed and safe:
-                            # st.experimental_rerun() # Use with caution from background threads
-                except types.generation_types.StopCandidateException:
-                    print("Gemini response stream ended (StopCandidateException).")
+                        if audio_data := chunk.data:
+                            if not audio_from_gemini_playback_q.full():
+                                audio_from_gemini_playback_q.put_nowait(audio_data)
+                            else: logging.warning("Audio playback queue full, discarding Gemini audio data.")
+                        if text_response := chunk.text:
+                            logging.info(f"Gemini text response: {text_response[:100]}")
+                            if 'chat_messages' not in st.session_state: st.session_state.chat_messages = []
+                            st.session_state.chat_messages = st.session_state.chat_messages + [{"role": "assistant", "content": text_response}]
+                            # Consider st.experimental_rerun() if a mechanism exists to call it from main thread
+                except types.generation_types.StopCandidateException: logging.info("Gemini response stream ended normally.")
                 except Exception as e:
-                    if self.is_running:
-                        print(f"Error receiving from Gemini: {e}")
-                        await asyncio.sleep(0.1)
-        except asyncio.CancelledError:
-            print("process_gemini_responses task cancelled.")
-        finally:
-            print("Gemini response processing task finished.")
+                    if self.is_running: logging.error(f"Error receiving from Gemini: {e}", exc_info=True)
+                    await asyncio.sleep(0.1)
+        except asyncio.CancelledError: logging.info("Task cancelled: process_gemini_responses.")
+        finally: logging.info("Task finished: process_gemini_responses.")
 
     async def play_gemini_audio(self):
+        logging.info("Task started: Play Gemini audio responses.")
         try:
             self.playback_stream = await asyncio.to_thread(
-                pya.open, format=FORMAT, channels=CHANNELS, rate=RECEIVE_SAMPLE_RATE, output=True
+                pya.open, format=PYAUDIO_FORMAT, channels=PYAUDIO_CHANNELS, rate=GEMINI_AUDIO_RECEIVE_SAMPLE_RATE, output=True, frames_per_buffer=PYAUDIO_PLAYBACK_CHUNK_SIZE
             )
-            print("Audio playback stream opened.")
+            logging.info(f"PyAudio playback stream opened at {GEMINI_AUDIO_RECEIVE_SAMPLE_RATE} Hz.")
             while self.is_running:
                 try:
-                    audio_chunk = await asyncio.wait_for(self.audio_playback_queue.get(), timeout=1.0)
-                    if audio_chunk: # Check if not None (sentinel)
+                    audio_chunk = await asyncio.wait_for(audio_from_gemini_playback_q.get(), timeout=1.0)
+                    if audio_chunk: # Not None (sentinel)
                         await asyncio.to_thread(self.playback_stream.write, audio_chunk)
-                    if audio_chunk: # Avoid task_done on None
-                        self.audio_playback_queue.task_done()
-                except asyncio.TimeoutError:
-                    continue 
-                except Exception as e:
-                    print(f"Error playing audio: {e}")
-                    await asyncio.sleep(0.01)
+                    if audio_chunk: audio_from_gemini_playback_q.task_done()
+                except asyncio.TimeoutError: continue
+                except Exception as e: logging.error(f"Error playing audio: {e}", exc_info=True); await asyncio.sleep(0.01)
         except Exception as e:
-            # st.error(f"Failed to open audio playback stream: {e}")
-            print(f"Failed to open audio playback stream: {e}\nTraceback: {traceback.format_exc()}")
+            logging.error(f"Failed to open PyAudio playback stream: {e}", exc_info=True)
             self.is_running = False
         finally:
             if self.playback_stream:
+                logging.info("Stopping and closing PyAudio playback stream.")
                 await asyncio.to_thread(self.playback_stream.stop_stream)
                 await asyncio.to_thread(self.playback_stream.close)
                 self.playback_stream = None
-            print("Audio playback task finished.")
+            logging.info("Task finished: play_gemini_audio.")
 
     def signal_stop(self):
-        print("Signal to stop AudioVisualLoop received.")
+        logging.info("Signal to stop GeminiInteractionLoop received.")
         self.is_running = False
-        if self.media_to_gemini_queue: self.media_to_gemini_queue.put_nowait(None) 
-        if self.audio_playback_queue: self.audio_playback_queue.put_nowait(None)
+        for q in [video_frames_to_gemini_q, audio_chunks_to_gemini_q, audio_from_gemini_playback_q]:
+            try: q.put_nowait(None) # Sentinel to unblock .get()
+            except asyncio.QueueFull: logging.warning(f"Queue was full when trying to put sentinel for stop signal.")
+            except Exception as e: logging.error(f"Error putting sentinel in queue: {e}", exc_info=True)
+
 
     async def run_main_loop(self):
         self.async_event_loop = asyncio.get_running_loop()
         self.is_running = True
-        # st.session_state['session_active'] = True # Set by caller in Streamlit thread
-
-        print("AudioVisualLoop starting...")
-        # Ensure client is not None before using (should be handled by global init)
+        logging.info("GeminiInteractionLoop run_main_loop starting...")
         if client is None:
-            print("Error: Gemini client is not initialized.")
-            st.error("Critical Error: Gemini client failed to initialize. Cannot start session.")
-            st.session_state['session_active'] = False
+            logging.critical("Gemini client is None in run_main_loop. Aborting.")
             return
 
         try:
-            # Use the global 'client' instance
             async with client.aio.live.connect(model=MODEL_NAME, config=LIVE_CONNECT_CONFIG) as session:
                 self.gemini_session = session
-                print("Gemini session established.")
-
+                logging.info("Gemini session established with API.")
                 try:
-                    print("Sending system prompt to Gemini...")
+                    logging.info("Sending system prompt to Gemini...")
                     await self.gemini_session.send(input=MEDICAL_ASSISTANT_SYSTEM_PROMPT, end_of_turn=False)
-                    print("System prompt sent successfully.")
+                    logging.info("System prompt sent successfully.")
                 except Exception as e:
-                    # st.error(f"Failed to send system prompt to Gemini: {e}")
-                    print(f"Failed to send system prompt to Gemini: {e}\nTraceback: {traceback.format_exc()}")
-                    self.is_running = False 
-                    return
-
-                self.audio_playback_queue = asyncio.Queue(maxsize=AUDIO_PLAYBACK_QUEUE_MAXSIZE)
-                self.media_to_gemini_queue = asyncio.Queue(maxsize=MEDIA_OUT_QUEUE_MAXSIZE)
+                    logging.error(f"Failed to send system prompt: {e}", exc_info=True)
+                    self.is_running = False; return
 
                 async with asyncio.TaskGroup() as tg:
-                    print("Creating async tasks...")
+                    logging.info("Creating async tasks for Gemini interaction...")
                     tg.create_task(self.stream_media_to_gemini(), name="stream_media_to_gemini")
-                    tg.create_task(self.capture_microphone_audio(), name="capture_microphone_audio")
-                    
-                    if self.video_mode == "camera":
-                        tg.create_task(self.stream_camera_frames(), name="stream_camera_frames")
-                    elif self.video_mode == "screen":
-                        tg.create_task(self.stream_screen_frames(), name="stream_screen_frames")
-
                     tg.create_task(self.process_gemini_responses(), name="process_gemini_responses")
                     tg.create_task(self.play_gemini_audio(), name="play_gemini_audio")
-                    print("All async tasks created in TaskGroup.")
-                
-                print("TaskGroup finished execution.")
-
-        except asyncio.CancelledError:
-            print("AudioVisualLoop.run_main_loop() was cancelled.")
+                    logging.info("All Gemini interaction tasks created in TaskGroup.")
+                logging.info("Gemini TaskGroup finished execution.")
+        except asyncio.CancelledError: logging.info("GeminiInteractionLoop.run_main_loop() was cancelled.")
         except ExceptionGroup as eg:
-            # st.error(f"An error occurred in one of the concurrent tasks: {eg.exceptions[0]}")
-            print(f"ExceptionGroup caught in AudioVisualLoop: {eg}")
+            logging.error(f"ExceptionGroup caught in GeminiInteractionLoop: {eg}")
             for i, exc in enumerate(eg.exceptions):
-                print(f"  Task Exception {i+1}/{len(eg.exceptions)}: {type(exc).__name__}: {exc}")
-                traceback.print_exception(type(exc), exc, exc.__traceback__)
-        except Exception as e: # This catches the ConnectionClosedError
-            # st.error(f"A critical error occurred in the main session loop: {e}")
-            print(f"General Exception in AudioVisualLoop: {type(e).__name__}: {e}")
-            traceback.print_exception(e) # Print full traceback for debugging
+                logging.error(f"  Task Exception {i+1}/{len(eg.exceptions)}: {type(exc).__name__}: {exc}", exc_info=exc)
+        except Exception as e:
+            logging.error(f"General Exception in GeminiInteractionLoop: {type(e).__name__}: {e}", exc_info=True)
         finally:
-            print("AudioVisualLoop.run_main_loop() finishing...")
-            self.is_running = False 
-            # st.session_state['session_active'] = False # Set by caller in Streamlit thread
-            self.gemini_session = None 
-            print("AudioVisualLoop finished.")
+            logging.info("GeminiInteractionLoop.run_main_loop() finishing...")
+            self.is_running = False
+            self.gemini_session = None
+            logging.info("GeminiInteractionLoop finished.")
+
+# --- WebRTC Media Processors ---
+class VideoProcessor(VideoProcessorBase):
+    def __init__(self):
+        self.frame_counter = 0
+        self.last_gemini_send_time = time.monotonic()
+        # No need to get loop here if create_task is used on the default loop
+
+    async def _process_and_queue_frame_async(self, frame_ndarray):
+        self.frame_counter += 1
+        current_time = time.monotonic()
+        if (current_time - self.last_gemini_send_time) < (1.0 / VIDEO_FPS_TO_GEMINI):
+            return
 
+        self.last_gemini_send_time = current_time
+        try:
+            img_rgb = cv2.cvtColor(frame_ndarray, cv2.COLOR_BGR2RGB)
+            pil_img = PIL.Image.fromarray(img_rgb)
+            pil_img.thumbnail(VIDEO_API_RESIZE)
+            image_io = io.BytesIO()
+            pil_img.save(image_io, format="jpeg")
+            image_bytes = image_io.getvalue()
+            api_data = {"mime_type": "image/jpeg", "data": base64.b64encode(image_bytes).decode()}
+            
+            if video_frames_to_gemini_q.full():
+                try: await asyncio.wait_for(video_frames_to_gemini_q.get(), timeout=0.01)
+                except asyncio.TimeoutError: logging.warning("Video queue full, frame dropped."); return
+            video_frames_to_gemini_q.put_nowait(api_data)
+        except Exception as e: logging.error(f"Error processing/queueing video frame: {e}", exc_info=True)
+
+    async def recv(self, frame): # Called by streamlit-webrtc
+        img_bgr = frame.to_ndarray(format="bgr24")
+        asyncio.create_task(self._process_and_queue_frame_async(img_bgr))
+        return frame # Return original frame for WebRTC to display
+
+class AudioProcessor(AudioProcessorBase):
+    async def _process_and_queue_audio_async(self, audio_frames):
+        for frame in audio_frames: # frame is an AudioFrame from aiortc
+            # frame.planes[0].to_bytes() is the raw audio data
+            # frame.sample_rate, frame.layout.channels
+            # logging.info(f"Audio frame: {len(frame.planes[0].to_bytes())} bytes, SR={frame.sample_rate}, C={frame.layout.channels}")
+            
+            # CRITICAL NOTE: This sends audio as received from WebRTC.
+            # If Gemini requires a specific sample rate (e.g., 16000 Hz) and WebRTC provides
+            # a different one (e.g., 48000 Hz), audio recognition may be poor.
+            # Proper solution: Implement resampling here. This is omitted for brevity.
+            audio_data = frame.planes[0].to_bytes()
+            # Mime type should reflect the actual data being sent.
+            # Example: "audio/L16;rate=48000;channels=1" if that's what WebRTC provides.
+            # Gemini documentation should specify what it accepts for PCM.
+            # Assuming "audio/pcm" is generic enough, or be more specific.
+            # Forcing L16 (16-bit linear PCM) as that's common.
+            mime_type = f"audio/L16;rate={frame.sample_rate};channels={frame.layout.channels}"
+            api_data = {"data": audio_data, "mime_type": mime_type}
+
+            try:
+                if audio_chunks_to_gemini_q.full():
+                    try: await asyncio.wait_for(audio_chunks_to_gemini_q.get(), timeout=0.01)
+                    except asyncio.TimeoutError: logging.warning("Audio queue full, chunk dropped."); continue
+                audio_chunks_to_gemini_q.put_nowait(api_data)
+            except Exception as e: logging.error(f"Error queueing audio chunk: {e}", exc_info=True)
+    
+    async def recv(self, frames): # Called by streamlit-webrtc
+        asyncio.create_task(self._process_and_queue_audio_async(frames))
+        return frames
+
+# --- Streamlit UI and Application Logic ---
+def initialize_app_session_state():
+    defaults = {
+        'gemini_session_active': False,
+        'gemini_loop_instance': None,
+        'chat_messages': [],
+        'webrtc_component_key': f"webrtc_streamer_key_{int(time.time())}", # Initial dynamic key
+    }
+    for key, value in defaults.items():
+        if key not in st.session_state:
+            st.session_state[key] = value
 
-# --- Streamlit UI ---
 def run_streamlit_app():
+    st.set_page_config(page_title="Live AI Medical Assistant (HF Spaces)", layout="wide")
+    initialize_app_session_state() # Ensure state is initialized
+
     st.title("Live AI Medical Assistant")
-    
+    st.markdown("Utilizing Gemini Live API via WebRTC on Hugging Face Spaces")
+
     with st.sidebar:
         st.header("Session Control")
-        video_mode_options = ["camera", "screen", "none"]
-        current_video_mode = st.session_state.get('video_mode_selection', DEFAULT_VIDEO_MODE)
-        selected_video_mode = st.selectbox(
-            "Video Source:",
-            video_mode_options,
-            index=video_mode_options.index(current_video_mode),
-            disabled=st.session_state.get('session_active', False) 
-        )
-        if selected_video_mode != current_video_mode: # Update if changed and not active
-             st.session_state['video_mode_selection'] = selected_video_mode
-
-        if not st.session_state.get('session_active', False):
-            if st.button("🚀 Start Session", type="primary", use_container_width=True):
-                st.session_state['session_active'] = True # Set active before starting thread
-                st.session_state['chat_messages'] = [{
-                    "role": "system", 
-                    "content": (
-                        "Medical Assistant is activating. The AI has been instructed on its role to visually assist you. "
-                        "Remember: This AI cannot provide medical diagnoses or replace consultation with a healthcare professional."
-                    )
-                }]
-                st.session_state['current_frame_preview'] = None
-
-                audio_loop_instance = AudioVisualLoop(video_mode_setting=st.session_state['video_mode_selection'])
-                st.session_state['audio_loop_instance'] = audio_loop_instance
+        if not st.session_state.gemini_session_active:
+            if st.button("🚀 Start Session", type="primary", use_container_width=True, key="start_session_btn"):
+                st.session_state.gemini_session_active = True
+                st.session_state.chat_messages = [{"role": "system", "content": "Assistant activating. Please allow camera/microphone access in your browser if prompted."}]
                 
-                threading.Thread(target=lambda: asyncio.run(audio_loop_instance.run_main_loop()), daemon=True).start()
-                st.success("Session starting... Please wait for initialization.")
-                time.sleep(1) 
+                for q in [video_frames_to_gemini_q, audio_chunks_to_gemini_q, audio_from_gemini_playback_q]:
+                    while not q.empty(): 
+                        try: q.get_nowait()
+                        except asyncio.QueueEmpty: break
+                
+                gemini_loop = GeminiInteractionLoop()
+                st.session_state.gemini_loop_instance = gemini_loop
+                threading.Thread(target=lambda: asyncio.run(gemini_loop.run_main_loop()), name="GeminiLoopThread", daemon=True).start()
+                st.success("Gemini session starting... WebRTC will attempt to connect.")
+                st.session_state.webrtc_component_key = f"webrtc_streamer_key_{int(time.time())}" # Force re-render of WebRTC
                 st.rerun()
-        else:
-            if st.button("��� Stop Session", type="secondary", use_container_width=True):
-                if st.session_state.get('audio_loop_instance'):
-                    st.session_state['audio_loop_instance'].signal_stop()
-                st.session_state['audio_loop_instance'] = None
-                st.session_state['session_active'] = False # Set inactive
-                st.warning("Session stopping... Please wait.")
-                time.sleep(1) 
+        else: # Session is active
+            if st.button("🛑 Stop Session", type="secondary", use_container_width=True, key="stop_session_btn"):
+                if st.session_state.gemini_loop_instance:
+                    st.session_state.gemini_loop_instance.signal_stop()
+                st.session_state.gemini_loop_instance = None
+                st.session_state.gemini_session_active = False
+                st.warning("Session stopping...")
+                time.sleep(0.5) 
                 st.rerun()
-    
-    col_video, col_chat = st.columns([2, 3]) 
-    
-    with col_video:
-        st.subheader("Live Feed")
-        if st.session_state.get('session_active', False) and st.session_state.get('video_mode_selection', DEFAULT_VIDEO_MODE) != "none":
-            if st.session_state.get('current_frame_preview') is not None:
-                st.image(st.session_state['current_frame_preview'], caption="Live Video Feed", use_column_width=True)
-            else:
-                st.info("Waiting for video feed...")
-        elif st.session_state.get('video_mode_selection', DEFAULT_VIDEO_MODE) != "none":
-            st.info("Video feed will appear here once the session starts.")
-        else:
-            st.info("Video input is disabled for this session.")
-    
-    with col_chat:
-        st.subheader("Chat with Assistant")
-        # Display chat messages from session state
+
+    if st.session_state.gemini_session_active:
+        st.subheader("Your Live Feed (from your browser)")
+        RTC_CONFIGURATION = RTCConfiguration({"iceServers": [{"urls": ["stun:stun.l.google.com:19302"]}]})
+        MEDIA_STREAM_CONSTRAINTS = {
+            "video": True, # Or specific constraints like {"width": 640, "height": 480}
+            "audio": { # Request specific audio format
+                "sampleRate": {"ideal": WEBRTC_REQUESTED_SEND_SAMPLE_RATE},
+                "channelCount": {"exact": WEBRTC_REQUESTED_AUDIO_CHANNELS},
+                "echoCancellation": True, # Recommended for voice
+                "noiseSuppression": True  # Recommended for voice
+            }
+        }
+        
+        webrtc_ctx = webrtc_streamer(
+            key=st.session_state.webrtc_component_key,
+            mode=WebRtcMode.SENDONLY,
+            rtc_configuration=RTC_CONFIGURATION,
+            media_stream_constraints=MEDIA_STREAM_CONSTRAINTS,
+            video_processor_factory=VideoProcessor,
+            audio_processor_factory=AudioProcessor,
+            async_processing=True,
+            # desired_playing_state=st.session_state.gemini_session_active # Let it be controlled by rendering
+        )
+
+        if webrtc_ctx.state.playing:
+            st.caption("WebRTC connected. Streaming your camera and microphone.")
+        elif st.session_state.gemini_session_active:
+            st.caption("WebRTC attempting to connect. Ensure camera/microphone permissions are granted in your browser.")
+            if webrtc_ctx.state.error:
+                st.error(f"WebRTC Connection Error: {webrtc_ctx.state.error}")
+    else:
+        st.info("Click 'Start Session' in the sidebar to enable the live feed and assistant.")
+
+    st.subheader("Chat with Assistant")
+    chat_placeholder = st.container() # Use a container for chat messages
+    with chat_placeholder:
         for msg in st.session_state.get('chat_messages', []):
             with st.chat_message(msg["role"]):
                 st.write(msg["content"])
-        
-        user_chat_input = st.chat_input(
-            "Type your message or ask about the video...", 
-            key="user_chat_input_box", 
-            disabled=not st.session_state.get('session_active', False)
-        )
+    
+    user_chat_input = st.chat_input(
+        "Type your message...",
+        key="user_chat_input_box",
+        disabled=not st.session_state.gemini_session_active
+    )
 
-        if user_chat_input:
-            # Append user message and rerun to display it immediately
-            st.session_state['chat_messages'] = st.session_state.get('chat_messages', []) + [{"role": "user", "content": user_chat_input}]
-            
-            loop_instance = st.session_state.get('audio_loop_instance')
-            if loop_instance and loop_instance.async_event_loop and loop_instance.gemini_session:
-                if loop_instance.async_event_loop.is_running():
-                    asyncio.run_coroutine_threadsafe(
-                        loop_instance.send_text_input_to_gemini(user_chat_input),
-                        loop_instance.async_event_loop
-                    )
-                else:
-                    st.error("Session event loop is not running. Cannot send message.")
-            elif not loop_instance or not st.session_state.get('session_active', False):
-                 st.error("Session is not active. Please start a session to send messages.")
-            else:
-                 st.warning("Session components not fully ready. Please wait a moment.")
-            st.rerun() # Rerun to show user message and any quick AI text response
+    if user_chat_input:
+        current_messages = st.session_state.get('chat_messages', [])
+        current_messages.append({"role": "user", "content": user_chat_input})
+        st.session_state.chat_messages = current_messages
+        
+        loop_instance = st.session_state.get('gemini_loop_instance')
+        if loop_instance and loop_instance.async_event_loop and loop_instance.gemini_session:
+            if loop_instance.async_event_loop.is_running():
+                asyncio.run_coroutine_threadsafe(
+                    loop_instance.send_text_input_to_gemini(user_chat_input),
+                    loop_instance.async_event_loop
+                )
+            else: st.error("Session event loop is not running. Cannot send message.")
+        elif not loop_instance or not st.session_state.gemini_session_active:
+            st.error("Session is not active. Please start a session to send messages.")
+        else: st.warning("Session components not fully ready. Please wait a moment.")
+        st.rerun()
 
 if __name__ == "__main__":
-    run_streamlit_app()
+    if client is None: # Final check before running
+        logging.critical("Gemini client could not be initialized. Application cannot start.")
+    else:
+        run_streamlit_app()