import requests
import sseclient  # Make sure this is from sseclient-py (pip install sseclient-py)
import logging
import base64
import wave
import json
import os
import sys

logging.basicConfig(level=logging.INFO, stream=sys.stdout)
logger = logging.getLogger("sse-debug")
API_URL = "http://localhost:7004/v1/audio"
TEXT = "<giggle> The quick brown fox jumps over the lazy dog"
MODEL = "orpheus"
OUT = "debug_audio.wav"

def debug_sse():
    params = {"text": TEXT, "model": MODEL}
    logger.info(f"Requesting: {API_URL} with {params}")
    try:
        response = requests.post(API_URL, params=params, stream=True)
        logger.info(f"HTTP {response.status_code}")
        logger.info(f"Headers: {dict(response.headers)}")
        logger.info(f"Content-Type: {response.headers.get('content-type','')}")
        if 'text/event-stream' not in response.headers.get('content-type', ''):
            logger.warning("Server did NOT send the expected text/event-stream Content-Type!")

        client = sseclient.SSEClient(response)  # Do not read from response first!
        print("Connected to SSE, waiting for events:")
        all_bytes = bytearray()
        sample_rate = None

        for event in client.events():
            logger.info(f"Received SSE event: {event.data[:70]}{'...' if len(event.data) > 70 else ''}")
            if event.data == '[DONE]':
                logger.info("Stream finished, got [DONE].")
                break
            try:
                data = json.loads(event.data)
                if 'audio' in data:
                    chunk = base64.b64decode(data['audio'])
                    all_bytes.extend(chunk)
                    if not sample_rate:
                        sample_rate = data.get('rate')
            except Exception as e:
                logger.error(f"Failed to parse SSE data: {event.data} ({e})")

        if all_bytes and sample_rate:
            with wave.open(OUT, 'wb') as f:
                f.setnchannels(1)
                f.setsampwidth(2)
                f.setframerate(sample_rate)
                f.writeframes(all_bytes)
            logger.info(f"Saved audio to {os.path.abspath(OUT)}")
        else:
            logger.warning("No audio chunks or sample rate detected.")

    except Exception as e:
        logger.exception(f"Error during SSE client run: {e}")

if __name__ == "__main__":
    debug_sse()