import os
import io
from pathlib import Path

import gradio as gr
import edge_tts
from fastapi import FastAPI, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse, JSONResponse

# ===== 1. CPU / SYSTEM RESOURCE DETECTOR =====
def effective_cpus() -> int:
    try:
        quota, period = Path("/sys/fs/cgroup/cpu.max").read_text().split()[:2]
        if quota != "max":
            return max(1, int(quota) // int(period))
    except Exception:
        pass
    try:
        return len(os.sched_getaffinity(0))
    except Exception:
        return os.cpu_count() or 2

def memory_limit_gb():
    try:
        v = Path("/sys/fs/cgroup/memory.max").read_text().strip()
        if v != "max":
            return round(int(v) / 1e9, 1)
    except Exception:
        pass
    return None

CORES = effective_cpus()
RAM = memory_limit_gb() or "?"
print(f"[resources] Effective cores: {CORES} | RAM limit: {RAM} GB")

# ===== 2. FASTAPI SERVER & CORS SETUP =====
fastapi_app = FastAPI(title="Edge-TTS-Server")
fastapi_app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# ===== 3. PARAMETER NORMALIZERS =====
def format_rate(rate_input):
    if rate_input is None:
        return "+0%"
    val = str(rate_input).strip()
    if val.endswith("%"):
        return val if val.startswith(("+", "-")) else f"+{val}"
    try:
        f = float(val)
        pct = int(round((f - 1.0) * 100))
        return f"+{pct}%" if pct >= 0 else f"{pct}%"
    except ValueError:
        return "+0%"

def format_pitch(pitch_input):
    if pitch_input is None:
        return "+0Hz"
    val = str(pitch_input).strip()
    if val.endswith("Hz") or val.endswith("%"):
        return val if val.startswith(("+", "-")) else f"+{val}"
    try:
        val_int = int(val)
        return f"+{val_int}Hz" if val_int >= 0 else f"{val_int}Hz"
    except ValueError:
        return "+0Hz"

# ===== 4. EXTERNAL STREAMING ENDPOINTS =====

# --- Direct Audio Stream (GET & POST) ---
@fastapi_app.api_route("/tts", methods=["GET", "POST"])
async def tts_stream(request: Request):
    if request.method == "POST":
        try:
            data = await request.json()
        except Exception:
            data = {}
    else:
        data = dict(request.query_params)

    text = data.get("text", "")
    if not text:
        return JSONResponse({"error": "Missing 'text' parameter"}, status_code=400)

    voice = data.get("voice", "en-US-AriaNeural")
    raw_rate = data.get("rate") or data.get("speed")
    rate = format_rate(raw_rate)
    pitch = format_pitch(data.get("pitch"))

    async def generate_audio():
        communicate = edge_tts.Communicate(text=text, voice=voice, rate=rate, pitch=pitch)
        async for chunk in communicate.stream():
            if chunk["type"] == "audio":
                yield chunk["data"]

    return StreamingResponse(
        generate_audio(),
        media_type="audio/mpeg",
        headers={
            "Cache-Control": "no-cache",
            "Content-Disposition": "inline; filename=tts.mp3"
        }
    )

# --- OpenAI Audio Speech Compatible Endpoint ---
@fastapi_app.api_route("/v1/audio/speech", methods=["POST"])
async def openai_speech(request: Request):
    try:
        data = await request.json()
    except Exception:
        data = {}

    text = data.get("input", "")
    voice = data.get("voice", "en-US-AriaNeural")
    speed = data.get("speed", 1.0)
    rate = format_rate(speed)

    async def generate_audio():
        communicate = edge_tts.Communicate(text=text, voice=voice, rate=rate)
        async for chunk in communicate.stream():
            if chunk["type"] == "audio":
                yield chunk["data"]

    return StreamingResponse(generate_audio(), media_type="audio/mpeg")

# --- Get All Available Voices ---
@fastapi_app.get("/tts/voices")
async def list_voices():
    voices = await edge_tts.list_voices()
    return JSONResponse(voices)

# ===== 5. GRADIO TEST INTERFACE =====
async def gradio_tts(text, voice, speed, pitch):
    if not text:
        return None
    rate_str = format_rate(speed)
    pitch_str = format_pitch(pitch)
    communicate = edge_tts.Communicate(text=text, voice=voice, rate=rate_str, pitch=pitch_str)
    
    buf = io.BytesIO()
    async for chunk in communicate.stream():
        if chunk["type"] == "audio":
            buf.write(chunk["data"])
    buf.seek(0)
    return buf.getvalue()

DEFAULT_VOICES = [
    "en-US-AriaNeural",
    "en-US-ChristopherNeural",
    "en-US-GuyNeural",
    "en-US-JennyNeural",
    "en-GB-SoniaNeural",
    "en-GB-RyanNeural",
    "es-ES-AlvaroNeural",
    "fr-FR-DeniseNeural",
    "de-DE-KatjaNeural",
    "zh-CN-XiaoxiaoNeural"
]

with gr.Blocks(title="High-Speed Edge-TTS API") as demo:
    gr.Markdown(
        f"# ⚡ High-Speed Edge-TTS Server\n"
        f"Running with {CORES} CPU Cores Allocated\n\n"
        f"**External API Endpoints:**\n"
        f"- `GET / POST /tts?text=...&voice=...&speed=1.0&pitch=+0Hz`\n"
        f"- `POST /v1/audio/speech` (OpenAI Compatible)\n"
        f"- `GET /tts/voices` (List All Available Edge-TTS Voices)"
    )
    with gr.Row():
        with gr.Column():
            text_input = gr.Textbox(label="Text", value="Hello! This is a real-time streaming test of Edge TTS.", lines=3)
            voice_dropdown = gr.Dropdown(choices=DEFAULT_VOICES, value="en-US-AriaNeural", label="Voice")
            speed_slider = gr.Slider(minimum=0.5, maximum=2.0, value=1.0, step=0.1, label="Speed / Rate")
            pitch_input = gr.Textbox(value="+0Hz", label="Pitch (e.g. +0Hz, +5Hz, -5Hz)")
            btn = gr.Button("Generate Speech", variant="primary")
        with gr.Column():
            audio_output = gr.Audio(label="Audio Output", autoplay=True)

    btn.click(fn=gradio_tts, inputs=[text_input, voice_dropdown, speed_slider, pitch_input], outputs=audio_output)

# ===== 6. MOUNT GRADIO ON FASTAPI & EXPORT =====
app = gr.mount_gradio_app(fastapi_app, demo, path="/")

if __name__ == "__main__":
    import uvicorn
    # Only run uvicorn locally when app.py is run directly
    port = int(os.environ.get("PORT", 7860))
    uvicorn.run("app:app", host="0.0.0.0", port=port, reload=False)