import os
import asyncio
import json
import base64
import requests
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from groq import Groq

app = FastAPI()

# Configuration
GROQ_API_KEY = os.environ.get("GROQ_API_KEY", "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")
SARVAM_API_KEY = os.environ.get("SARVAM_API_KEY", "sk_ibddkc1o_5ByGCOrIePXVZihmXubVUgTm")
groq_client = Groq(api_key=GROQ_API_KEY)

conversation_history = [
    {"role": "system", "content": "You are a real-time conversational voice assistant. Keep answers incredibly brief, punchy, and conversational."}
]

@app.websocket("/conversational-stream")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    print("🚀 Client connected to real-time voice stream.")
    
    # Task management for handling interruptions smoothly
    tts_task = None

    try:
        while True:
            # Wait for incoming messages from the client (Laptop/Mobile)
            data = await websocket.receive_text()
            message = json.loads(data)
            
            # --- TALK TO INTERRUPT LOGIC ---
            # If the client transmits a "USER_INTERRUPT" signal or new audio data 
            # while the assistant is currently speaking, kill the speaking task immediately.
            if message.get("type") == "INTERRUPT" or message.get("type") == "AUDIO_START":
                if tts_task and not tts_task.done():
                    tts_task.cancel()
                    print("🛑 AI Interrupted by user speech! Stopping current audio output.")
                if message.get("type") == "INTERRUPT":
                    continue

            # Process incoming transcript from user
            if message.get("type") == "TEXT_INPUT":
                user_text = message.get("text")
                conversation_history.append({"role": "user", "content": user_text})
                
                # Generate LLM response
                completion = groq_client.chat.completions.create(
                    model="openai/gpt-oss-20b",
                    messages=conversation_history
                )
                ai_text = completion.choices[0].message.content
                conversation_history.append({"role": "assistant", "content": ai_text})
                
                # Start speaking via Sarvam AI in a cancelable async task
                tts_task = asyncio.create_task(stream_tts_to_client(websocket, ai_text))

    except WebSocketDisconnect:
        print("🔌 Client disconnected.")
    except asyncio.CancelledError:
        pass

async def stream_tts_to_client(websocket: WebSocket, text: str):
    """Generates audio via Sarvam and pushes chunks down the WebSocket securely."""
    try:
        # Requesting audio generation
        headers = {"api-subscription-key": SARVAM_API_KEY, "Content-Type": "application/json"}
        payload = {"text": text, "voice": "bulbul:v3", "language_code": "en-IN", "output_format": "wav"}
        
        response = requests.post("https://api.sarvam.ai/text-to-speech", headers=headers, json=payload)
        if response.status_code == 200:
            audio_base64 = response.json().get("audios", [None])[0]
            if audio_base64:
                # To facilitate fast interruption delivery, we break the raw audio down into 
                # micro-chunks and pipe them iteratively with quick async breathers.
                audio_bytes = base64.b64decode(audio_base64)
                chunk_size = 4096 
                
                for i in range(0, len(audio_bytes), chunk_size):
                    chunk = audio_bytes[i:i+chunk_size]
                    await websocket.send_json({
                        "type": "AUDIO_CHUNK",
                        "audio": base64.b64encode(chunk).decode('utf-8')
                    })
                    # Yield control briefly to let the event loop process incoming interruption events
                    await asyncio.sleep(0.01) 
                    
                await websocket.send_json({"type": "AUDIO_END"})
    except asyncio.CancelledError:
        # This clean-up segment fires automatically when tts_task.cancel() is invoked
        print("🧼 TTS Streaming task cleanly aborted.")