import io
import os
import subprocess
import time
import wave
import numpy as np
import gradio as gr
import requests
import torch
from functools import lru_cache
from fastapi import FastAPI, HTTPException, UploadFile, File
from fastapi.responses import FileResponse, JSONResponse, Response
from zipvoice.luxvoice import LuxTTS

from pydantic import BaseModel

GROQ_API_KEY = os.environ.get("GROQ_API_KEY", "")

# Tracked so /health can report uptime and request volume - useful for
# correlating a slow day with "how long has this container been running"
# or "how many requests has it handled" without digging through raw logs.
START_TIME = time.time()
speak_request_count = 0
transcribe_request_count = 0


class ChatBody(BaseModel):
    messages: list

# Init Model (weights already downloaded at build time)
device = "cuda" if torch.cuda.is_available() else "cpu"
lux_tts = LuxTTS("YatharthS/LuxTTS", device=device, threads=2)


@lru_cache(maxsize=8)
def get_encoded_prompt(audio_prompt, ref_duration, rms):
    # encode_prompt() is deterministic for a given (file, duration, rms)
    # combo, and me.wav never changes between calls from the chat bot -
    # so we only need to pay this cost once instead of on every message.
    return lux_tts.encode_prompt(audio_prompt, duration=ref_duration, rms=rms)


# Warm the cache at startup so the very first real request doesn't
# also have to pay the reference-encoding cost.
try:
    get_encoded_prompt("me.wav", 5, 0.01)
    print("Reference prompt pre-encoded and cached.")
except Exception as e:
    print(f"Could not pre-warm reference prompt cache: {e}")


def run_tts(
    text,
    audio_prompt=None,
    rms=0.01,
    ref_duration=5,
    t_shift=0.9,
    num_steps=4,
    speed=0.8,
    return_smooth=False,
):
    # Shared core used by both the Gradio UI (infer) and the /speak
    # FastAPI route, so there's one code path for actual generation.
    if not audio_prompt:
        audio_prompt = "me.wav"

    start_time = time.time()

    encoded_prompt = get_encoded_prompt(audio_prompt, ref_duration, rms)

    final_wav = lux_tts.generate_speech(
        text,
        encoded_prompt,
        num_steps=int(num_steps),
        t_shift=t_shift,
        speed=speed,
        return_smooth=return_smooth,
    )

    duration = round(time.time() - start_time, 2)

    final_wav = final_wav.cpu().squeeze(0).numpy()
    final_wav = (np.clip(final_wav, -1.0, 1.0) * 32767).astype(np.int16)

    return final_wav, duration


def wav_bytes_from_array(int16_array, sample_rate=48000):
    # Wraps a raw int16 PCM array in a proper WAV header, in-memory -
    # no temp file on disk, so nothing to serve back over a second
    # HTTP round trip.
    buf = io.BytesIO()
    with wave.open(buf, "wb") as wf:
        wf.setnchannels(1)
        wf.setsampwidth(2)
        wf.setframerate(sample_rate)
        wf.writeframes(int16_array.tobytes())
    return buf.getvalue()


def compress_to_aac(wav_bytes, bitrate="40k"):
    # AAC-in-fragmented-MP4 instead of Opus-in-Ogg: Safari/iOS/WebKit
    # never added Ogg container support (confirmed still true as of
    # 2026), so Opus playback silently fails on iPad/iPhone even
    # though it works fine on Chrome (laptop, Android). AAC in MP4
    # plays everywhere - Safari, Chrome, Firefox, Android - with one
    # code path instead of detecting device/browser and branching.
    # "frag_keyframe+empty_moov+default_base_moof" produces a
    # streamable fragmented MP4 that doesn't need to seek back to
    # write a moov atom at the end, since we're writing to a pipe.
    # Runs entirely via pipes, no temp files.
    proc = subprocess.run(
        [
            "ffmpeg", "-i", "pipe:0",
            "-c:a", "aac", "-b:a", bitrate, "-vn",
            "-movflags", "frag_keyframe+empty_moov+default_base_moof",
            "-f", "mp4", "pipe:1",
        ],
        input=wav_bytes,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        check=True,
    )
    return proc.stdout


