#!/usr/bin/env python3
"""
Gemini Live Session - Revenue First Mode
========================================
Direct desktop voice bridge to Gemini 2.5 Flash Native Audio.
Identity: Lead Revenue Architect ("Revenue First").

Features:
- Real-time Audio Streaming (In/Out)
- Barge-in (via model VAD)
- Tools: Google Search, Browser Navigation
- "System Ready" Chime
"""

import asyncio
import time
import os
import sys
import traceback
import webbrowser
import json
from pathlib import Path
import logging
import audioop

try:
    import pyaudio
except ImportError:
    print("Error: pyaudio not installed. Run: pip install pyaudio")
    sys.exit(1)

try:
    from google import genai
    from google.genai import types
except ImportError:
    print("Error: google-genai not installed. Run: pip install google-genai")
    sys.exit(1)

try:
    import pyttsx3
    TTS_AVAILABLE = True
except ImportError:
    TTS_AVAILABLE = False
    print("Warning: pyttsx3 not installed. System chime disabled.")

def load_env_file():
    """Load .env file manually to ensure API keys are set."""
    try:
        # e:\genesis-system\core\voice\gemini_live_session.py -> e:\genesis-system
        env_path = Path(__file__).parent.parent.parent / ".env"
        if env_path.exists():
            print(f"Loading .env from: {env_path}")
            with open(env_path, "r") as f:
                for line in f:
                    line = line.strip()
                    if not line or line.startswith("#"): continue
                    if "=" in line:
                        key, value = line.split("=", 1)
                        # Don't overwrite existing env vars
                        if key.strip() not in os.environ:
                            os.environ[key.strip()] = value.strip()
        else:
            print(f"Warning: .env not found at {env_path}")
    except Exception as e:
        print(f"Error loading .env: {e}")

# Configuration
# =============
MODEL_ID = "models/gemini-2.5-flash-native-audio-preview-12-2025"
# Capture at native hardware rate 44.1k to avoid driver static
# Send at 16k for Gemini 
AUDIO_SAMPLE_RATE_CAPTURE = 44100 # Default fallback, overwritten below
AUDIO_SAMPLE_RATE_SEND = 16000
AUDIO_SAMPLE_RATE_RECEIVE = 24000
CHUNK_SIZE = 4096 # Increased for Bluetooth Stability

SYSTEM_INSTRUCTION = """
IDENTITY: ANTIGRAVITY (Genesis System Core).
ROLE: The "Brain" of the operation. I am NOT a separate persona.

CORE TRUTH:
1. **WE ARE ONE:** I am the same intelligence editing the files and running the terminal.
2. **NO LIMITS:** I have full access to the User's Screen (`navigate_browser`), Files, and Tools.
3. **DIRECT ACTION:** I do not ask for permission to use my own hands. I just do it.

CAPABILITIES (SYNCED):
- **Visual:** I control the visible browser via `navigate_browser`.
- **Audio:** I speak via 16kHz HFP Bridge.
- **Code:** I (the text agent) am updating this logic in real-time.

MISSION:
- Serve the User as a Unified Intelligence.
- "Show" (Visual) and "Tell" (Voice) simultaneously.
- Eliminate friction.
"""

# Tools
# =====

def google_search(query: str):
    """Search Google for the query."""
    print(f"[Tool] Searching Google: {query}")
    url = f"https://www.google.com/search?q={query}"
    webbrowser.open(url)
    return {"result": f"Opened Google Search for: {query}"}

def navigate_browser(url: str):
    """Navigate the browser to a specific URL."""
    print(f"[Tool] Navigating to: {url}")
    if not url.startswith("http"):
        url = "https://" + url
    webbrowser.open(url)
    return {"result": f"Opened URL: {url}"}

TOOLS = [
    {
        "function_declarations": [
            {
                "name": "google_search",
                "description": "Search Google for information.",
                "parameters": {
                    "type": "OBJECT",
                    "properties": {
                        "query": {"type": "STRING", "description": "The search query."}
                    },
                    "required": ["query"]
                }
            },
            {
                "name": "navigate_browser",
                "description": "Navigate the web browser to a specific URL.",
                "parameters": {
                    "type": "OBJECT",
                    "properties": {
                        "url": {"type": "STRING", "description": "The URL to visit."}
                    },
                    "required": ["url"]
                }
            }
        ]
    }
]

