import os
from pathlib import Path
import gradio as gr
import tempfile
import soundfile as sf
from models import Tokenizer, Kokoro
from fastapi import FastAPI, Request
from fastapi.responses import FileResponse
import uvicorn

# --- EXISTING LOGIC (UNCHANGED) ---

BASE_DIR = Path(__file__).resolve().parent
tokenizer_cache = None
kokoro_cache = {}
model_status = "not_loaded"
model_error = None

def get_style_vector_choices(directory="voices"):
    directory_path = BASE_DIR / directory
    return [file.name for file in directory_path.iterdir() if file.suffix == ".pt"]

def get_onnx_models(directory="weights"):
    directory_path = BASE_DIR / directory
    return [file.name for file in directory_path.iterdir() if file.suffix == ".onnx"]

def local_tts(
        text: str,
        model_path: str,
        style_vector: str,
        output_file_format: str = "wav",
        speed: float = 1.0
):
    global tokenizer_cache, model_status, model_error

    if len(text) > 0:
        try:
            style_vector_path = str(BASE_DIR / "voices" / style_vector)
            model_path_full = str(BASE_DIR / "weights" / model_path)
            cache_key = (model_path_full, style_vector_path)

            if tokenizer_cache is None:
                tokenizer_cache = Tokenizer()

            if cache_key not in kokoro_cache:
                model_status = "loading"
                model_error = None
                kokoro_cache[cache_key] = Kokoro(model_path_full, style_vector_path, tokenizer=tokenizer_cache, lang='en-us')
                model_status = "ready"

            inference = kokoro_cache[cache_key]

            audio, sample_rate = inference.generate_audio(text, speed=speed)

            with tempfile.NamedTemporaryFile(suffix=f".{output_file_format}", delete=False) as temp_file:
                sf.write(temp_file.name, audio, sample_rate)
                temp_file_path = temp_file.name

            return temp_file_path

        except Exception as e:
            model_status = "failed"
            model_error = str(e)
            raise gr.Error(f"An error occurred during TTS inference: {str(e)}")
    else:
        raise gr.Error("Input text cannot be empty.")

style_vector_choices = get_style_vector_choices()
onnx_models_choices = get_onnx_models()

sample_outputs = [
    ("Educational Note", "Machine learning models rely on large datasets and complex algorithms to identify patterns and make predictions.", str(BASE_DIR / "assets" / "edu_note.wav")),
    ("Fun Fact", "Did you know that honey never spoils? Archaeologists have found pots of honey in ancient Egyptian tombs that are over 3,000 years old and still edible!", str(BASE_DIR / "assets" / "fun_fact.wav")),
    ("Thanks", "Thank you for listening to this audio. It was generated by the Kokoro TTS model.", str(BASE_DIR / "assets" / "thanks.wav"))
]

example_texts = [
    ["Machine learning models rely on large datasets and complex algorithms to identify patterns and make predictions."],
    ["Did you know that honey never spoils? Archaeologists have found pots of honey in ancient Egyptian tombs that are over 3,000 years old and still edible!"],
    ["Thank you for listening to this audio. It was generated by the Kokoro TTS model."]
]

# --- GRADIO INTERFACE (UNCHANGED) ---

with gr.Blocks() as demo:
    gr.Markdown("## <center> Kokoro TTS ONNX Inference | [GitHub Link](https://github.com/yakhyo/kokoro-onnx) </center>")
    with gr.Row(variant="panel"):
        model_path = gr.Dropdown(choices=onnx_models_choices, label="ONNX Model Path", value=onnx_models_choices[0])
        style_vector = gr.Dropdown(choices=style_vector_choices, label="Style Vector", value=style_vector_choices[0])
        output_file_format = gr.Dropdown(choices=["wav", "mp3"], label="Output Format", value="wav")
        speed = gr.Slider(minimum=0.5, maximum=2.0, value=1.0, step=0.1, label="Speed")

    text = gr.Textbox(label="Input Text", placeholder="Enter text to convert to speech.")
    btn = gr.Button("Generate Speech")
    output_audio = gr.Audio(label="Generated Audio", type="filepath")

    btn.click(fn=local_tts, inputs=[text, model_path, style_vector, output_file_format, speed], outputs=output_audio)

    gr.Examples(examples=example_texts, inputs=[text], label="Click an example to populate the input text")
    gr.Markdown("### Sample Texts and Audio")
    for topic, sample_text, sample_audio in sample_outputs:
        with gr.Row():
            gr.Textbox(value=sample_text, label=topic, interactive=False)
            gr.Audio(value=sample_audio, label="Example Audio", type="filepath", interactive=False)

# --- FASTAPI WRAPPER & STARTUP ---

app = FastAPI()

@app.post("/v1/audio/speech")
async def api_speech(request: Request):
    """
    OpenAI-compatible /v1/audio/speech endpoint for automation.
    Expects JSON: {"input": "text", "voice": "voice_file.pt", "model": "model_file.onnx"}
    """
    data = await request.json()
    input_text = data.get("input", "")
    m_path = data.get("model", onnx_models_choices[0])
    s_vec = data.get("voice", style_vector_choices[0])
    spd = float(data.get("speed", 1.0))
    
    file_path = local_tts(input_text, m_path, s_vec, speed=spd)
    return FileResponse(file_path, media_type="audio/wav")

# Mount Gradio into the FastAPI app
app = gr.mount_gradio_app(app, demo, path="/")

if __name__ == "__main__":
    # Runs on port 7860 as expected by the Dockerfile/Hugging Face
    uvicorn.run(app, host="0.0.0.0", port=7860)