src/streamlit_app.py

# -*- coding: utf-8 -*-
import streamlit as st
import os
import asyncio
import base64
import io
import threading
import traceback
import atexit # Correctly imported
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 # Not needed directly

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

# Audio configuration
PYAUDIO_FORMAT = pyaudio.paInt16
PYAUDIO_CHANNELS = 1
WEBRTC_REQUESTED_AUDIO_CHANNELS = 1
WEBRTC_REQUESTED_SEND_SAMPLE_RATE = 16000
GEMINI_AUDIO_RECEIVE_SAMPLE_RATE = 24000
PYAUDIO_PLAYBACK_CHUNK_SIZE = 1024
AUDIO_PLAYBACK_QUEUE_MAXSIZE = 50
MEDIA_TO_GEMINI_QUEUE_MAXSIZE = 30

# Video configuration
VIDEO_FPS_TO_GEMINI = 2
VIDEO_API_RESIZE = (1024, 1024)

# !!! IMPORTANT: Verify this model name is correct for the Live 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:
        pya.terminate()
atexit.register(cleanup_pyaudio)

# --- Global Queues - Declare as None, initialize later ---
video_frames_to_gemini_q: asyncio.Queue = None
audio_chunks_to_gemini_q: asyncio.Queue = None
audio_from_gemini_playback_q: asyncio.Queue = None

# --- 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:
        st.error(f"Failed to initialize Gemini client: {e}")
        logging.critical(f"Gemini client initialization failed: {e}", exc_info=True)
        st.stop()
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()

# Gemini LiveConnectConfig - Simplified for debugging ConnectionClosedError
# Removed speech_config temporarily. If connection works now, the issue was voice config.
# If it still fails, check MODEL_NAME and API Key permissions.
LIVE_CONNECT_CONFIG = types.LiveConnectConfig(
    response_modalities=["audio", "text"],
    # speech_config=types.SpeechConfig( # Temporarily commented out for debugging
    #     voice_config=types.VoiceConfig(prebuilt_voice_config=types.PrebuiltVoiceConfig(voice_name="Zephyr"))
    # ),
)
logging.info(f"Using LiveConnectConfig: {LIVE_CONNECT_CONFIG}") # Log the config being used

# --- 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():
            if video_frames_to_gemini_q is None or audio_chunks_to_gemini_q is None:
                await asyncio.sleep(0.1); return None
            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
                if audio_from_gemini_playback_q is None: 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 = []
                            # This direct update might cause issues if multiple updates happen before rerun
                            st.session_state.chat_messages = st.session_state.chat_messages + [{"role": "assistant", "content": text_response}]
                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.")
        # Note: This task might fail in environments without proper ALSA setup (like HF Spaces default)
        # Error handling below will log the error but allow other tasks to continue.
        try:
            while audio_from_gemini_playback_q is None and self.is_running: await asyncio.sleep(0.1)
            if not self.is_running: return

            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: 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 chunk: {e}", exc_info=True); await asyncio.sleep(0.01)
        except Exception as e:
            logging.error(f"Failed to open or use PyAudio playback stream: {e}", exc_info=True)
            # Don't set self.is_running = False here, just log the playback failure
        finally:
            if self.playback_stream:
                logging.info("Stopping and closing PyAudio playback stream.")
                try:
                    await asyncio.to_thread(self.playback_stream.stop_stream)
                    await asyncio.to_thread(self.playback_stream.close)
                except Exception as e_close:
                     logging.error(f"Error closing playback stream: {e_close}", exc_info=True)
                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]:
            if q:
                try: q.put_nowait(None)
                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):
        global video_frames_to_gemini_q, audio_chunks_to_gemini_q, audio_from_gemini_playback_q

        self.async_event_loop = asyncio.get_running_loop()
        self.is_running = True
        logging.info("GeminiInteractionLoop run_main_loop starting...")

        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)
        logging.info("Asyncio queues initialized in GeminiInteractionLoop.")

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

        session_task_group = None # Define before try block
        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:
                    session_task_group = tg # Assign task group
                    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.")
        # Removed ExceptionGroup handling as it's not available in Python 3.9
        except Exception as e:
            # Log general exceptions, including the ConnectionClosedError if it happens here
            logging.error(f"Exception in GeminiInteractionLoop run_main_loop: {type(e).__name__}: {e}", exc_info=True)
        finally:
            logging.info("GeminiInteractionLoop.run_main_loop() finishing...")
            self.is_running = False # Ensure flag is set
            # Cancel any potentially lingering tasks if TaskGroup didn't exit cleanly (though it should)
            if session_task_group and not session_task_group._closed:
                 logging.warning("TaskGroup did not close cleanly, attempting cancellation.")
                 try:
                     session_task_group.cancel() # Attempt cancellation if needed
                 except Exception as e_cancel:
                     logging.error(f"Error cancelling task group: {e_cancel}")

            self.gemini_session = None # Session closed by async with
            # Clear global queues
            video_frames_to_gemini_q = None
            audio_chunks_to_gemini_q = None
            audio_from_gemini_playback_q = None
            logging.info("GeminiInteractionLoop finished and queues cleared.")

