import os
import re
import tempfile
from pathlib import Path

import gradio as gr
import numpy as np
import soundfile as sf
import torch
from huggingface_hub import hf_hub_download
from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM

# ---- EGTTS (Coqui XTTS) ----
from TTS.tts.configs.xtts_config import XttsConfig
from TTS.tts.models.xtts import Xtts

# =========================================================
# CONFIG
# =========================================================
# Translator (MSA <-> Egyptian)
TRANSLATOR_MODEL = os.getenv("TRANSLATOR_MODEL", "oddadmix/Masrawy-BiLingual-v1")

# LLM (CPU default: Qwen 1.5B is more realistic than 3B on CPU Spaces)
LLM_MODEL = os.getenv("LLM_MODEL", "Qwen/Qwen2.5-1.5B-Instruct")

# Translation length controls (NEW)
TRANSLATE_MAX_LEN = int(os.getenv("TRANSLATE_MAX_LEN", "256"))  # زوديها لو عايزة
TRANSLATE_MIN_LEN = int(os.getenv("TRANSLATE_MIN_LEN", "1"))

# EGTTS repo/files
EGTTS_REPO = os.getenv("EGTTS_REPO", "OmarSamir/EGTTS-V0.1")
EGTTS_CONFIG = os.getenv("EGTTS_CONFIG", "config.json")
EGTTS_VOCAB = os.getenv("EGTTS_VOCAB", "vocab.json")
EGTTS_MODEL = os.getenv("EGTTS_MODEL", "model.pth")
EGTTS_DEFAULT_SPEAKER = os.getenv("EGTTS_DEFAULT_SPEAKER", "speaker_reference.wav")

# Paths
BASE_DIR = Path(__file__).parent
ASSETS_DIR = BASE_DIR / "assets"
ASSETS_DIR.mkdir(parents=True, exist_ok=True)

# Device
USE_GPU = torch.cuda.is_available()
DEVICE_STR = "cuda" if USE_GPU else "cpu"
PIPE_DEVICE = 0 if USE_GPU else -1

# =========================================================
# Load models (once)
# =========================================================
translator = pipeline("translation", model=TRANSLATOR_MODEL, device=PIPE_DEVICE)

tokenizer = AutoTokenizer.from_pretrained(LLM_MODEL, trust_remote_code=True)
llm = AutoModelForCausalLM.from_pretrained(
    LLM_MODEL,
    torch_dtype="auto",
    device_map="auto" if USE_GPU else None,
    trust_remote_code=True,
)
if not USE_GPU:
    llm = llm.to("cpu")
llm.eval()

_tts = None
_default_ref_path = str(ASSETS_DIR / EGTTS_DEFAULT_SPEAKER)

# =========================================================
# Helpers
# =========================================================
def _wrap_msa_prompt(msa: str) -> str:
    """Optional light normalization (kept minimal)."""
    msa = (msa or "").strip()
    msa = re.sub(r"\s+", " ", msa).strip()
    return msa

def _safe_int(x, default=0):
    try:
        return int(x)
    except Exception:
        return default

# =========================================================
# Translator helpers (explicit direction, non-ambiguous)
# NOTE: We explicitly set max_length/min_length to avoid truncation.
# =========================================================
def to_msa(text: str) -> str:
    text = (text or "").strip()
    if not text:
        return ""
    out = translator(
        text + " <ar>",
        max_length=TRANSLATE_MAX_LEN,
        min_length=TRANSLATE_MIN_LEN,
        truncation=True,
    )
    return out[0]["translation_text"]

def to_egyptian(text: str) -> str:
    text = (text or "").strip()
    if not text:
        return ""
    out = translator(
        text + " <arz>",
        max_length=TRANSLATE_MAX_LEN,
        min_length=TRANSLATE_MIN_LEN,
        truncation=True,
    )
    return out[0]["translation_text"]

# =========================================================
# Output cleaning (reduce meta + repetition)
# =========================================================
_BANNED_PHRASES = [
    "كمساعد", "كمساعد ذكي", "اعتذر", "آسف", "مش عارف", "لا أستطيع", "غير قادر",
    "لا يمكنني", "لا أقدر", "لا أملك معلومات", "قد لا يكون", "ربما", "عادةً", "بشكل عام",
]

