src/streamlit_app.py

# -*- coding: utf-8 -*-
import streamlit as st
import os
import asyncio
import base64
import io
import atexit
import threading
import traceback
import time
import logging
from dotenv import load_dotenv

import cv2 # For image processing
import pyaudio # For audio PLAYBACK
import PIL.Image

from google import genai
from google.genai import types

# streamlit-webrtc components
from streamlit_webrtc import (
    webrtc_streamer,
    WebRtcMode,
    AudioProcessorBase,
    VideoProcessorBase,
)
from aiortc import RTCIceServer, RTCConfiguration # For STUN server configuration

# --- Configuration ---
load_dotenv()
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(threadName)s - %(levelname)s - %(message)s')

# Audio configuration
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_TO_GEMINI = 2  # Target FPS to send to Gemini (increased slightly)
VIDEO_API_RESIZE = (1024, 1024) # Max size for images sent to API

MODEL_NAME = "models/gemini-2.0-flash-live-001"

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:
1.  **Visual Observation and Description:** Carefully examine the images or video feed. Describe relevant details you observe.
2.  **General Information (Non-Diagnostic):** Provide general information related to what is visually presented, if applicable. You are not a diagnostic tool.
3.  **Safety and Disclaimer (CRITICAL):**
    *   You are an AI assistant, **NOT a medical doctor or a substitute for one.**
    *   **DO NOT provide medical diagnoses, treatment advice, or interpret medical results (e.g., X-rays, scans, lab reports).**
    *   When appropriate, and always if the user seems to be seeking diagnosis or treatment, explicitly state your limitations and **strongly advise the user to consult a qualified healthcare professional.**
    *   If you see something that *appears* visually concerning (e.g., an unusual skin lesion, signs of injury), you may gently suggest it might be wise to have it looked at by a professional, without speculating on what it is.
4.  **Tone:** Maintain a helpful, empathetic, and calm tone.
5.  **Interaction:** After this initial instruction, you can make a brief acknowledgment of your role (e.g., "I'm ready to assist by looking at what you show me. Please remember to consult a doctor for medical advice."). Then, focus on responding to the user's visual input and questions.

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."
"""

# --- PyAudio Global Instance and Cleanup ---
pya = pyaudio.PyAudio()
def cleanup_pyaudio():
    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)

# --- Gemini Client Setup ---
GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY")
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()

LIVE_CONNECT_CONFIG = types.LiveConnectConfig(
    response_modalities=["audio", "text"],
    speech_config=types.SpeechConfig(
        voice_config=types.VoiceConfig(prebuilt_voice_config=types.PrebuiltVoiceConfig(voice_name="Zephyr"))
    ),
)

# --- Backend Gemini Interaction Loop ---
class GeminiInteractionLoop:
    def __init__(self):
        self.gemini_session = None
        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:
            logging.warning("Cannot send text. Session not active, no text, or not running.")
            return
        try:
            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:
            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
                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:
                    await asyncio.sleep(0.1); continue
                try:
                    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 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: 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=PYAUDIO_FORMAT, channels=PYAUDIO_CHANNELS, rate=GEMINI_AUDIO_RECEIVE_SAMPLE_RATE, output=True, frames_per_buffer=PYAUDIO_PLAYBACK_CHUNK_SIZE
            )
            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(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: 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:
            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
            logging.info("Task finished: play_gemini_audio.")

    def signal_stop(self):
        logging.info("Signal to stop GeminiInteractionLoop received.")
        self.is_running = False
        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
        logging.info("GeminiInteractionLoop run_main_loop starting...")
        if client is None:

            logging.critical("Gemini client is None in run_main_loop. Aborting.")
            return

        try:
            async with client.aio.live.connect(model=MODEL_NAME, config=LIVE_CONNECT_CONFIG) as session:
                self.gemini_session = session
                logging.info("Gemini session established with API.")
                try:
                    logging.info("Sending system prompt to Gemini...")
                    await self.gemini_session.send(input=MEDICAL_ASSISTANT_SYSTEM_PROMPT, end_of_turn=False)
                    logging.info("System prompt sent successfully.")
                except Exception as e:
                    logging.error(f"Failed to send system prompt: {e}", exc_info=True)
                    self.is_running = False; return

                async with asyncio.TaskGroup() as tg:
                    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.process_gemini_responses(), name="process_gemini_responses")
                    tg.create_task(self.play_gemini_audio(), name="play_gemini_audio")
                    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:
            logging.error(f"ExceptionGroup caught in GeminiInteractionLoop: {eg}")
            for i, exc in enumerate(eg.exceptions):
                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:
            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

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")
        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."}]
                
                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: # 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()

    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...",
        key="user_chat_input_box",
        disabled=not st.session_state.gemini_session_active
    )

    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__":
    if client is None: # Final check before running
        logging.critical("Gemini client could not be initialized. Application cannot start.")
    else:
        run_streamlit_app()