import sys
import json
import base64
import io
import os
import soundfile as sf
from kokoro_onnx import Kokoro

def main():
    try:
        # Read input from stdin
        input_data = sys.stdin.read()
        if not input_data.strip():
            print(json.dumps({"error": "No input provided"}))
            sys.exit(1)
            
        params = json.loads(input_data)
        text = params.get("text", "")
        voice_name = params.get("voice_name", "af_heart")
        
        if not text:
            print(json.dumps({"error": "No text provided"}))
            sys.exit(1)
            
        # Get absolute paths of model files relative to this script
        script_dir = os.path.dirname(os.path.abspath(__file__))
        model_path = os.path.join(script_dir, "kokoro-v1.0.onnx")
        voices_path = os.path.join(script_dir, "voices-v1.0.bin")
        
        if not os.path.exists(model_path) or not os.path.exists(voices_path):
            print(json.dumps({"error": f"Model files missing at {script_dir}"}))
            sys.exit(1)
            
        # Initialize Kokoro
        kokoro = Kokoro(model_path, voices_path)
        
        # Generate audio samples
        samples, sample_rate = kokoro.create(
            text, 
            voice=voice_name, 
            speed=1.0, 
            lang="en-us"
        )
        
        # Convert to WAV in memory
        buffer = io.BytesIO()
        sf.write(buffer, samples, sample_rate, format='WAV')
        buffer.seek(0)
        
        # Encode to base64
        audio_base64 = base64.b64encode(buffer.read()).decode('utf-8')
        
        # Output JSON response
        print(json.dumps({
            "audio_base64": audio_base64,
            "sample_rate": sample_rate,
            "voice_name": voice_name
        }))
        
    except Exception as e:
        print(json.dumps({"error": str(e)}))
        sys.exit(1)

if __name__ == "__main__":
    main()