def infer(
    text,
    audio_prompt,
    rms,
    ref_duration,
    t_shift,
    num_steps,
    speed,
    return_smooth,
):
    if not text:
        return None, "Please provide text."

    final_wav, duration = run_tts(
        text, audio_prompt, rms, ref_duration, t_shift, num_steps, speed, return_smooth,
    )

    stats_msg = f"✨ Generation complete in **{duration}s**."
    return (48000, final_wav), stats_msg


with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🎙️ Papa Jerry Voice (LuxTTS)")

    gr.Markdown(
        """
        > **Note:** This demo runs on a **2-core CPU**, so expect slower inference.  
        > **Tip:** If words get cut off, lower **Speed** or increase **Ref Duration**.
        """
    )

    with gr.Row():
        with gr.Column():
            input_text = gr.Textbox(
                label="Text to Synthesize",
                value="Hey, what's up? I'm feeling really great!",
            )
            input_audio = gr.Audio(
                label="Reference Audio (.wav)",
                type="filepath",
                value="me.wav",
            )

            with gr.Row():
                rms_val = gr.Number(value=0.01, label="RMS (Loudness)")
                ref_duration_val = gr.Number(
                    value=5,
                    label="Reference Duration (sec)",
                    info="Lower = faster. Set ~1000 if you hear artifacts.",
                )
                t_shift_val = gr.Number(value=0.9, label="T-Shift")

            with gr.Row():
                steps_val = gr.Slider(1, 10, value=4, step=1, label="Num Steps")
                speed_val = gr.Slider(
                    0.5, 2.0, value=0.8, step=0.1,
                    label="Speed (Lower = Longer / Clearer)",
                )
                smooth_val = gr.Checkbox(label="Return Smooth", value=False)

            btn = gr.Button("Generate Speech", variant="primary")

        with gr.Column():
            audio_out = gr.Audio(label="Result")
            status_text = gr.Markdown("Ready to generate...")

    btn.click(
        fn=infer,
        inputs=[
            input_text, input_audio, rms_val, ref_duration_val,
            t_shift_val, steps_val, speed_val, smooth_val,
        ],
        outputs=[audio_out, status_text],
        api_name="predict",
    )

app = FastAPI()

@app.get("/")
def read_index():
    return FileResponse("index.html")

@app.get("/speak")
def speak(text: str):
    # GET + query param (instead of POST + JSON body) so the browser's
    # <audio> element can set this URL directly as its src and stream
    # the response progressively, playing as bytes arrive instead of
    # the page having to fetch() the whole body into a blob first.
    #
    # Also compresses to Opus before sending - measured egress out of
    # this Space is throttled to roughly 20-30 KB/s regardless of
    # payload size, and raw WAV needs ~96 KB/s to play in real time.
    # Opus at 32kbps needs ~4 KB/s, so it actually fits the available
    # bandwidth instead of just being smaller.
    if not text:
        raise HTTPException(status_code=400, detail="No text provided")

    global speak_request_count
    speak_request_count += 1
    request_start = time.time()

    final_wav, duration = run_tts(text)
    wav_bytes = wav_bytes_from_array(final_wav)

    try:
        audio_bytes = compress_to_aac(wav_bytes)
        media_type = "audio/mp4"
    except Exception as e:
        print(f"AAC compression failed, falling back to raw wav: {e}")
        audio_bytes = wav_bytes
        media_type = "audio/wav"

    total_seconds = round(time.time() - request_start, 2)
    uptime_hours = round((time.time() - START_TIME) / 3600, 2)
    print(
        f"[/speak] request #{speak_request_count} | "
        f"generation={duration}s | total={total_seconds}s | "
        f"wav={len(wav_bytes)}B -> {media_type}={len(audio_bytes)}B | "
        f"container_uptime={uptime_hours}h"
    )

    return Response(
        content=audio_bytes,
        media_type=media_type,
        headers={
            "X-Generation-Seconds": str(duration),
            # Lets the browser reuse an already-downloaded clip for the
            # same exact text (e.g. a "Replay Voice" click) instead of
            # re-requesting and regenerating it from scratch.
            "Cache-Control": "public, max-age=86400",
        },
    )


