"""
N-ATLaS Text Gateway — Gradio Space (CPU / no GPU required).

Runs the REAL N-ATLaS model from a GGUF quantisation using llama-cpp-python on
CPU. No @spaces.GPU needed. Exposes a Gradio API named `chat` that the cPanel
PHP gateway calls.

To keep answers inside cPanel's ~120 s request limit on free CPU:
  - use a small quant (default Q3_K_M); set NATLAS_GGUF_FILE to change it
  - cap output length (NATLAS_MAX_TOKENS, default 200)
"""

import os

import gradio as gr
from huggingface_hub import hf_hub_download
from llama_cpp import Llama

# If the Space happens to run on ZeroGPU we must expose a @spaces.GPU function;
# on plain CPU the decorator is a no-op.
try:
    import spaces
    _gpu = spaces.GPU
except Exception:
    def _gpu(fn=None, **_kwargs):
        return fn if callable(fn) else (lambda f: f)

REPO = os.environ.get("NATLAS_GGUF_REPO", "QuantFactory/N-ATLaS-GGUF")
FILE = os.environ.get("NATLAS_GGUF_FILE", "N-ATLaS.Q3_K_M.gguf")   # smaller = faster on CPU
N_CTX = int(os.environ.get("NATLAS_CTX", "2048"))
N_THREADS = int(os.environ.get("NATLAS_THREADS", str(max(2, os.cpu_count() or 2))))
MAX_TOKENS = int(os.environ.get("NATLAS_MAX_TOKENS", "200"))

SYSTEM = (
    "You are N-ATLaS, Nigeria's premier sovereign multilingual AI foundation model "
    "developed by FMCIDE, NCAIR, and Awarri Technologies. You are fluent in Hausa, Yoruba, "
    "Igbo, Nigerian Pidgin, and Nigerian English, and you answer in the user's language. "
    "Be concise. If unsure of a medical, legal, or financial fact, say so and advise "
    "consulting a qualified professional. Attribution: N-ATLaS is an initiative of the "
    "Federal Ministry of Communications, Innovation and Digital Economy (FMCIDE), powered "
    "by Awarri Technologies."
)

print(f"[N-ATLaS] loading {REPO} / {FILE} ...", flush=True)
_model_path = hf_hub_download(REPO, FILE)
llm = Llama(model_path=_model_path, n_ctx=N_CTX, n_threads=N_THREADS,
            chat_format="llama-3", verbose=False)


@_gpu
def natlas_reply(message: str) -> str:
    message = (message or "").strip()
    if not message:
        return "Please type a message."
    try:
        out = llm.create_chat_completion(
            messages=[{"role": "system", "content": SYSTEM},
                      {"role": "user", "content": message}],
            temperature=0.7,
            max_tokens=MAX_TOKENS,
        )
        return out["choices"][0]["message"]["content"]
    except Exception as e:  # noqa: BLE001
        return f"[N-ATLaS error] {e}"


demo = gr.Interface(
    fn=natlas_reply,
    inputs=gr.Textbox(label="Message", lines=3, placeholder="Sannu / Ẹ káàbọ̀ / Ndeewo / How far"),
    outputs=gr.Textbox(label="N-ATLaS", lines=8),
    title="N-ATLaS Text Gateway (CPU)",
    description="Nigeria's sovereign multilingual AI (Hausa, Yoruba, Igbo, Pidgin, English). "
                "Powered by Awarri Technologies.",
    api_name="chat",
)

demo.launch()