ui.py

import gradio as gr
import json
import time
import os
import asyncio
import websockets
import logging
from fastrtc import RTCComponent
import threading

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

RENDER_SIGNALING_URL = os.getenv("RENDER_SIGNALING_URL", "wss://render-signal-audio.onrender.com/stream")
HF_SPACE_URL = os.getenv("HF_SPACE_URL", "androidguy-speaker-diarization.hf.space")
WS_TRANSCRIPTION_URL = f"wss://{HF_SPACE_URL}/ws_transcription"

class TranscriptionClient:
    """Client to handle WebSocket connection to transcription service"""
    
    def __init__(self, url, on_message, on_error=None, on_close=None):
        self.url = url
        self.on_message = on_message
        self.on_error = on_error or (lambda e: logger.error(f"WebSocket error: {e}"))
        self.on_close = on_close or (lambda: logger.info("WebSocket closed"))
        self.ws = None
        self.running = False
        self.connected = False
        self.reconnect_task = None
        self.thread = None

    async def connect_async(self):
        """Connect to WebSocket server asynchronously"""
        try:
            self.ws = await websockets.connect(self.url)
            self.connected = True
            logger.info(f"Connected to {self.url}")
            
            # Start listening for messages
            while self.running:
                try:
                    message = await self.ws.recv()
                    self.on_message(message)
                except websockets.exceptions.ConnectionClosed:
                    logger.warning("Connection closed")
                    self.connected = False
                    break
                except Exception as e:
                    self.on_error(e)
                    break
                    
            # Handle connection closed
            self.connected = False
            self.on_close()
            
        except Exception as e:
            logger.error(f"Connection error: {e}")
            self.connected = False
            self.on_error(e)
    
    def connect(self):
        """Start connection in a separate thread"""
        if self.running:
            return
            
        self.running = True
        
        def run_async_loop():
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)
            loop.run_until_complete(self.connect_async())
            loop.close()
        
        self.thread = threading.Thread(target=run_async_loop, daemon=True)
        self.thread.start()
    
    def disconnect(self):
        """Disconnect from WebSocket server"""
        self.running = False
        if self.ws:
            asyncio.run(self.ws.close())
            self.ws = None
        self.connected = False

class SpeakerDiarizationUI:
    """Main UI for speaker diarization"""
    
    def __init__(self):
        self.transcription_client = None
        self.conversation_html = ""
        self.system_status = {"status": "disconnected"}
        self.webrtc_state = "stopped"
        
    def handle_transcription_message(self, message):
        """Handle incoming transcription messages"""
        try:
            data = json.loads(message)
            message_type = data.get("type", "unknown")
            
            if message_type == "conversation_update":
                # Update the conversation display
                self.conversation_html = data.get("conversation_html", "")
                # Update system status if available
                if "status" in data:
                    self.system_status = data.get("status", {})
            
            elif message_type == "connection":
                # Handle connection status update
                status = data.get("status", "unknown")
                logger.info(f"Connection status: {status}")
                if "hf_space_status" in data:
                    self.system_status["hf_space_status"] = data["hf_space_status"]
            
            elif message_type == "error":
                # Handle error message
                error_msg = data.get("message", "Unknown error")
                logger.error(f"Error from server: {error_msg}")
                
        except json.JSONDecodeError:
            logger.warning(f"Received invalid JSON: {message}")
        except Exception as e:
            logger.error(f"Error handling message: {e}")
    
    def handle_transcription_error(self, error):
        """Handle WebSocket errors"""
        logger.error(f"WebSocket error: {error}")
    
    def handle_transcription_close(self):
        """Handle WebSocket connection closure"""
        logger.info("WebSocket connection closed")
        self.system_status["status"] = "disconnected"
    
    def connect_to_transcription(self):
        """Connect to transcription WebSocket"""
        if self.transcription_client and self.transcription_client.connected:
            return
            
        self.transcription_client = TranscriptionClient(
            url=WS_TRANSCRIPTION_URL,
            on_message=self.handle_transcription_message,
            on_error=self.handle_transcription_error,
            on_close=self.handle_transcription_close
        )
        self.transcription_client.connect()
    
    def disconnect_from_transcription(self):
        """Disconnect from transcription WebSocket"""
        if self.transcription_client:
            self.transcription_client.disconnect()
            self.transcription_client = None
    
    def start_listening(self):
        """Start listening to audio and connect to services"""
        self.connect_to_transcription()
        self.webrtc_state = "started"
        return {
            webrtc: gr.update(streaming=True),
            status_display: gr.update(value=f"Status: Connected and listening"),
            start_button: gr.update(visible=False),
            stop_button: gr.update(visible=True),
            clear_button: gr.update(visible=True)
        }
    
    def stop_listening(self):
        """Stop listening to audio and disconnect from services"""
        self.disconnect_from_transcription()
        self.webrtc_state = "stopped"
        return {
            webrtc: gr.update(streaming=False),
            status_display: gr.update(value=f"Status: Disconnected"),
            start_button: gr.update(visible=True),
            stop_button: gr.update(visible=False),
            clear_button: gr.update(visible=True)
        }
    
    def clear_conversation(self):
        """Clear the conversation display"""
        # Call API to clear conversation
        import requests
        try:
            response = requests.post(f"https://{HF_SPACE_URL}/clear")
            if response.status_code == 200:
                logger.info("Conversation cleared")
            else:
                logger.error(f"Failed to clear conversation: {response.status_code}")
        except Exception as e:
            logger.error(f"Error clearing conversation: {e}")
        
        # Clear local display
        self.conversation_html = ""
        return {
            conversation_display: gr.update(value="<div class='conversation-container'><p>Conversation cleared</p></div>")
        }
    
    def update_display(self):
        """Update conversation display - called periodically"""
        status_text = f"Status: "
        if self.webrtc_state == "started":
            status_text += "Connected and listening"
        else:
            status_text += "Disconnected"
            
        if self.system_status.get("hf_space_status"):
            status_text += f" | HF Space: {self.system_status['hf_space_status']}"
            
        return {
            conversation_display: gr.update(value=self.conversation_html if self.conversation_html else "<div class='conversation-container'><p>No conversation yet. Start speaking to begin transcription.</p></div>"),
            status_display: gr.update(value=status_text)
        }

