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
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')
logging.info("Application starting up...")

# 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"
logging.info(f"Using Gemini Model: {MODEL_NAME}")

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 = None
try:
    pya = pyaudio.PyAudio()
    def cleanup_pyaudio():
        logging.info("Terminating PyAudio instance.")
        if pya: pya.terminate()
    atexit.register(cleanup_pyaudio)
    logging.info("PyAudio initialized successfully.")
except Exception as e_pyaudio:
    logging.warning(f"PyAudio initialization failed (expected in some server environments): {e_pyaudio}")
    pya = None

# --- 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)
        logging.info("Gemini client initialized successfully.")
    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 - Using audio response and Puck voice as in your latest code
# Ensure this configuration is valid for your API key and model.
LIVE_CONNECT_CONFIG = types.LiveConnectConfig(
    response_modalities=["audio"], # Requesting audio response
    speech_config=types.SpeechConfig(
        voice_config=types.VoiceConfig(
            prebuilt_voice_config=types.PrebuiltVoiceConfig(voice_name="Puck") # Using Puck voice
        )
    ) # <---------------------------------- CORRECTED: Added missing closing parenthesis
)
logging.info(f"Attempting connection with LiveConnectConfig: {LIVE_CONNECT_CONFIG}")


# --- 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]}...'")
            # Use the specific method as suggested by the deprecation warning if possible
            # For now, keeping session.send as it was working functionally
            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)
                if video_frame is None: return None # Sentinel received
                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)
                if audio_chunk is None: return None # Sentinel received
                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 is None and not self.is_running: break # Sentinel and stop signal
                if media_data and self.gemini_session and self.is_running:
                    try:
                        # Use the specific method as suggested by the deprecation warning if possible
                        # For now, keeping session.send as it was working functionally
                        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) # No data, yield
        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 = []
                            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.")
        if pya is None:
             logging.warning("PyAudio not available. Audio playback task will not run.")
             return

        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 is None and not self.is_running: break # Sentinel and stop signal
                    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 (might be expected in this environment): {e}", exc_info=True)
        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_name, q_obj_ref in [("video_q", video_frames_to_gemini_q),
                                  ("audio_in_q", audio_chunks_to_gemini_q),
                                  ("audio_out_q", audio_from_gemini_playback_q)]:
            if q_obj_ref:
                try: q_obj_ref.put_nowait(None)
                except asyncio.QueueFull: logging.warning(f"Queue {q_name} was full when trying to put sentinel for stop signal.")
                except Exception as e: logging.error(f"Error putting sentinel in {q_name}: {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

        try:
            async with client.aio.live.connect(model=MODEL_NAME, config=LIVE_CONNECT_CONFIG) as session:
                self.gemini_session = session
                logging.info(f"Gemini session established with API for model {MODEL_NAME}.")
                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

                # Using asyncio.gather for Python 3.9 compatibility
                tasks = []
                try:
                    logging.info("Creating async tasks for Gemini interaction...")
                    tasks.append(asyncio.create_task(self.stream_media_to_gemini(), name="stream_media_to_gemini"))
                    tasks.append(asyncio.create_task(self.process_gemini_responses(), name="process_gemini_responses"))
                    tasks.append(asyncio.create_task(self.play_gemini_audio(), name="play_gemini_audio"))
                    logging.info("All Gemini interaction tasks created.")
                    # Wait for tasks to complete or raise an exception
                    done, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
                    # Check results of completed tasks for errors
                    for future in done:
                        try:
                            future.result() # Raise exception if task failed
                        except Exception as task_exc:
                            logging.error(f"Task {future.get_name()} failed: {task_exc}", exc_info=True)
                            # Optionally cancel remaining tasks if one fails critically
                            for p_task in pending: p_task.cancel()
                    # If loop completes normally (e.g., user stops), pending tasks will be handled by finally block
                except Exception as e_gather: # Catch errors during task creation/gathering
                    logging.error(f"Error during task management: {e_gather}", exc_info=True)
                    for task in tasks:
                        if not task.done(): task.cancel()
                    # Wait for cancellations to complete
                    await asyncio.gather(*tasks, return_exceptions=True)
                logging.info("Gemini interaction tasks finished or cancelled.")

        except asyncio.CancelledError: logging.info("GeminiInteractionLoop.run_main_loop() was cancelled.")
        except Exception as e: # General catch-all, including ConnectionClosedError
            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
            self.signal_stop() # Ensure sentinels are sent
            # Clean up any remaining tasks (important if gather didn't complete)
            # current_tasks = [t for t in asyncio.all_tasks(self.async_event_loop) if t is not asyncio.current_task()]
            # if current_tasks:
            #     logging.info(f"Cancelling {len(current_tasks)} remaining tasks...")
            #     for task in current_tasks: task.cancel()
            #     await asyncio.gather(*current_tasks, return_exceptions=True)

            self.gemini_session = None
            video_frames_to_gemini_q = None
            audio_chunks_to_gemini_q = None
            audio_from_gemini_playback_q = None
            logging.info("GeminiInteractionLoop finished and global queues set to None.")


# --- 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."}]
                
                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")
    chat_container = st.container()
    with chat_container:
        messages = st.session_state.get('chat_messages', [])
        for msg in 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():
                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) 
                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.")
        st.rerun()

if __name__ == "__main__":
    if client is None:
        logging.critical("Gemini client could not be initialized. Application cannot start.")
    else:
        run_streamlit_app()