@app.post("/transcribe")
async def transcribe(file: UploadFile = File(...)):
    # Proxies the Groq Whisper call server-side, same reason /chat
    # proxies the LLM call - GROQ_API_KEY never has to live in
    # index.html or be visible via "View Page Source".
    #
    # Accepts whatever the browser's MediaRecorder produced (webm/opus
    # by default in Chrome) - Groq's Whisper endpoint auto-detects the
    # input format from the uploaded file, no client-side transcoding
    # needed.
    if not GROQ_API_KEY:
        raise HTTPException(status_code=500, detail="GROQ_API_KEY not configured on server")

    global transcribe_request_count
    transcribe_request_count += 1
    request_start = time.time()

    audio_bytes = await file.read()

    try:
        r = requests.post(
            "https://api.groq.com/openai/v1/audio/transcriptions",
            headers={"Authorization": f"Bearer {GROQ_API_KEY}"},
            files={"file": (file.filename or "audio.webm", audio_bytes, file.content_type or "audio/webm")},
            data={"model": "whisper-large-v3-turbo"},
            timeout=30,
        )
    except requests.exceptions.RequestException as e:
        raise HTTPException(status_code=502, detail=f"Groq request failed: {e}")

    total_seconds = round(time.time() - request_start, 2)
    print(
        f"[/transcribe] request #{transcribe_request_count} | "
        f"audio={len(audio_bytes)}B | total={total_seconds}s | "
        f"groq_status={r.status_code}"
    )

    if r.status_code != 200:
        raise HTTPException(status_code=r.status_code, detail=r.text)

    return JSONResponse(content=r.json())


@app.get("/health")
def health():
    uptime_seconds = round(time.time() - START_TIME, 1)
    return JSONResponse({
        "status": "ok",
        "device": device,
        "uptime_seconds": uptime_seconds,
        "uptime_hours": round(uptime_seconds / 3600, 2),
        "speak_requests_served": speak_request_count,
        "transcribe_requests_served": transcribe_request_count,
    })


@app.get("/{filename}")
def read_static_file(filename: str):
    # Serves any other flat file next to index.html/app.py in the repo -
    # me.jpg, facts.json, etc. This only matches a single path segment
    # (no slashes), so it won't intercept the /voice/... Gradio routes.
    # IMPORTANT: this catch-all must stay registered AFTER /speak and
    # /transcribe (and any other single-segment route) - FastAPI matches
    # routes in registration order, so an earlier catch-all would
    # shadow it.
    if os.path.isfile(filename):
        return FileResponse(filename)
    raise HTTPException(status_code=404, detail="Not found")


@app.post("/chat")
def chat(body: ChatBody):
    # Proxies the Groq call server-side so GROQ_API_KEY never has to
    # live in index.html or be visible via "View Page Source".
    # Plain (non-async) def so FastAPI runs this blocking network call
    # in a threadpool instead of freezing the whole server on it.
    if not GROQ_API_KEY:
        raise HTTPException(status_code=500, detail="GROQ_API_KEY not configured on server")

    r = requests.post(
        "https://api.groq.com/openai/v1/chat/completions",
        headers={
            "Content-Type": "application/json",
            "Authorization": f"Bearer {GROQ_API_KEY}",
        },
        json={
            "model": "llama-3.3-70b-versatile",
            "max_tokens": 300,
            "messages": body.messages,
        },
        timeout=30,
    )
    return JSONResponse(content=r.json(), status_code=r.status_code)


app = gr.mount_gradio_app(app, demo, path="/voice")

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