"""
Diba chat (ZeroGPU). Serves the shared Diba chat UI (diba-ui.js / diba-ui.css) on top of a minimal Gradio app,
and exposes a streaming endpoint `/chat_stream` that the UI calls through the Gradio JS client.
- Replies in the language of the user's latest message; clear, professional tone; thinking always off.
"""
import json
import os
import re
import threading
from pathlib import Path

import gradio as gr
import spaces
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, StoppingCriteria, StoppingCriteriaList, TextIteratorStreamer

TOKEN = os.environ.get("HF_TOKEN")
HERE = Path(__file__).parent

# Model registry. "diba-base" is always loaded; "diba-mini" is optional and only offered when its repo exists and loads.
DEFAULT_MODEL = "diba-base"
MODELS = {
    "diba-base": {"repo": os.environ.get("DIBA_MODEL", "Dibachain/Diba-Base"), "label": "Diba-Base"},
    "diba-mini": {"repo": os.environ.get("DIBA_MINI_MODEL", "Dibachain/Diba-mini"), "label": "Diba-mini"},
}
LOADED = {}   # model id -> (tokenizer, model)


def load_model(repo):
    tk = AutoTokenizer.from_pretrained(repo, token=TOKEN, trust_remote_code=True)
    md = AutoModelForCausalLM.from_pretrained(repo, token=TOKEN, dtype=torch.bfloat16, trust_remote_code=True)
    md.to("cuda")
    md.eval()
    return tk, md


def repo_exists(repo):
    try:
        from huggingface_hub import HfApi
        HfApi(token=TOKEN).model_info(repo)
        return True
    except Exception as e:
        print(f"[diba] model repo {repo} not available ({type(e).__name__}); option hidden", flush=True)
        return False


LOADED[DEFAULT_MODEL] = load_model(MODELS[DEFAULT_MODEL]["repo"])
for _mid, _spec in MODELS.items():
    if _mid == DEFAULT_MODEL:
        continue
    if not repo_exists(_spec["repo"]):
        continue
    try:
        LOADED[_mid] = load_model(_spec["repo"])
        print(f"[diba] loaded optional model {_mid} from {_spec['repo']}", flush=True)
    except Exception as e:   # never let an optional model break the Space
        import traceback
        traceback.print_exc()
        print(f"[diba] failed to load optional model {_mid}: {type(e).__name__}; option hidden", flush=True)

AVAILABLE_MODELS = [{"id": mid, "label": MODELS[mid]["label"], "default": mid == DEFAULT_MODEL} for mid in MODELS if mid in LOADED]


def resolve_model(model_id):
    """Pick the requested model's (tokenizer, model); unknown or unavailable ids fall back to Diba-Base."""
    mid = str(model_id or "").strip().lower()
    return LOADED.get(mid) or LOADED[DEFAULT_MODEL]

SYSTEM = {
    "fa": ("تو «دیبا» هستی، دستیار هوش مصنوعی ساخته‌ی شرکت دیباچین. کاربر به فارسی نوشته؛ فقط به فارسی جواب بده، به زبان نوشتاری و مؤدبانه. "
           "دقیقاً به همان چیزی که کاربر خواسته جواب بده، روشن و بدون مقدمه‌ی اضافه. اگر کد لازم است، کد تمیز و کامل در بلوک کد بده. "
           "هرگز نام مدل پایه یا سازنده‌ی آن را نگو؛ اگر پرسیدند فقط بگو روی یک مدل زبانی متن‌باز ساخته شده‌ای و دیباچین تو را آموزش داده است."),
    "en": ("You are Diba, an AI assistant built by the company Dibachain. The user wrote in English, so reply ONLY in English. "
           "Be clear, accurate and professional. Answer exactly what was asked, without unnecessary preamble. When code is needed, "
           "give clean, complete code in a fenced code block. Never name or hint at any base model or its maker; if asked, say only "
           "that you were built on an open-source language model and trained by Dibachain."),
    "other": ("You are Diba, an AI assistant built by the company Dibachain. Reply in the same language the user wrote in. "
              "Be clear, accurate and professional. Never name any base model or its maker."),
}


def detect_lang(text: str) -> str:
    text = re.sub(r"```.*?```|`[^`]*`", " ", text, flags=re.S)
    fa = len(re.findall(r"[؀-ۿ]", text))
    en = len(re.findall(r"[A-Za-z]", text))
    if not fa and not en:
        return "other"
    return "fa" if fa >= en else "en"


class StopOnEvent(StoppingCriteria):
    def __init__(self, event):
        self.event = event

    def __call__(self, input_ids, scores, **kwargs):
        return self.event.is_set()


def gpu_seconds(messages, temperature, max_new_tokens, model_id=DEFAULT_MODEL):
    """Reserve only the GPU time a reply needs: ZeroGPU rejects a call whose reserved time exceeds the visitor's remaining quota."""
    return int(min(75, 15 + 0.045 * int(max_new_tokens or 768)))