# Create UI instance
ui = SpeakerDiarizationUI()

# Custom CSS for better styling
css = """
.conversation-container {
    border: 1px solid #ddd;
    border-radius: 10px;
    padding: 15px;
    margin-bottom: 10px;
    max-height: 500px;
    overflow-y: auto;
    background-color: white;
}
.speaker {
    margin-bottom: 12px;
    border-radius: 8px;
    padding: 8px 12px;
}
.speaker-label {
    font-weight: bold;
    margin-bottom: 5px;
}
.status-display {
    margin-top: 10px;
    padding: 5px 10px;
    background-color: #f0f0f0;
    border-radius: 5px;
    font-size: 0.9rem;
}
"""

# Create Gradio interface as a function to avoid clashing
def create_interface():
    with gr.Blocks(css=css) as interface:
        gr.Markdown("# Real-Time Speaker Diarization")
        gr.Markdown("This app performs real-time speaker diarization on your audio. It automatically transcribes speech and identifies different speakers.")
        
        with gr.Row():
            with gr.Column(scale=2):
                conversation_display = gr.HTML("<div class='conversation-container'><p>No conversation yet. Start speaking to begin transcription.</p></div>")
            
            with gr.Column(scale=1):
                status_display = gr.Markdown("Status: Disconnected", elem_classes=["status-display"])
                webrtc = RTCComponent(url=RENDER_SIGNALING_URL, streaming=False, modality="audio", mode="send-receive")
                
                with gr.Row():
                    start_button = gr.Button("Start Listening", variant="primary")
                    stop_button = gr.Button("Stop Listening", variant="secondary", visible=False)
                    clear_button = gr.Button("Clear Conversation", visible=True)
                
                with gr.Accordion("Advanced Settings", open=False):
                    speaker_threshold = gr.Slider(0.5, 0.9, value=0.65, label="Speaker Change Threshold")
                    max_speakers = gr.Slider(2, 8, value=4, step=1, label="Maximum Number of Speakers")
                    
                    def update_settings(threshold, speakers):
                        import requests
                        try:
                            response = requests.post(
                                f"https://{HF_SPACE_URL}/settings",
                                params={"threshold": threshold, "max_speakers": speakers}
                            )
                            if response.status_code == 200:
                                return gr.update(value=f"Settings updated: Threshold={threshold}, Max Speakers={speakers}")
                            else:
                                return gr.update(value=f"Failed to update settings: {response.status_code}")
                        except Exception as e:
                            return gr.update(value=f"Error updating settings: {e}")
                    
                    settings_button = gr.Button("Update Settings")
                    settings_status = gr.Markdown("", elem_classes=["status-display"])
                    
                    settings_button.click(
                        update_settings,
                        [speaker_threshold, max_speakers],
                        [settings_status]
                    )
        
        # Set up event handlers
        start_button.click(
            ui.start_listening,
            [],
            [webrtc, status_display, start_button, stop_button, clear_button]
        )
        
        stop_button.click(
            ui.stop_listening,
            [],
            [webrtc, status_display, start_button, stop_button, clear_button]
        )
        
        clear_button.click(
            ui.clear_conversation,
            [],
            [conversation_display]
        )
        
        # Periodic update every 0.5 seconds
        interface.load(
            ui.update_display,
            [],
            [conversation_display, status_display],
            every=0.5
        )
        
        return interface

# Global interface instance
interface = None

# Launch the app
if __name__ == "__main__":
    interface = create_interface()
    interface.launch()

# Add mount_ui function for integration with FastAPI
def mount_ui(app):
    """Mount the Gradio interface at /ui path of the FastAPI app"""
    global interface
    
    # Create interface if it doesn't exist yet
    if interface is None:
        interface = create_interface()
    
    # Mount Gradio app at /ui path
    interface.mount_in_app(app, path="/ui")
    
    return app