def clean_egyptian(text: str) -> str:
    t = (text or "").strip()

    for p in _BANNED_PHRASES:
        t = t.replace(p, "")

    t = re.sub(r"\s+", " ", t).strip()

    # prevent weird repetition loops (soft)
    t = re.sub(r"(\b\w+\b)(?:\s+\1){3,}", r"\1", t)

    # normalize punctuation
    t = re.sub(r"[.،]{3,}", "…", t).strip()

    if not t:
        t = "تمام — قولي محتاج إيه بالظبط؟"

    return t

# =========================================================
# Qwen generation (MSA for stability)
# Fix truncation/early-stop:
# - don't pass eos_token_id
# - soften repetition constraints
# - optional: add min_new_tokens
# =========================================================
def qwen_generate_msa(msa_prompt: str, max_new_tokens: int, temperature: float, top_p: float) -> str:
    msa_prompt = _wrap_msa_prompt(msa_prompt)
    if not msa_prompt:
        return ""

    system_msg = (
        "أنت مساعد تنفيذي عملي ومباشر. "
        "إذا طُلبت طريقة/شرح: اكتب خطوات كاملة من البداية للنهاية. "
        "بدون اعتذار وبدون كلام ميتا. "
        "اكتب بالعربية الفصحى البسيطة فقط."
    )

    messages = [
        {"role": "system", "content": system_msg},
        {"role": "user", "content": msa_prompt},
    ]

    input_ids = tokenizer.apply_chat_template(
        messages, add_generation_prompt=True, return_tensors="pt"
    )
    input_ids = input_ids.to(llm.device) if USE_GPU else input_ids.to("cpu")

    pad_id = tokenizer.pad_token_id or tokenizer.eos_token_id

    with torch.no_grad():
        output_ids = llm.generate(
            input_ids,
            max_new_tokens=int(max_new_tokens),
            min_new_tokens=32,  # help avoid tiny answers

            do_sample=True,
            temperature=float(temperature),
            top_p=float(top_p),

            repetition_penalty=1.06,
            no_repeat_ngram_size=3,

            pad_token_id=pad_id,
            # IMPORTANT: do NOT set eos_token_id -> prevents early stop
        )

    gen_ids = output_ids[0][input_ids.shape[-1]:]
    text = tokenizer.decode(gen_ids, skip_special_tokens=True).strip()
    return text

# =========================================================
# EGTTS assets + loader
# =========================================================
def _ensure_egtts_assets():
    cfg_path = hf_hub_download(EGTTS_REPO, EGTTS_CONFIG)
    vocab_path = hf_hub_download(EGTTS_REPO, EGTTS_VOCAB)
    model_path = hf_hub_download(EGTTS_REPO, EGTTS_MODEL)
    spk_path = hf_hub_download(EGTTS_REPO, EGTTS_DEFAULT_SPEAKER)

    for src in [cfg_path, vocab_path, model_path, spk_path]:
        dst = ASSETS_DIR / Path(src).name
        if not dst.exists():
            Path(dst).write_bytes(Path(src).read_bytes())

def load_tts():
    global _tts
    if _tts is not None:
        return _tts

    _ensure_egtts_assets()

    cfg = XttsConfig()
    cfg.load_json(str(ASSETS_DIR / EGTTS_CONFIG))

    model = Xtts.init_from_config(cfg)
    model.load_checkpoint(
        cfg,
        checkpoint_path=str(ASSETS_DIR / EGTTS_MODEL),
        vocab_path=str(ASSETS_DIR / EGTTS_VOCAB),
        use_deepspeed=False,
        eval=True,
    )
    model.to(DEVICE_STR)
    model.eval()

    _tts = model
    return _tts

# =========================================================
# Speak (TTS) utility: speaks FINAL LLM Egyptian response
# =========================================================
def _split_text(text: str, max_chars: int = 220):
    text = (text or "").strip()
    if not text:
        return []
    parts = re.split(r"([.!؟\n،])", text)
    chunks, cur = [], ""
    for i in range(0, len(parts), 2):
        seg = parts[i].strip()
        punc = parts[i + 1] if i + 1 < len(parts) else ""
        piece = (seg + punc).strip()
        if not piece:
            continue
        if len(cur) + len(piece) + 1 <= max_chars:
            cur = (cur + " " + piece).strip()
        else:
            if cur:
                chunks.append(cur)
            cur = piece
    if cur:
        chunks.append(cur)
    return chunks

