import torch
import numpy as np
from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import Response
import soundfile as sf
import io
import json
import os

app = FastAPI(title="AniTTS Vocoder Server")

# Зареждане на MioCodec - опит с различни методи
codec = None

# Метод 1: Оригинален MioCodec
try:
    from miocodec import MioCodecModel
    codec = MioCodecModel.from_pretrained("Aratako/MioCodec-25Hz-24kHz")
    codec = codec.eval()
    print("✅ MioCodec loaded (method 1)")
except Exception as e:
    print(f"⚠️ Method 1 failed: {e}")

# Метод 2: CodecV6 от BgTTS
if codec is None:
    try:
        # Добавяне на текущата директория
        import sys
        sys.path.append(os.getcwd())
        from codec import CodecV6
        codec = CodecV6(device="cpu")
        codec.model.eval()
        print("✅ CodecV6 loaded (method 2)")
    except Exception as e:
        print(f"⚠️ Method 2 failed: {e}")

# Метод 3: Dummy vocoder (винаги работи, но само бип)
if codec is None:
    print("⚠️ Using dummy vocoder - will return beep only")
    codec = "dummy"

@app.post("/vocoder")
async def vocoder_endpoint(request: Request):
    global codec
    
    # Вземане на данните от заявката
    try:
        body = await request.body()
        body_str = body.decode('utf-8')
        
        # Парсване на form data
        tokens_str = None
        embedding_str = None
        
        for part in body_str.split('&'):
            if '=' in part:
                key, val = part.split('=', 1)
                if key == 'tokens':
                    tokens_str = val
                elif key == 'embedding':
                    embedding_str = val
        
        if tokens_str is None or embedding_str is None:
            raise HTTPException(status_code=400, detail="Missing tokens or embedding")
        
        # URL декодиране
        import urllib.parse
        tokens_str = urllib.parse.unquote(tokens_str)
        embedding_str = urllib.parse.unquote(embedding_str)
        
        tokens = json.loads(tokens_str)
        speaker_emb = json.loads(embedding_str)
        
    except Exception as e:
        print(f"Parse error: {e}")
        raise HTTPException(status_code=400, detail=f"Invalid request: {e}")
    
    # Dummy vocoder (бип)
    if codec == "dummy" or codec is None:
        print(f"⚠️ Dummy vocoder for {len(tokens)} tokens")
        duration = 0.3
        sample_rate = 24000
        t = np.linspace(0, duration, int(sample_rate * duration))
        beep = 0.3 * np.sin(2 * np.pi * 440 * t)
        beep = beep * np.hanning(len(beep))
        buffer = io.BytesIO()
        sf.write(buffer, beep, sample_rate, format='wav')
        buffer.seek(0)
        return Response(content=buffer.read(), media_type="audio/wav")
    
    # Реален vocoder
    try:
        tokens_tensor = torch.tensor(tokens, dtype=torch.long)
        speaker_emb_tensor = torch.tensor(speaker_emb, dtype=torch.float32)
        
        print(f"Processing {len(tokens)} tokens...")
        
        with torch.no_grad():
            waveform = codec.decode(
                global_embedding=speaker_emb_tensor,
                content_token_indices=tokens_tensor
            )
        
        if torch.is_tensor(waveform):
            waveform = waveform.cpu().numpy()
        
        if waveform.ndim > 1:
            waveform = waveform.squeeze()
        
        # Нормализиране
        max_val = np.abs(waveform).max()
        if max_val > 0:
            waveform = waveform / max_val * 0.95
        
        buffer = io.BytesIO()
        sf.write(buffer, waveform, 24000, format='wav')
        buffer.seek(0)
        
        return Response(content=buffer.read(), media_type="audio/wav")
        
    except Exception as e:
        print(f"Decode error: {e}")
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/health")
async def health():
    return {"status": "ok", "codec_available": codec is not None}

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)