import os, re, json, tempfile
import numpy as np

from fastapi import FastAPI, HTTPException
from fastapi.responses import FileResponse, HTMLResponse, JSONResponse
from pydantic import BaseModel
import onnxruntime as ort
import soundfile as sf
from huggingface_hub import hf_hub_download

SAMPLE_RATE=24000; DEC_N_LAYERS=8; DEC_D_MODEL=384; DEC_N_HEADS=6
SPEAKER_EMB_DIM=128; VOCAB_SIZE=12955; AUDIO_OFFSET=155; EOS_TOKEN=1; BOS_TOKEN=0

app = FastAPI()

def load_voices():
    voices={}
    d="voices"
    if not os.path.exists(d): return voices
    for f in sorted(os.listdir(d)):
        if not f.endswith(".json"): continue
        with open(os.path.join(d,f)) as fh: data=json.load(fh)
        label=data.get("name",f.replace(".json","").upper())
        voices[label]=np.array(data["embedding"],dtype=np.float32)
    return voices

VOICES=load_voices(); VOICE_NAMES=list(VOICES.keys())

class TTSTokenizer:
    SPECIAL=["<pad>","<eos>","<bos>","<unk>","<sep>","<mask>","<blank>","<sil>","<br>"]
    TEXT_CHARS=("АБВГДЕЖЗИЙКЛМНОПРСТУФХЦЧШЩЪЬЮЯабвгдежзийклмнопрстуфхцчшщъьюя"
        "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
        "0123456789 .,!?-:;\"'()…–—\n\t")
    def __init__(self):
        self.vocab=self.SPECIAL+list(self.TEXT_CHARS)
        self.char2id={c:i for i,c in enumerate(self.vocab)}
        self.unk_id=self.char2id["<unk>"]
    def encode(self,text):
        return [BOS_TOKEN]+[self.char2id.get(c,self.unk_id) for c in text]+[EOS_TOKEN]
    def split_sentences(self,text,max_chars=200):
        parts=re.split(r'(?<=[.!?\n])\s+',text.strip())
        result,current=[],""
        for p in parts:
            if len(current)+len(p)+1<=max_chars: current=(current+" "+p).strip()
            else:
                if current: result.append(current)
                current=p[:max_chars]
        if current: result.append(current)
        return result or [text[:max_chars]]

_sessions={}
def get_session(name):
    if name not in _sessions:
        path = hf_hub_download(repo_id="ssasio/AniTTS-Onix", filename=name, repo_type="dataset")
        model_dir = os.path.dirname(path)

        data_file = name + ".data"
        try:
            data_path = hf_hub_download(repo_id="ssasio/AniTTS-Onix", filename=data_file, repo_type="dataset")
            expected = os.path.join(model_dir, data_file)
            if not os.path.exists(expected) and data_path != expected:
                os.symlink(data_path, expected)
        except Exception:
            pass

        opts = ort.SessionOptions()
        opts.inter_op_num_threads = 4
        opts.intra_op_num_threads = 4
        _sessions[name] = ort.InferenceSession(
            path, sess_options=opts, providers=["CPUExecutionProvider"]
        )
    return _sessions[name]

def softmax(x):
    e=np.exp(x-x.max()); return e/e.sum()

def sample_token(logits,top_k,top_p,temperature,rep_penalty,recent):
    logits=logits.copy().astype(np.float32)
    if temperature>0: logits/=temperature
    for t in set(recent or []):
        logits[t]=logits[t]/rep_penalty if logits[t]>0 else logits[t]*rep_penalty
    if top_k>0:
        cutoff=np.sort(logits)[::-1][min(top_k,len(logits)-1)]
        logits[logits<cutoff]=-np.inf
    idx=np.argsort(logits)[::-1]; cum=np.cumsum(softmax(logits[idx]))
    rem=np.roll(cum>top_p,1); rem[0]=False; logits[idx[rem]]=-np.inf
    probs=softmax(logits); return int(np.random.choice(len(probs),p=probs))

def encode_text(input_ids,attn_mask):
    return get_session("encoder.onnx").run(["enc_out"],{"input_ids":input_ids,"attention_mask":attn_mask})[0]

def decode_step(token,enc_out,enc_mask,spk,past_kv):
    feed={"input_ids":np.array([[token]],dtype=np.int64),"enc_out":enc_out,"enc_mask":enc_mask,"speaker_emb":spk}
    for i,(sk,ck) in enumerate(past_kv):
        b=i*4; feed[f"p{b}"]=sk[0]; feed[f"p{b+1}"]=sk[1]; feed[f"p{b+2}"]=ck[0]; feed[f"p{b+3}"]=ck[1]
    outs=get_session("decoder.onnx").run(["logits"]+[f"v{i}" for i in range(32)],feed)
    new_kv=[((outs[1+i*4],outs[2+i*4]),(outs[3+i*4],outs[4+i*4])) for i in range(DEC_N_LAYERS)]
    return outs[0][0,0],new_kv

