from fastapi import FastAPI, File, Query, Response, UploadFile
from fastapi.encoders import jsonable_encoder
from fastapi.responses import FileResponse
from config import EDGETTS_VOICE, TOKEN
import edge_tts
import hmac
import model
from pydantic import BaseModel
from typing import Annotated, List, Optional
import tempfile

from rvc.main import song_cover_pipeline

app = FastAPI()

class Message(BaseModel):
    role: str
    content: str

@app.post("/")
async def root(token: str,
    messages: List[Message],
    response: Response,
    max_new_tokens: Optional[int] = 64,
    temperature: Optional[float] = 0.9,
    repetition_penalty: Optional[float] = 1.2):
    if not hmac.compare_digest(token, TOKEN):
        response.status_code = 401
        return {"error": "Bad token"}

    dict_in = jsonable_encoder(messages)
    output = model.inference(dict_in, max_new_tokens=max_new_tokens, temperature=temperature, repetition_penalty=repetition_penalty)
    return {"raw": output}

@app.post("/rvc")
async def rvc(token: str,
    file: Annotated[bytes, File()],
    response: Response,
    pitch_change_oct: Annotated[int, Query()] = 1,
    pitch_change_sem: Annotated[int, Query()] = 0):
    if not hmac.compare_digest(token, TOKEN):
        response.status_code = 401
        return {"error": "Bad token"}

    with tempfile.NamedTemporaryFile() as tmp:
        tmp.write(file)
        ai_vocals_path = song_cover_pipeline(tmp.name, pitch_change_oct, voice_model='miku', pitch_change_sem=pitch_change_sem)

    return FileResponse(ai_vocals_path)

@app.post("/tts")
async def tts(token: str,
    text: str,
    response: Response,
    pitch_change_oct: Annotated[int, Query()] = 1,
    pitch_change_sem: Annotated[int, Query()] = 0):
    if not hmac.compare_digest(token, TOKEN):
        response.status_code = 401
        return {"error": "Bad token"}
    
    with tempfile.NamedTemporaryFile() as tmp:
        communicate = edge_tts.Communicate(text, EDGETTS_VOICE)
        await communicate.save(tmp.name)
        ai_vocals_path = song_cover_pipeline(tmp.name, pitch_change_oct, voice_model='miku', pitch_change_sem=pitch_change_sem)

    return FileResponse(ai_vocals_path)


@app.get("/ping")
def ping():
    return {"message": "pong"}