import base64
import os
import tempfile
import uuid
from pathlib import Path
from threading import Lock
from typing import Dict, Optional

import requests
import torch
import torchaudio
from torchaudio.transforms import Resample
from fastapi import BackgroundTasks, Body, FastAPI, Header, HTTPException
from fastapi.responses import FileResponse, JSONResponse
from pydantic import BaseModel, Field, HttpUrl

SPACE_API_KEY = os.getenv("SPACE_API_KEY")
DUB_API_KEY = os.getenv("DUB_API_KEY")
AUTH_KEYS = {key for key in (SPACE_API_KEY, DUB_API_KEY) if key}
HF_TOKEN = (
    os.getenv("HUGGING_FACE_HUB_TOKEN")
    or os.getenv("HUGGINGFACEHUB_API_TOKEN")
    or os.getenv("HF_TOKEN")
)

MODEL_REPO = "IndexTeam/IndexTTS-2"
MODEL_DIR = os.getenv("MODEL_DIR", "/data/indextts2")
MAX_TEXT_LENGTH = 1000
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
SPEAKER_DIR = Path(tempfile.gettempdir(), "indextts2-speakers")
SPEAKER_DIR.mkdir(parents=True, exist_ok=True)

JOBS: Dict[str, Dict[str, str]] = {}
JOB_LOCK = Lock()
INFER_LOCK = Lock()

if HF_TOKEN:
    os.environ["HUGGING_FACE_HUB_TOKEN"] = HF_TOKEN
    os.environ["HF_TOKEN"] = HF_TOKEN

os.makedirs(MODEL_DIR, exist_ok=True)
from huggingface_hub import snapshot_download

if not Path(MODEL_DIR, "config.yaml").exists():
    snapshot_download(repo_id=MODEL_REPO, local_dir=MODEL_DIR, token=HF_TOKEN)

from indextts.infer_v2 import IndexTTS2

tts_model = IndexTTS2(
    cfg_path=os.path.join(MODEL_DIR, "config.yaml"),
    model_dir=MODEL_DIR,
    use_fp16=torch.cuda.is_available(),
    use_cuda_kernel=False,
    use_deepspeed=False,
)

app = FastAPI(title="indextts2-api", version="2.0.0")


class SpeakerRequest(BaseModel):
    speaker_wav: str = Field(..., description="HTTPS URL or base64 WAV")


class GenerateRequest(BaseModel):
    text: str = Field(..., min_length=1, max_length=MAX_TEXT_LENGTH)
    speaker_id: Optional[str] = None
    speaker_wav: Optional[str] = None
    emotion_wav: Optional[str] = None
    emotion_alpha: float = Field(0.9, ge=0.0, le=1.0)


def _require_api_key(x_api_key: Optional[str]):
    if AUTH_KEYS and x_api_key not in AUTH_KEYS:
        raise HTTPException(status_code=401, detail="Unauthorized")


def _write_temp_audio_from_url(url: str) -> str:
    parsed = HttpUrl(url)
    response = requests.get(str(parsed), stream=True, timeout=60)
    if response.status_code >= 400:
        raise HTTPException(status_code=400, detail="Could not fetch audio")
    suffix = Path(parsed.path).suffix or ".wav"
    with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
        for chunk in response.iter_content(chunk_size=1024 * 1024):
            if chunk:
                tmp.write(chunk)
    return tmp.name


def _write_temp_audio_from_base64(payload: str) -> str:
    if payload.startswith("data:"):
        payload = payload.split(",", 1)[-1]
    try:
        raw = base64.b64decode(payload, validate=True)
    except Exception as exc:
        raise HTTPException(status_code=400, detail="Invalid base64 audio") from exc
    with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp:
        tmp.write(raw)
    return tmp.name


def _temp_audio(value: str) -> str:
    if value.startswith("http://") or value.startswith("https://"):
        return _write_temp_audio_from_url(value)
    return _write_temp_audio_from_base64(value)


def _preprocess_audio(path: str, target_sr: int = 24000, target_peak: float = 0.98) -> str:
    wav, sr = torchaudio.load(path)
    if wav.shape[0] > 1:
        wav = wav.mean(dim=0, keepdim=True)
    if sr != target_sr:
        wav = Resample(orig_freq=sr, new_freq=target_sr)(wav)
        sr = target_sr
    peak = wav.abs().max().item() if wav.numel() else 0.0
    if peak > 0:
        wav = wav * min(target_peak / peak, 1.0)
    torchaudio.save(path, wav, sr, bits_per_sample=16)
    return path


def _cleanup_files(*files: Optional[str]):
    for file_path in files:
        if file_path:
            Path(file_path).unlink(missing_ok=True)


def _set_job(job_id: str, **kwargs):
    with JOB_LOCK:
        JOBS[job_id] = {**JOBS.get(job_id, {}), **kwargs}


def _get_job(job_id: str) -> Optional[Dict[str, str]]:
    with JOB_LOCK:
        job = JOBS.get(job_id)
        return dict(job) if job else None


def _pop_job(job_id: str) -> Optional[Dict[str, str]]:
    with JOB_LOCK:
        return JOBS.pop(job_id, None)