def generate_sentence(text,spk_emb,temperature,top_k,top_p,rep_penalty,max_tokens):
    tok=TTSTokenizer()
    ids=tok.encode(text)
    
    ENC_LEN = 128
    PAST_LEN = 64
    
    T = len(ids)
    
    print(f"Original text length: {len(text)}, tokens: {T}")
    
    if T > ENC_LEN:
        ids = tok.encode(text[:ENC_LEN - 2])
        T = len(ids)
        print(f"Truncated to {T} tokens")
    
    if T < ENC_LEN:
        pad_len = ENC_LEN - T
        input_ids = np.array([ids + [0] * pad_len], dtype=np.int64)
        attn_mask = np.array([[1] * T + [0] * pad_len], dtype=np.int64)
    else:
        input_ids = np.array([ids[:ENC_LEN]], dtype=np.int64)
        attn_mask = np.array([[1] * ENC_LEN], dtype=np.int64)
    
    print(f"Final input_ids shape: {input_ids.shape}")
    
    enc_out = encode_text(input_ids, attn_mask)
    
    H = DEC_D_MODEL // DEC_N_HEADS
    
    past_kv = []
    for _ in range(DEC_N_LAYERS):
        sk_self = (np.zeros((1, DEC_N_HEADS, PAST_LEN, H), dtype=np.float32),
                   np.zeros((1, DEC_N_HEADS, PAST_LEN, H), dtype=np.float32))
        ck_self = (np.zeros((1, DEC_N_HEADS, ENC_LEN, H), dtype=np.float32),
                   np.zeros((1, DEC_N_HEADS, ENC_LEN, H), dtype=np.float32))
        past_kv.append((sk_self, ck_self))
    
    cur, generated, recent = BOS_TOKEN, [], []
    spk = spk_emb.reshape(1, SPEAKER_EMB_DIM)
    
    for _ in range(max_tokens):
        logits, past_kv_new = decode_step(cur, enc_out, attn_mask, spk, past_kv)
        
        past_kv_updated = []
        for i in range(DEC_N_LAYERS):
            (sk_k, sk_v), ck = past_kv_new[i]
            cur_len = sk_k.shape[2]
            if cur_len > PAST_LEN:
                sk_k = sk_k[:, :, -PAST_LEN:, :]
                sk_v = sk_v[:, :, -PAST_LEN:, :]
            elif cur_len < PAST_LEN:
                pad = PAST_LEN - cur_len
                pad_k = np.zeros((1, DEC_N_HEADS, pad, H), dtype=np.float32)
                pad_v = np.zeros((1, DEC_N_HEADS, pad, H), dtype=np.float32)
                sk_k = np.concatenate([pad_k, sk_k], axis=2)
                sk_v = np.concatenate([pad_v, sk_v], axis=2)
            past_kv_updated.append(((sk_k, sk_v), ck))
        past_kv = past_kv_updated
        
        mask = np.full(VOCAB_SIZE, -np.inf, np.float32)
        mask[EOS_TOKEN] = logits[EOS_TOKEN]
        mask[AUDIO_OFFSET:] = logits[AUDIO_OFFSET:]
        cur = sample_token(mask, top_k, top_p, temperature, rep_penalty, recent[-50:])
        
        if cur == EOS_TOKEN:
            break
        
        idx = cur - AUDIO_OFFSET
        if 0 <= idx < 12800:
            generated.append(idx)
            recent.append(cur)
    
    return np.array(generated, dtype=np.int64) if generated else None

class SynthRequest(BaseModel):
    text: str
    voice: str
    temperature: float = 0.3
    top_k: int = 250
    top_p: float = 0.95
    rep_penalty: float = 1.1
    max_tokens: int = 512

@app.get("/voices")
def list_voices():
    return {"voices": VOICE_NAMES}

@app.post("/synthesize")
def synthesize(req: SynthRequest):
    if not req.text.strip():
        return {"error": "Моля въведете текст."}
    if not VOICES:
        return {"error": "Няма гласове в папка voices/."}
    if req.voice not in VOICES:
        return {"error": "Непознат глас."}

    spk_emb = VOICES[req.voice]
    tok = TTSTokenizer()
    sentences = tok.split_sentences(req.text, max_chars=200)
    all_tokens = []

    for sent in sentences:
        print(f"Processing sentence: {sent[:50]}...")
        tokens = generate_sentence(sent, spk_emb, req.temperature, req.top_k,
                                   req.top_p, req.rep_penalty, req.max_tokens)
        if tokens is None or len(tokens) == 0:
            continue
        all_tokens.extend(tokens.tolist())

    if not all_tokens:
        return {"error": "Генерирането не върна токени."}

    # Връщаме JSON с токените
    return JSONResponse(content={
        "tokens": all_tokens,
        "speaker_emb": spk_emb.tolist(),
        "sample_rate": SAMPLE_RATE,
        "voice": req.voice
    })

@app.get("/")
async def root():
    # Опитваме се да прочетем index.html, ако съществува
    if os.path.exists("index.html"):
        with open("index.html", "r", encoding="utf-8") as f:
            return HTMLResponse(f.read())
    else:
        return HTMLResponse("""
        <!DOCTYPE html>
        <html>
        <head><title>AniTTS ONNX - Token Generator</title></head>
        <body>
            <h1>🎙️ AniTTS ONNX - Token Generator</h1>
            <p>✅ Сървърът работи!</p>
            <p>📡 Използвайте <code>POST /synthesize</code> за да получите JSON с токени.</p>
            <p>🎤 Vocoder сървър: <code>https://ssasio-ani-bg.hf.space/vocoder</code></p>
            <hr>
            <h3>📱 За Android приложението:</h3>
            <p>Използвайте URL: <code>https://ssasio-anitts-onix.hf.space/synthesize</code></p>
        </body>
        </html>
        """)

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