def egtss_speak(text_egy: str, ref_path: str | None, tts_temp: float):
    tts = load_tts()
    ref = ref_path if (ref_path and Path(ref_path).exists()) else _default_ref_path

    gpt_cond_latent, speaker_embedding = tts.get_conditioning_latents(audio_path=[ref])

    chunks = _split_text(text_egy, max_chars=220)
    if not chunks:
        chunks = ["تمام"]

    wavs = []
    for ch in chunks:
        out = tts.inference(
            text=ch,
            language="ar",
            gpt_cond_latent=gpt_cond_latent,
            speaker_embedding=speaker_embedding,
            temperature=float(tts_temp),
        )
        wavs.append(np.array(out["wav"], dtype=np.float32))

    audio = np.concatenate(wavs) if len(wavs) > 1 else wavs[0]
    return audio, 24000

# =========================================================
# Core pipeline: Input -> (optional to MSA) -> LLM -> to Egyptian -> clean -> TTS
# =========================================================
def run_full_pipeline(
    user_text: str,
    use_translator: bool,
    max_new_tokens: int,
    llm_temp: float,
    top_p: float,
    speaker_ref_path: str | None,
    tts_temp: float,
    show_debug: bool,
):
    user_text = (user_text or "").strip()
    if not user_text:
        return "", None, "", ""

    msa_in = to_msa(user_text) if use_translator else user_text
    llm_msa = qwen_generate_msa(msa_in, max_new_tokens, llm_temp, top_p)

    # Always enforce Egyptian output via translator to avoid mixed style
    final_egy = clean_egyptian(to_egyptian(llm_msa))

    audio_np, sr = egtss_speak(final_egy, speaker_ref_path, tts_temp)

    out_wav = tempfile.NamedTemporaryFile(delete=False, suffix=".wav")
    out_wav.close()
    sf.write(out_wav.name, audio_np, sr)

    if show_debug:
        return final_egy, out_wav.name, msa_in, llm_msa
    return final_egy, out_wav.name, "", ""

# =========================================================
# Gradio UI
# =========================================================
with gr.Blocks(title="Egyptian Arabic Assistant — Qwen (LLM) → EGTTS (TTS)") as demo:
    gr.Markdown(
        "## Egyptian Arabic Assistant — Qwen (LLM) → EGTTS (TTS)\n"
        "الـ TTS بينطق **رد الـ LLM** (مش input).\n\n"
        f"Translation max_length={TRANSLATE_MAX_LEN}, min_length={TRANSLATE_MIN_LEN}"
    )

    with gr.Row():
        with gr.Column(scale=1):
            user_in = gr.Textbox(lines=6, label="User message (Egyptian / Arabic)", placeholder="اكتب رسالتك هنا…")
            use_translator = gr.Checkbox(value=True, label="Use translator (Egyptian↔MSA)")

            max_new_tokens = gr.Slider(64, 768, value=384, step=16, label="Max new tokens")
            llm_temp = gr.Slider(0.1, 1.2, value=0.7, step=0.05, label="LLM Temperature")
            top_p = gr.Slider(0.5, 1.0, value=0.9, step=0.05, label="Top-p")

            speaker_ref = gr.Audio(type="filepath", label="Speaker reference (optional)")
            tts_temp = gr.Slider(0.1, 1.2, value=0.75, step=0.05, label="TTS Temperature")

            show_debug = gr.Checkbox(value=False, label="Show debug outputs")
            btn = gr.Button("Generate")

        with gr.Column(scale=1):
            out_text = gr.Textbox(lines=10, label="LLM Final Response (Egyptian)")
            out_audio = gr.Audio(label="TTS Output (speaks LLM response)", type="filepath")

            dbg_msa = gr.Textbox(lines=3, label="(Debug) Input after to_msa")
            dbg_llm = gr.Textbox(lines=8, label="(Debug) LLM output (before to_egy)")

    btn.click(
        fn=run_full_pipeline,
        inputs=[user_in, use_translator, max_new_tokens, llm_temp, top_p, speaker_ref, tts_temp, show_debug],
        outputs=[out_text, out_audio, dbg_msa, dbg_llm],
    )

demo.launch(server_name="0.0.0.0", server_port=7860)