# --- WebRTC Media Processors ---
class VideoProcessor(VideoProcessorBase):
    def __init__(self):
        self.frame_counter = 0
        self.last_gemini_send_time = time.monotonic()

    async def _process_and_queue_frame_async(self, frame_ndarray):
        if video_frames_to_gemini_q is None: return
        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):
        img_bgr = frame.to_ndarray(format="bgr24")
        try:
            loop = asyncio.get_running_loop()
            loop.create_task(self._process_and_queue_frame_async(img_bgr))
        except RuntimeError: logging.error("VideoProcessor.recv: No running asyncio loop in current thread for create_task.")
        return frame

class AudioProcessor(AudioProcessorBase):
    async def _process_and_queue_audio_async(self, audio_frames):
        if audio_chunks_to_gemini_q is None: return
        for frame in audio_frames:
            audio_data = frame.planes[0].to_bytes()
            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):
        try:
            loop = asyncio.get_running_loop()
            loop.create_task(self._process_and_queue_audio_async(frames))
        except RuntimeError: logging.error("AudioProcessor.recv: No running asyncio loop in current thread for create_task.")
        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())}",
    }
    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()

    st.title("Live AI Medical Assistant")
    st.markdown("Utilizing Gemini Live API via WebRTC on Hugging Face Spaces")
    st.info("Remember: This AI cannot provide medical diagnoses. Always consult a healthcare professional for medical advice.")


    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."}]
                
                # Queues will be initialized inside GeminiInteractionLoop's thread
                
                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())}"
                st.rerun()
        else:
            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)")
        
        MEDIA_STREAM_CONSTRAINTS = {
            "video": True, 
            "audio": { 
                "sampleRate": {"ideal": WEBRTC_REQUESTED_SEND_SAMPLE_RATE},
                "channelCount": {"exact": WEBRTC_REQUESTED_AUDIO_CHANNELS},
                "echoCancellation": True, 
                "noiseSuppression": True  
            }
        }
        
        webrtc_ctx = webrtc_streamer(
            key=st.session_state.webrtc_component_key,
            mode=WebRtcMode.SENDONLY,
            rtc_configuration={ 
                "iceServers": [{"urls": ["stun:stun.l.google.com:19302"]}]
            },
            media_stream_constraints=MEDIA_STREAM_CONSTRAINTS,
            video_processor_factory=VideoProcessor,
            audio_processor_factory=AudioProcessor,
            async_processing=True,
        )

        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 hasattr(webrtc_ctx.state, 'error') and 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")
    # Use a container with a fixed height for scrollable chat
    chat_container = st.container()
    with chat_container:
        # Display messages from session state
        messages = st.session_state.get('chat_messages', [])
        for msg in messages:
            with st.chat_message(msg["role"]):
                st.write(msg["content"])

    # Chat input outside the container
    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 immediately for responsiveness
        st.session_state.chat_messages = st.session_state.get('chat_messages', []) + [{"role": "user", "content": user_chat_input}]
        
        # Send to backend
        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():
                # Use run_coroutine_threadsafe to call async func from Streamlit thread
                future = asyncio.run_coroutine_threadsafe(
                    loop_instance.send_text_input_to_gemini(user_chat_input),
                    loop_instance.async_event_loop
                )
                try:
                    future.result(timeout=2) # Optional: wait briefly for confirmation
                except TimeoutError:
                    logging.warning("Timed out waiting for send_text_input_to_gemini confirmation.")
                except Exception as e:
                     logging.error(f"Error calling send_text_input_to_gemini: {e}", exc_info=True)

            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.")
        
        # Rerun to display the user message and potentially any quick text response from AI
        st.rerun()

if __name__ == "__main__":
    # Final check before running the app
    if client is None:
        # Log critical error if client is still None (should have been caught earlier)
        logging.critical("Gemini client could not be initialized. Application cannot start.")
        # Display error in Streamlit if possible (might not render if client init failed early)
        st.error("CRITICAL ERROR: Gemini client failed to initialize. Check API Key and logs.")
    else:
        run_streamlit_app()