import os, re, hmac, shutil, threading, traceback
from pathlib import Path
import numpy as np
import torch
from fastapi import FastAPI, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, StreamingResponse
from pydantic import BaseModel
from typing import Optional

torch.set_num_threads(int(os.environ.get("TORCH_THREADS", "2")))

HERE = os.path.dirname(os.path.abspath(__file__))
VOICES_DIR = os.path.join(HERE, "voices")
AUDIO_EXTS = (".wav", ".mp3", ".flac", ".ogg")
SECRET = os.environ.get("TTS_SECRET_KEY", "")           # khali he to koi key nahi lagegi
DEFAULT_VOICE = os.environ.get("DEFAULT_VOICE", "")     # chahe to default awaaz ka naam yaha daalo
MODEL_REPO = "saryps-labs/pocket-tts-hindi"
LOCAL_DIR = "/home/user/hindi_model"
CONFIG_NAME = "hindi_saryps"                            # library ke config folder me is naam se jayega
FALLBACK_VOICE = "hf://kyutai/tts-voices/alba-mackenna/casual.wav"
MAX_CHARS = int(os.environ.get("MAX_CHARS", "6000"))
GAP_SECONDS = 0.22

model = None
SR = 24000
state_msg = "loading"
last_error = ""
gen_lock = threading.Lock()
voice_cache = {}


# ---------------- voices ----------------
def discover_voices() -> dict:
    """voices/ folder aur root me padi saari audio files -> {naam: path}"""
    found = {}
    for folder in (VOICES_DIR, HERE):
        if not os.path.isdir(folder):
            continue
        for fn in sorted(os.listdir(folder)):
            if fn.lower().endswith(AUDIO_EXTS):
                base = os.path.splitext(fn)[0]
                name, i = base, 2
                while name in found:
                    name = f"{base}_{i}"
                    i += 1
                found[name] = os.path.join(folder, fn)
    return found


def resolve_voice(name: Optional[str]):
    voices = discover_voices()
    if not voices:
        return "default", FALLBACK_VOICE
    if not name or name not in voices:
        name = DEFAULT_VOICE if DEFAULT_VOICE in voices else next(iter(voices))
    return name, voices[name]


def get_voice_state(path: str):
    key = path if path.startswith("hf://") else f"{path}:{os.path.getmtime(path)}"
    if key not in voice_cache:
        print("Voice process ho rahi he:", path)
        voice_cache[key] = model.get_state_for_audio_prompt(path)
    return voice_cache[key]


# ---------------- model ----------------
def download_config() -> str:
    cfg = os.path.join(LOCAL_DIR, "config.yaml")
    if not os.path.exists(cfg):
        from huggingface_hub import hf_hub_download
        hf_hub_download(MODEL_REPO, "config.yaml", local_dir=LOCAL_DIR)
    return cfg


def try_load(TTSModel, arg):
    try:
        return TTSModel.load_model(arg, temp=0.3, lsd_decode_steps=1, eos_threshold=-1.0)
    except TypeError:
        return TTSModel.load_model(arg)


def load_everything():
    global model, SR, state_msg, last_error
    errs = []
    try:
        import pocket_tts
        from pocket_tts import TTSModel

        cfg = download_config()
        print("Config file:", cfg)

        # config ko library ke apne config folder me copy karo, taaki naam se load ho sake
        pkg_cfg_dir = os.path.join(os.path.dirname(pocket_tts.__file__), "config")
        os.makedirs(pkg_cfg_dir, exist_ok=True)
        shutil.copyfile(cfg, os.path.join(pkg_cfg_dir, CONFIG_NAME + ".yaml"))
        print("Config copy hui:", os.path.join(pkg_cfg_dir, CONFIG_NAME + ".yaml"))

        attempts = [
            ("naam se (" + CONFIG_NAME + ")", CONFIG_NAME),
            ("Path object", Path(cfg)),
            ("bina .yaml ke path", os.path.splitext(cfg)[0]),
        ]
        m = None
        for label, arg in attempts:
            try:
                m = try_load(TTSModel, arg)
                print("Model load hua:", label)
                break
            except Exception as e:
                errs.append(f"[{label}] {type(e).__name__}: {str(e)[:300]}")
                traceback.print_exc()
        if m is None:
            raise RuntimeError(" || ".join(errs))

        model = m
        SR = int(getattr(m, "sample_rate", 24000))
        _, path = resolve_voice(DEFAULT_VOICE)
        get_voice_state(path)                 # default awaaz pehle se taiyar
        state_msg = "ready"
        print("Pocket TTS Hindi ready, sample rate", SR)
    except Exception as e:
        state_msg = "error"
        last_error = str(e) if errs else f"{type(e).__name__}: {e}"
        traceback.print_exc()


