#!/usr/bin/env python3
"""
Stream LLM output to Soprano+RVC voice synthesis
Streams tokens from llama.cpp → buffered sentences → TTS → RVC → audio output
"""

import asyncio
import aiohttp
import json
import requests
import re
import time
from collections import deque

# Configuration
LLAMA_URL = "http://localhost:8091"  # llama-swap-amd container
SOPRANO_RVC_URL = "http://localhost:8765"
CONTEXT_SIZE = 2048
BATCH_SIZE = 32
TEMPERATURE = 0.8
MAX_TOKENS = 512

class StreamingVoiceGenerator:
    def __init__(self):
        self.sentence_buffer = ""
        self.sentence_endings = re.compile(r'[.!?]\s+|[.!?]$')
        self.audio_queue = deque()
        
    def _should_flush_sentence(self, text):
        """Check if we have a complete sentence to synthesize"""
        # Flush on sentence endings, or if buffer is getting long (>150 chars)
        if self.sentence_endings.search(text):
            return True
        if len(text) > 150:  # Long buffer without punctuation
            # Try to break on natural boundaries (comma, semicolon)
            if ',' in text or ';' in text:
                return True
        return False
    
    async def synthesize_sentence(self, text):
        """Send a sentence to Soprano+RVC for synthesis"""
        if not text.strip():
            return
            
        print(f"🎵 Synthesizing: {text[:60]}...")
        try:
            async with aiohttp.ClientSession() as session:
                payload = {"text": text.strip()}
                async with session.post(
                    f"{SOPRANO_RVC_URL}/api/speak",
                    json=payload,
                    timeout=aiohttp.ClientTimeout(total=30)
                ) as response:
                    if response.status == 200:
                        result = await response.json()
                        job_id = result.get('job_id')
                        print(f"✓ Queued synthesis job: {job_id[:8]}")
                    else:
                        print(f"✗ TTS error: {response.status}")
        except Exception as e:
            print(f"✗ Synthesis error: {e}")
    
    async def process_token_stream(self, tokens_generator):
        """Process tokens as they arrive, buffering into sentences"""
        synthesis_tasks = []
        
        async for token in tokens_generator:
            # Add token to buffer
            self.sentence_buffer += token
            print(token, end='', flush=True)  # Echo to console
            
            # Check if we should flush
            if self._should_flush_sentence(self.sentence_buffer):
                # Extract complete sentence
                sentence = self.sentence_buffer.strip()
                self.sentence_buffer = ""
                
                # Start synthesis task (non-blocking)
                task = asyncio.create_task(self.synthesize_sentence(sentence))
                synthesis_tasks.append(task)
        
        # Flush any remaining buffer
        if self.sentence_buffer.strip():
            task = asyncio.create_task(self.synthesize_sentence(self.sentence_buffer))
            synthesis_tasks.append(task)
        
        # Wait for all synthesis tasks to complete
        if synthesis_tasks:
            print(f"\n\n⏳ Waiting for {len(synthesis_tasks)} synthesis jobs...")
            await asyncio.gather(*synthesis_tasks)
            print("✓ All synthesis jobs queued")

async def stream_llm_tokens(prompt, system_prompt=None):
    """
    Stream tokens from llama.cpp server
    Yields tokens as they arrive
    """
    headers = {'Content-Type': 'application/json'}
    
    messages = []
    if system_prompt:
        messages.append({"role": "system", "content": system_prompt})
    messages.append({"role": "user", "content": prompt})
    
    payload = {
        "model": "llama3.1",  # Model name
        "messages": messages,
        "stream": True,  # Enable streaming
        "temperature": TEMPERATURE,
        "max_tokens": MAX_TOKENS,
        "n_predict": MAX_TOKENS,  # llama.cpp parameter
    }
    
    async with aiohttp.ClientSession() as session:
        try:
            async with session.post(
                f"{LLAMA_URL}/v1/chat/completions",
                json=payload,
                headers=headers,
                timeout=aiohttp.ClientTimeout(total=600)
            ) as response:
                if response.status != 200:
                    error = await response.text()
                    print(f"✗ LLM error: {response.status} - {error}")
                    return
                
                # Process SSE stream
                async for line in response.content:
                    line = line.decode('utf-8').strip()
                    
                    # SSE format: "data: {...}"
                    if line.startswith('data: '):
                        data_str = line[6:]  # Remove "data: " prefix
                        
                        # Check for stream end
                        if data_str == '[DONE]':
                            break
                        
                        try:
                            data = json.loads(data_str)
                            delta = data.get('choices', [{}])[0].get('delta', {})
                            token = delta.get('content', '')
                            
                            if token:
                                yield token
                        except json.JSONDecodeError:
                            continue
        
        except Exception as e:
            print(f"✗ Streaming error: {e}")

async def main():
    """Test the streaming voice pipeline"""
    
    # Test prompt
    prompt = input("Enter your message to Miku: ").strip()
    if not prompt:
        prompt = "Tell me about your favorite song!"
    
    # Miku's character
    system_prompt = """You are Hatsune Miku, the famous virtual idol and Vocaloid.
You are energetic, cheerful, and passionate about music and your fans.
Respond in Miku's voice - friendly, enthusiastic, and occasionally playful.
Keep responses natural and conversational."""
    
    print("\n" + "="*60)
    print("🎤 STREAMING VOICE SYNTHESIS TEST")
    print("="*60)
    print(f"📝 Prompt: {prompt}")
    print(f"🤖 Model: llama.cpp @ {LLAMA_URL}")
    print(f"🎵 Voice: Soprano+RVC @ {SOPRANO_RVC_URL}")
    print("="*60 + "\n")
    
    print("💬 Miku says:\n")
    
    # Create voice generator
    voice_gen = StreamingVoiceGenerator()
    
    # Stream tokens and synthesize
    start_time = time.time()
    token_generator = stream_llm_tokens(prompt, system_prompt)
    await voice_gen.process_token_stream(token_generator)
    
    elapsed = time.time() - start_time
    print(f"\n\n⏱️  Total time: {elapsed:.2f}s")
    print("✓ Done!")

if __name__ == "__main__":
    asyncio.run(main())