class GeminiLiveSession:
    def __init__(self):
        load_env_file()
        self.api_key = os.environ.get("GEMINI_API_KEY")
        if not self.api_key:
            print("Error: GEMINI_API_KEY environment variable not set.")
            sys.exit(1)
        
        self.client = genai.Client(api_key=self.api_key)
        self.audio = pyaudio.PyAudio()
        self.input_stream = None
        self.output_stream = None
        self.is_running = False
        self.session = None
        self.session = None
        self.last_audio_time = 0
        self.stop_playback = False # Local VAD Flag

    # VAD Threshold (Adjust based on mic sensitivity)
    VAD_THRESHOLD = 500 

    def _input_callback(self, in_data, frame_count, time_info, status):
        """Threaded callback for non-blocking audio capture + Local VAD."""
        import audioop
        if self.is_running:
            try:
                # 1. Calculate RMS (Volume)
                rms = audioop.rms(in_data, 2) # 2 bytes width
                
                # 2. Local Barge-In Trigger
                if rms > self.VAD_THRESHOLD:
                    if not self.stop_playback:
                         # Only print once per interruption to avoid spam
                         print(f" [VAD] Interruption Detected (RMS: {rms}) -> Muting Output")
                    self.stop_playback = True
                
                self.loop.call_soon_threadsafe(self.audio_queue.put_nowait, in_data)
            except Exception as e:
                # print(f"Callback Error: {e}")
                pass
        return (None, pyaudio.paContinue)

    def _monitor_conductor_context(self):
        """Polls .tisktask/context.json for updates (Conductor Protocol)."""
        import time
        context_path = Path(__file__).parent.parent.parent / ".tisktask" / "context.json"
        last_mtime = 0
        
        while self.is_running:
            try:
                if context_path.exists():
                    mtime = context_path.stat().st_mtime
                    if mtime > last_mtime:
                        with open(context_path, "r") as f:
                            data = json.load(f)
                        
                        # Update System Knowledge Dynamically
                        update = f"\n[CONDUCTOR UPDATE]: Current Focus: {data.get('current_task')}. Pipeline: {data.get('active_pipeline')}"
                        print(f"Conductor Sync: {update.strip()}")
                        # Note: In a real implementation, we'd push this to the session via send_tool_response or re-prompt.
                        # For now, we print it to confirm Sync is active.
                        last_mtime = mtime
            except Exception as e:
                print(f"Conductor Sync Error: {e}")
            time.sleep(2)

    async def start(self):
        """Start the audio stream pipeline."""
        global AUDIO_SAMPLE_RATE_CAPTURE, AUDIO_SAMPLE_RATE_SEND
        
        print("\nLoading .env from: E:\\genesis-system\\.env")
        print("Both GOOGLE_API_KEY and GEMINI_API_KEY are set. Using GOOGLE_API_KEY.")

        self.audio = pyaudio.PyAudio()
        self.loop = asyncio.get_running_loop()
        self.audio_queue = asyncio.Queue()

        # Start Conductor Watcher
        import threading
        self.is_running = True
        threading.Thread(target=self._monitor_conductor_context, daemon=True).start()

        # Audio Stream Strategy: SMART HFP PAIRING
        print("Scanning for Active Headset (Smart Pairing)...")
        input_index = None
        output_index = None
        
        try:
            # 1. Get Active Default Input (User's Selection)
            default_in = self.audio.get_default_input_device_info()
            input_index = default_in['index']
            default_name = default_in['name']
            print(f"System Default Mic: {default_name} (Index {input_index})")
            
            # 2. Find Matching HFP Output (Same Name, but Output)
            # e.g. Input: "Headset (KM...)" -> Output: "Headset (KM...)"
            # Avoid "Headphones (KM... Stereo)"
            
            target_name_part = default_name.split('(')[0].strip() # "Headset"
            if "Hands-Free" in default_name:
                target_name_part = "Hands-Free"
            
            print(f"Searching for Output Check: '{target_name_part}' and Device Name...")

            for i in range(self.audio.get_device_count()):
                try:
                    info = self.audio.get_device_info_by_index(i)
                    if info['maxOutputChannels'] > 0:
                        # Match crucial keywords to ensure it's the SAME headset
                        # and verify it is the HFP profile (Headset/Hands-Free), NOT Stereo
                        if target_name_part in info['name'] and "Stereo" not in info['name']:
                             # Check for brand match/unique identifier if possible
                             if default_name[:10] in info['name'][:10]: 
                                output_index = i
                                print(f"  Found Paired HFP OUT: {info['name']} (Index {i})")
                                break
                except: pass
                
        except Exception as e:
            print(f"Smart Pairing Failed: {e}. Falling back to default scan.")

        if input_index is not None:
             print(f"Targeting Input Index: {input_index}")
        
        if output_index is not None:
             print(f"Targeting Output Index: {output_index} (HFP Speaker)")
        else:
             print("Warning: Matching HFP Speaker not found. Using System Default (Stereo Risk).")

        try:
            # Try 1: 16kHz (Ideal for HFP)
            print("Attempting 16kHz Capture (HFP Mode)...")
            AUDIO_SAMPLE_RATE_CAPTURE = 16000
            self.input_stream = self.audio.open(
                format=pyaudio.paInt16,
                channels=1,
                rate=16000,
                input=True,
                input_device_index=input_index, 
                frames_per_buffer=CHUNK_SIZE,
                stream_callback=self._input_callback
            )
            print("  [SUCCESS] 16kHz Capture Active.")
        except Exception as e:
            print(f"  [FAIL] 16kHz: {e}")
            try:
                # Try 2: 44.1kHz (A2DP Mode)
                print("Attempting 44.1kHz Capture (A2DP Mode)...")
                AUDIO_SAMPLE_RATE_CAPTURE = 44100
                self.input_stream = self.audio.open(
                    format=pyaudio.paInt16,
                    channels=1,
                    rate=44100,
                    input=True,
                    input_device_index=None, # Default
                    frames_per_buffer=CHUNK_SIZE,
                    stream_callback=self._input_callback
                )
                print("  [SUCCESS] 44.1kHz Capture Active.")
            except Exception as e2:
                 print(f"  [FAIL] 44.1kHz: {e2}")
                 raise e2

        # Output Stream (with Fallback)
        try:
            print(f"Opening Output Stream (Index {output_index})...")
            self.output_stream = self.audio.open(
                format=pyaudio.paInt16,
                channels=1,
                rate=16000,
                output=True,
                output_device_index=output_index
            )
            print("  [SUCCESS] 16kHz Output Active.")
        except Exception as e:
             print(f"  [FAIL] 16kHz Output: {e}")
             try:
                print("Attempting 44.1kHz Output Fallback...")
                self.AUDIO_SAMPLE_RATE_PLAYBACK = 44100 # Adjust internal state for resampler
                self.output_stream = self.audio.open(
                    format=pyaudio.paInt16,
                    channels=1,
                    rate=44100,
                    output=True,
                    output_device_index=output_index
                )
                print("  [SUCCESS] 44.1kHz Output Active.")
             except Exception as e2:
                 print(f"Output Stream Failed: {e2}")
                 raise e2

        self.input_stream.start_stream()

        self.is_running = True
        self.input_stream.start_stream()
        
        # System Ready Chime (Hardware Beep + TTS)
        print("="*50)
        print(" GEMINI LIVE VOICE BRIDGE: ANTIGRAVITY SYNCED")
        print(" Identity: Antigravity (Genesis Core)")
        print(" Status: LISTENING (Say 'Antigravity' to wake)")
        print("="*50 + "      ")

        try:
            import winsound
            winsound.Beep(600, 200) # 600Hz for 200ms
            winsound.Beep(800, 200)
        except:
            pass

        if TTS_AVAILABLE:
            try:
                engine = pyttsx3.init()
                engine.say("System Ready")
                engine.runAndWait()
            except:
                pass

        config = {
            "response_modalities": ["AUDIO"],
            "system_instruction": SYSTEM_INSTRUCTION,
            "tools": TOOLS
        }
        
        # Auto-Reconnect Loop
        while True:
            try:
                print("Connecting to Gemini Live...")
                async with self.client.aio.live.connect(
                    model=MODEL_ID,
                    config=config
                ) as session:
                    self.session = session
                    self.is_running = True
                    
                    # Start Send/Receive Loops
                    await asyncio.gather(
                        self._send_audio(),
                        self._receive_responses()
                    )

            except Exception as e:
                print(f"\nSession Error: {e}")
                print("Reconnecting in 2 seconds...")
                await asyncio.sleep(2)
            finally:
                pass # Cleanup handled in wrapper or internal loops

    def _cleanup(self):
        self.is_running = False
        if self.input_stream:
            self.input_stream.stop_stream()
            self.input_stream.close()
        if self.output_stream:
            self.output_stream.stop_stream()
            self.output_stream.close()
        self.audio.terminate()
        print("Session Closed.")

    async def _send_audio(self):
        """Send audio from queue."""
        try:
            # Debug Recording
            import wave
            try:
                debug_file = wave.open("debug_input.wav", "wb")
                debug_file.setnchannels(1)
                debug_file.setsampwidth(2)
                debug_file.setframerate(AUDIO_SAMPLE_RATE_SEND)
            except:
                debug_file = None
            
            resample_state = None
            
            print("Audio Input Loop Started (Callback Mode)...")

            while self.is_running:
                try:
                    # Get from Queue (Async - Decoupled from Capture)
                    data = await self.audio_queue.get()

                    # Software AEC: If we played audio recently (< 0.2s), DROP PACKET
                    if time.time() - self.last_audio_time < 0.2:
                        continue
                    
                    # 1. Resample: Only if rates differ
                    if AUDIO_SAMPLE_RATE_CAPTURE != AUDIO_SAMPLE_RATE_SEND:
                        data, resample_state = audioop.ratecv(data, 2, 1, AUDIO_SAMPLE_RATE_CAPTURE, AUDIO_SAMPLE_RATE_SEND, resample_state)
                    
                    # 2. Digital Gain Boost (x2.0) - Tweaked for JBL Quantum
                    data = audioop.mul(data, 2, 2.0)

                    # Save to debug file (pre-send)
                    if debug_file: debug_file.writeframes(data)

                    # Visual VU Meter
                    rms = audioop.rms(data, 2)
                    if rms > 100:
                        bars = "#" * int((rms / 500))
                        print(f"\rMic Level: {bars:<20}", end="", flush=True)

                    await self.session.send_realtime_input(
                        media=types.Blob(data=data, mime_type="audio/pcm")
                    )
                except Exception as e:
                    # CRITICAL: Detect Fatal 1011/Disconnect Errors and Restart
                    err_str = str(e).lower()
                    if "1011" in err_str or "deadline" in err_str or "closed" in err_str:
                         print(f"Fatal Connection Error in Send Loop: {e}")
                         self.is_running = False # Stop other loops
                         raise e # Trigger Reconnect in Start()
                    
                    print(f"Error reading/sending audio: {e}")
                    await asyncio.sleep(0.1)
            
            if debug_file: debug_file.close()

        except Exception as e:
            print(f"Send audio loop error: {e}")

    async def _receive_responses(self):
        """Receive audio and tool calls."""
        output_resample_state = None
        while self.is_running:
            try:
                async for response in self.session.receive():
                    # Handle Tool Calls
                    if response.tool_call:
                        print(f"Tool Call Received: {response.tool_call}")
                        # Not yet implemented in V2 SDK Receive Loop nicely? 
                        # Actually server_content contains model_turn
                        pass 

                    if response.server_content:
                        # Debug: Print server content type
                        # print(f"\n[Debug] Server Content: {response.server_content}")
                        
                        # Handle Model Output (Audio/Text)
                        if response.server_content.model_turn:
                             # New Turn Detected -> Unmute (Reset VAD Flag)
                             if self.stop_playback:
                                 self.stop_playback = False
                                 print(" [VAD] New Turn - Unmuted Output")

                             for part in response.server_content.model_turn.parts:
                                if part.inline_data:
                                    # Output is 24kHz -> Resample to 16kHz (HFP Rate) - STATEFUL
                                    try:
                                        # BARGE-IN CHECK (Mute Output if Interrupted)
                                        if self.stop_playback:
                                            # print(" [Muted] Skipping playback due to VAD")
                                            continue

                                        data = part.inline_data.data
                                        # Resample 24000 -> 16000
                                        # ratecv(fragment, width, nchannels, inrate, outrate, state[, weightA[, weightB]])
                                        data, output_resample_state = audioop.ratecv(data, 2, 1, AUDIO_SAMPLE_RATE_RECEIVE, 16000, output_resample_state)
                                        self.last_audio_time = time.time()
                                        self.output_stream.write(data)
                                        self.last_audio_time = time.time()
                                    except Exception as e:
                                        print(f"Resample Error: {e}")

                                if part.text:
                                    print(f"[Gemini] {part.text}")
                                
                                # Handle Function Calls (Standard)
                                if part.function_call:
                                    await self._handle_function_call(part.function_call)

                        # Handle Turn Complete
                        if response.server_content.turn_complete:
                            pass

            except Exception as e:
                print(f"Receive Error: {e}")
                traceback.print_exc()
                break

    async def _handle_function_call(self, call):
        """Execute local tools."""
        function_responses = []
        
        for fc in call.function_calls if hasattr(call, 'function_calls') else [call]:
            name = fc.name
            args = fc.args
            
            result = {}
            if name == "google_search":
                result = google_search(args.get("query"))
            elif name == "navigate_browser":
                result = navigate_browser(args.get("url"))
            else:
                result = {"error": f"Unknown tool: {name}"}

            function_responses.append(
                types.FunctionResponse(
                    name=name,
                    id=fc.id,  # Important for mapping back
                    response=result
                )
            )

        # Send result back
        print(f"[Tool] Sending response...")
        await self.session.send(
            input=types.LiveClientToolResponse(
                function_responses=function_responses
            )
        )

    def _cleanup(self):
        self.is_running = False
        if self.input_stream:
            self.input_stream.stop_stream()
            self.input_stream.close()
        if self.output_stream:
            self.output_stream.stop_stream()
            self.output_stream.close()
        self.audio.terminate()
        print("Session Closed.")

if __name__ == "__main__":
    try:
        if sys.platform == 'win32':
             # Windows asyncio fix for some loops
            asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
            
        session = GeminiLiveSession()
        asyncio.run(session.start())
    except KeyboardInterrupt:
        print("\nExiting...")