threading.Thread(target=load_everything, daemon=True).start()

app = FastAPI()
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
    expose_headers=["X-Sample-Rate"],
)


def authorized(request: Request) -> bool:
    if not SECRET:
        return True
    cands = [
        request.headers.get("x-secret-key", ""),
        request.query_params.get("key", ""),
    ]
    auth = request.headers.get("authorization", "")
    if auth.startswith("Bearer "):
        cands.append(auth[7:])
    return any(hmac.compare_digest(c.encode(), SECRET.encode()) for c in cands if c)


# ---------------- text / audio ----------------
def split_clean(text: str):
    text = text.replace("\r", "\n")
    parts = re.split(r"[\u0964\u0965\.\?\!\n;:]+", text)
    out = []
    for p in parts:
        p = re.sub(r"[^\u0900-\u0963\u0966-\u097F0-9\s]", " ", p)
        p = re.sub(r"\s+", " ", p).strip()
        if len(p) < 3:
            continue
        while len(p) > 220:
            cut = p.rfind(" ", 0, 200)
            cut = cut if cut > 50 else 200
            out.append(p[:cut].strip())
            p = p[cut:].strip()
        if p:
            out.append(p)
    return out


def iter_audio(text: str, vstate):
    if hasattr(model, "generate_audio_stream"):
        for a in model.generate_audio_stream(vstate, text):
            yield a
    else:
        yield model.generate_audio(vstate, text)


def to_pcm16(a) -> bytes:
    arr = a.detach().cpu().numpy() if hasattr(a, "detach") else np.asarray(a)
    arr = np.clip(arr.astype(np.float32).reshape(-1), -1.0, 1.0)
    return (arr * 32767.0).astype("<i2").tobytes()


def pcm_stream(sentences, vstate):
    silence = np.zeros(int(GAP_SECONDS * SR), dtype="<i2").tobytes()
    with gen_lock:                        # ek samay me ek hi generation
        for s in sentences:
            try:
                for a in iter_audio(s, vstate):
                    yield to_pcm16(a)
                yield silence
            except Exception:
                print("Sentence fail:", s[:60])
                traceback.print_exc()


# ---------------- API ----------------
class TTSReq(BaseModel):
    text: str
    voice: Optional[str] = None
    speed: Optional[float] = None


@app.get("/")
def root():
    return {"status": "ok"}


@app.get("/health")
def health():
    return {"state": state_msg, "error": last_error}


@app.get("/voices")
def voices(request: Request):
    if not authorized(request):
        return JSONResponse({"error": "unauthorized"}, status_code=401)
    names = list(discover_voices().keys()) or ["default"]
    default, _ = resolve_voice(DEFAULT_VOICE)
    return {"voices": names, "default": default}


@app.post("/tts")
def tts(req: TTSReq, request: Request):
    if not authorized(request):
        return JSONResponse({"error": "unauthorized"}, status_code=401)
    if state_msg == "loading":
        return JSONResponse({"error": "model abhi load ho raha he, 1 minute baad try karo"}, status_code=503)
    if state_msg == "error" or model is None:
        return JSONResponse({"error": "model load fail hua: " + last_error}, status_code=500)
    if len(req.text) > MAX_CHARS:
        return JSONResponse({"error": f"text bahut lamba he (max {MAX_CHARS} characters)"}, status_code=413)
    sentences = split_clean(req.text)
    if not sentences:
        return JSONResponse({"error": "text me Hindi (देवनागरी) shabd nahi mile"}, status_code=400)
    name, path = resolve_voice(req.voice)
    try:
        with gen_lock:                    # nayi awaaz pehli baar process hogi (kuch second lagenge)
            vstate = get_voice_state(path)
    except Exception as e:
        traceback.print_exc()
        return JSONResponse({"error": f"awaaz '{name}' load nahi hui: {e}"}, status_code=500)
    return StreamingResponse(
        pcm_stream(sentences, vstate),
        media_type="application/octet-stream",
        headers={"X-Sample-Rate": str(SR), "Cache-Control": "no-store"},
    )