def _registered_speaker_path(speaker_id: str) -> Path:
    if not speaker_id or any(ch not in "0123456789abcdef-" for ch in speaker_id.lower()):
        raise HTTPException(status_code=400, detail="Invalid speaker_id")
    path = SPEAKER_DIR / f"{speaker_id}.wav"
    if not path.exists():
        raise HTTPException(status_code=404, detail="Speaker not found; register it again")
    return path


@app.post("/speakers")
def register_speaker(payload: SpeakerRequest, x_api_key: Optional[str] = Header(default=None)):
    _require_api_key(x_api_key)
    temp_path = _temp_audio(payload.speaker_wav)
    try:
        _preprocess_audio(temp_path)
        speaker_id = str(uuid.uuid4())
        final_path = SPEAKER_DIR / f"{speaker_id}.wav"
        Path(temp_path).replace(final_path)
        return {"speaker_id": speaker_id}
    except Exception:
        _cleanup_files(temp_path)
        raise


@app.delete("/speakers/{speaker_id}")
def delete_speaker(speaker_id: str, x_api_key: Optional[str] = Header(default=None)):
    _require_api_key(x_api_key)
    path = _registered_speaker_path(speaker_id)
    path.unlink(missing_ok=True)
    return {"deleted": speaker_id}


def _run_generate_job(job_id: str, payload: dict):
    speaker_temp = None
    emotion_temp = None
    output_file = None
    _set_job(job_id, status="processing")
    try:
        if payload.get("speaker_id"):
            speaker_file = str(_registered_speaker_path(payload["speaker_id"]))
        elif payload.get("speaker_wav"):
            speaker_temp = _preprocess_audio(_temp_audio(payload["speaker_wav"]))
            speaker_file = speaker_temp
        else:
            raise HTTPException(status_code=400, detail="speaker_id or speaker_wav is required")

        if payload.get("emotion_wav"):
            emotion_temp = _preprocess_audio(_temp_audio(payload["emotion_wav"]))

        output_file = os.path.join(tempfile.gettempdir(), f"indextts2-{uuid.uuid4()}.wav")
        with INFER_LOCK:
            tts_model.infer(
                spk_audio_prompt=speaker_file,
                text=payload["text"],
                output_path=output_file,
                emo_audio_prompt=emotion_temp,
                emo_alpha=float(payload.get("emotion_alpha", 0.9)),
                use_random=False,
                interval_silence=80,
                verbose=False,
            )
        _preprocess_audio(output_file)
        if not Path(output_file).exists():
            raise RuntimeError("TTS output was not created")
        _cleanup_files(speaker_temp, emotion_temp)
        _set_job(job_id, status="completed", output_file=output_file)
    except Exception as exc:
        _cleanup_files(speaker_temp, emotion_temp, output_file)
        _set_job(job_id, status="error", error=str(exc))


@app.post("/generate")
def generate(
    payload: GenerateRequest = Body(...),
    background_tasks: BackgroundTasks = BackgroundTasks(),
    x_api_key: Optional[str] = Header(default=None),
):
    _require_api_key(x_api_key)
    if not payload.speaker_id and not payload.speaker_wav:
        raise HTTPException(status_code=400, detail="speaker_id or speaker_wav is required")
    job_id = str(uuid.uuid4())
    _set_job(job_id, status="queued")
    data = payload.model_dump() if hasattr(payload, "model_dump") else payload.dict()
    background_tasks.add_task(_run_generate_job, job_id, data)
    return JSONResponse(
        status_code=202,
        content={
            "job_id": job_id,
            "status": "queued",
            "status_url": f"/status/{job_id}",
            "result_url": f"/result/{job_id}",
        },
    )


@app.get("/status/{job_id}")
def job_status(job_id: str, x_api_key: Optional[str] = Header(default=None)):
    _require_api_key(x_api_key)
    job = _get_job(job_id)
    if not job:
        raise HTTPException(status_code=404, detail="Job not found")
    response = {"job_id": job_id, "status": job.get("status", "unknown")}
    if "error" in job:
        response["error"] = job["error"]
    return response


@app.get("/result/{job_id}")
def job_result(
    job_id: str,
    background_tasks: BackgroundTasks = BackgroundTasks(),
    x_api_key: Optional[str] = Header(default=None),
):
    _require_api_key(x_api_key)
    job = _get_job(job_id)
    if not job:
        raise HTTPException(status_code=404, detail="Job not found")
    if job.get("status") != "completed":
        raise HTTPException(status_code=409, detail=f"Job not ready (status={job.get('status')})")
    output_file = job.get("output_file")
    if not output_file or not Path(output_file).exists():
        _pop_job(job_id)
        raise HTTPException(status_code=410, detail="Result expired or missing")
    _pop_job(job_id)
    background_tasks.add_task(_cleanup_files, output_file)
    return FileResponse(output_file, media_type="audio/wav", filename="output.wav")


@app.api_route("/health", methods=["GET", "POST"])
def health(x_api_key: Optional[str] = Header(default=None)):
    _require_api_key(x_api_key)
    return {"status": "ok", "model": "indextts2", "version": "2.0.0", "device": DEVICE}


@app.get("/")
def root():
    return {"name": "indextts2-api", "version": "2.0.0"}