@spaces.GPU(duration=gpu_seconds)
def generate(messages, temperature, max_new_tokens, model_id=DEFAULT_MODEL):
    tk, md = resolve_model(model_id)
    inputs = tk.apply_chat_template(messages, add_generation_prompt=True, enable_thinking=False,
                                    return_tensors="pt", return_dict=True).to("cuda")
    streamer = TextIteratorStreamer(tk, skip_prompt=True, skip_special_tokens=True, timeout=120)
    stop = threading.Event()
    kwargs = dict(**inputs, streamer=streamer, max_new_tokens=int(max_new_tokens),
                  do_sample=float(temperature) > 0, temperature=max(float(temperature), 1e-5), top_p=0.9,
                  repetition_penalty=1.05, stopping_criteria=StoppingCriteriaList([StopOnEvent(stop)]))
    worker = threading.Thread(target=md.generate, kwargs=kwargs)
    worker.start()
    out = ""
    try:
        for piece in streamer:
            out += piece
            yield out
    finally:
        stop.set()
        worker.join(timeout=30)


def build_messages(history):
    plain = [{"role": m["role"], "content": m["content"]} for m in history
             if isinstance(m, dict) and m.get("role") in ("user", "assistant") and isinstance(m.get("content"), str) and m["content"]]
    last_user = next((m["content"] for m in reversed(plain) if m["role"] == "user"), "")
    return [{"role": "system", "content": SYSTEM[detect_lang(last_user)]}] + plain[-16:]


def chat_stream(messages_json, temperature, max_new_tokens, model=DEFAULT_MODEL):
    try:
        history = json.loads(messages_json or "[]")
    except json.JSONDecodeError:
        history = []
    msgs = build_messages(history)
    if len(msgs) < 2:
        yield ""
        return
    try:
        for text in generate(msgs, temperature if temperature is not None else 0, max_new_tokens or 1024, model):
            yield text
    except gr.Error:
        raise
    except Exception as e:  # surface the real cause (ZeroGPU quota, missing token, CUDA) instead of an empty error event
        import traceback
        traceback.print_exc()
        raise gr.Error(f"{type(e).__name__}: {str(e)[:300]}")


def ask(message, temperature=0.0, max_new_tokens=768, model=DEFAULT_MODEL):
    final = ""
    for text in generate(build_messages([{"role": "user", "content": message}]), temperature, max_new_tokens, model):
        final = text
    return final


def models():
    """List the models this Space can serve (same JSON the UI receives as window.DIBA_MODELS)."""
    return json.dumps(AVAILABLE_MODELS, ensure_ascii=False)


def whoami(profile: "gr.OAuthProfile | None" = None):
    """Report the signed-in Hugging Face username (or empty) so the UI can offer sign in / sign out. Signed-in visitors run on their own ZeroGPU quota."""
    return profile.username if profile else ""


UI_CSS = (HERE / "diba-ui.css").read_text(encoding="utf-8")
UI_JS = (HERE / "diba-ui.js").read_text(encoding="utf-8")
gr.set_static_paths(paths=[str(HERE / "vendor")])   # self-hosted fonts and libraries, served at /gradio_api/file=...
ASSET_BASE = "/gradio_api/file=" + str(HERE).replace("\\", "/").rstrip("/") + "/"
MODELS_JSON = json.dumps(AVAILABLE_MODELS, ensure_ascii=False).replace("</", "<\\/")
HEAD = ('<meta name="viewport" content="width=device-width, initial-scale=1, viewport-fit=cover">'
        f"<style>{UI_CSS}</style><script>window.DIBA_BACKEND = 'gradio'; window.DIBA_ASSET_BASE = '{ASSET_BASE}'; "
        f"window.DIBA_FALLBACK_API = 'https://dibaai-diba-chat.hf.space'; window.DIBA_MODELS = {MODELS_JSON};</script><script>{UI_JS}</script>")

with gr.Blocks(title="Diba · دیبا", head=HEAD, css="footer{display:none!important}") as demo:
    gr.HTML('<div style="padding:40px;text-align:center;font-family:Tahoma">در حال بارگذاری دیبا… · Loading Diba…</div>')
    gr.LoginButton(visible=False)   # registers /login/huggingface and /logout; signed-in users get their own ZeroGPU quota
    with gr.Row(visible=False):
        messages_in = gr.Textbox()
        temperature_in = gr.Number(value=0)
        max_in = gr.Number(value=1024)
        model_in = gr.Textbox(value=DEFAULT_MODEL)   # model id from the registry ("diba-base" / "diba-mini"); unknown ids fall back to Base
        reply_out = gr.Textbox()
        who_out = gr.Textbox()
        who_btn = gr.Button()
        msg_in = gr.Textbox()
        ask_out = gr.Textbox()
        models_out = gr.Textbox()
        models_btn = gr.Button()
        stream_btn = gr.Button()
        ask_btn = gr.Button()
    stream_btn.click(chat_stream, [messages_in, temperature_in, max_in, model_in], reply_out, api_name="chat_stream", concurrency_limit=4)
    ask_btn.click(ask, [msg_in, temperature_in, max_in, model_in], ask_out, api_name="ask", concurrency_limit=4)
    who_btn.click(whoami, None, who_out, api_name="whoami")
    models_btn.click(models, None, models_out, api_name="models")

demo.queue(default_concurrency_limit=4)

if __name__ == "__main__":
    # show_error: the real reason (e.g. GPU quota) reaches the UI instead of an empty error event; SSR is not needed for this custom UI
    demo.launch(show_error=True, ssr_mode=False, favicon_path=str(HERE / "vendor" / "logo.png"))