# ---------------------------------------------------------------------------
# FunctionGemma Mobile Actions — Hugging Face Space
# Runs litert-community/functiongemma-270m-ft-mobile-actions on-device (CPU)
# via the LiteRT-LM Python runtime. Gradio UI.
# Body matches the project's original code; header constants re-assembled.
# ---------------------------------------------------------------------------
import datetime as _dt
import html
import json
import os
import re
import threading

import gradio as gr

MODEL_REPO = "litert-community/functiongemma-270m-ft-mobile-actions"
MODEL_FILE = "mobile_actions_q8_ekv1024.litertlm"
MAX_TOKENS = 128

FUNCTIONS = {
    "send_email": {
        "description": "Sends an email to a recipient.",
        "params": {
            "to": "The email address of the recipient.",
            "subject": "The subject of the email.",
            "body": "The body of the email.",
        },
        "required": ["to", "subject", "body"],
    },
    "turn_on_flashlight": {
        "description": "Turns the flashlight on.",
        "params": {},
        "required": [],
    },
    "turn_off_flashlight": {
        "description": "Turns the flashlight off.",
        "params": {},
        "required": [],
    },
    "open_wifi_settings": {
        "description": "Opens the WiFi settings.",
        "params": {},
        "required": [],
    },
    "create_calendar_event": {
        "description": "Creates a new calendar event.",
        "params": {
            "title": "The title of the event.",
            "datetime": "The datetime of the event in ISO format.",
        },
        "required": ["title", "datetime"],
    },
    "show_map": {
        "description": "Shows a location on the map.",
        "params": {"query": "The location query to show on the map."},
        "required": ["query"],
    },
    "create_contact": {
        "description": "Creates a new contact.",
        "params": {
            "first_name": "The first name of the contact.",
            "last_name": "The last name of the contact.",
            "email": "The email of the contact.",
            "phone_number": "The phone number of the contact.",
        },
        "required": ["first_name", "last_name"],
    },
}
ORDER = list(FUNCTIONS)


def _schema_decl(name, info):
    if info["params"]:
        props = ",".join(
            "{k}:{{description:<escape>{v}<escape>,type:<escape>STRING<escape>}}".format(
                k=k, v=info["params"][k]
            )
            for k in info["params"]
        )
        required = ",".join("<escape>{0}<escape>".format(r) for r in info["required"])
        parameters = (
            "parameters:{{properties:{{{0}}},required:[{1}],"
            "type:<escape>OBJECT<escape>}}".format(props, required)
        )
    else:
        parameters = (
            "parameters:{properties:{},required:[],type:<escape>OBJECT<escape>}"
        )
    return (
        "<start_function_declaration>declaration:{name}"
        "{{description:<escape>{desc}<escape>,{params}}}"
        "<end_function_declaration>".format(name=name, desc=info["description"], params=parameters)
    )


# ---------------------------------------------------------------------------
# Prompt builder — exactly the FunctionGemma fine-tune template
# ---------------------------------------------------------------------------
def build_prompt(user_query, when=None):
    when = when or _dt.datetime.now()
    decls = "".join(_schema_decl(n, FUNCTIONS[n]) for n in ORDER)
    return (
        "<bos><start_of_turn>developer\n"
        "Current date and time given in YYYY-MM-DDTHH:MM:SS format: {ts}\n"
        "Day of week is {dow}\n"
        "You are a model that can do function calling with the following functions"
        "{decls}\n<end_of_turn>\n"
        "<start_of_turn>user\n{query}<end_of_turn>\n"
        "<start_of_turn>model\n"
    ).format(ts=when.strftime("%Y-%m-%dT%H:%M:%S"), dow=when.strftime("%A"),
             decls=decls, query=user_query)


# ---------------------------------------------------------------------------
# Output parser — <start_function_call>call:name{args}<end_function_call>
# ---------------------------------------------------------------------------
_CALL_RE = re.compile(r"<start_function_call>(.*?)<end_function_call>", re.S)
_NAME_RE = re.compile(r"call:\s*([a-zA-Z_]\w*)\s*\{(.*)\}", re.S)


def parse_function_call(text):
    """Return (function_name_or_None, args_dict) parsed from model output."""
    if not text:
        return None, {}
    m = _CALL_RE.search(text)
    if not m:
        return None, {}
    cm = _NAME_RE.match(m.group(1).strip())
    if not cm:
        return None, {}
    return cm.group(1), _parse_args(cm.group(2))


def _parse_args(raw):
    raw = (raw or "").strip()
    if not raw:
        return {}
    # FunctionGemma emits JSON-ish args where string values are wrapped in
    # <escape> markers and keys are bare. Normalize to real JSON:
    s = raw
    s = re.sub(r"<escape>\s*(\S[\s\S]*?\S?)\s*<escape>", r'"\1"', s)   # quote values
    s = re.sub(r'(?<!["\:\w])([A-Za-z_]\w*)\s*:', r'"\1":', s)           # quote keys
    s = re.sub(r"<escape>", "", s)                                       # stragglers
    try:
        v = json.loads(s)
        return v if isinstance(v, dict) else {}
    except Exception:
        out = {}
        for mm in re.finditer(
            r"([A-Za-z_]\w*)\s*:\s*([^,}]*?)\s*(?=,\s*[A-Za-z_]\w*\s*:|$)", raw, re.S
        ):
            out[mm.group(1)] = mm.group(2).replace("<escape>", "").strip()
        return out


# ---------------------------------------------------------------------------
# LiteRT-LM backend (verified against litert-lm 0.16.x package source)
# ---------------------------------------------------------------------------
_engine = {"fn": None, "backend": None, "error": None, "model_path": None}
_model_event = threading.Event()
_lock = threading.Lock()


def _download_model():
    token = os.getenv("HF_TOKEN") or None
    try:
        if token:
            from huggingface_hub import hf_hub_download
            _engine["model_path"] = hf_hub_download(
                repo_id=MODEL_REPO, filename=MODEL_FILE, token=token
            )
    except Exception as e:  # gated / offline / no token
        _engine["error"] = "model download failed: {0}".format(e)
    finally:
        _model_event.set()


def _generate_with_litert_lm(prompt):
    """Verified litert-lm 0.16.x API (from package source):
    Engine(model_path) -> create_session -> run_prefill -> run_decode."""
    from litert_lm import Engine
    from litert_lm.interfaces import Backend, SamplerConfig

    engine = Engine(model_path=_engine["model_path"], backend=Backend.CPU())
    try:
        session = engine.create_session(
            sampler_config=SamplerConfig(top_k=1, temperature=0.0),
            max_output_tokens=MAX_TOKENS,
        )
        session.run_prefill([prompt])
        resp = session.run_decode()
        return resp.texts[0] if resp.texts else ""
    finally:
        try:
            engine.close()
        except Exception:
            pass


def _init_backend():
    with _lock:
        if _engine["backend"] is not None or _engine["error"]:
            return
        try:
            # import check only; the engine is created per request (stateless)
            from litert_lm import Engine  # noqa: F401
            from litert_lm.interfaces import Backend, SamplerConfig  # noqa: F401
            _engine["backend"] = "litert-lm (CPU / XNNPACK)"
            _engine["fn"] = _generate_with_litert_lm
        except Exception as e:
            _engine["error"] = "litert-lm unavailable: {0}".format(e)


threading.Thread(target=_init_backend, daemon=True).start()
threading.Thread(target=_download_model, daemon=True).start()


# ---------------------------------------------------------------------------
# Rule-based intent mapper (fallback / prompt-format demo mode)
# ---------------------------------------------------------------------------
def demo_intent(q):
    ql = q.lower()
    if any(w in ql for w in ("flashlight", "flash", "torch")):
        off = any(w in ql for w in ("off", "stop", "disable", "turn off", "turn-off"))
        return ("turn_off_flashlight" if off else "turn_on_flashlight"), {}
    if "email" in ql or ("send" in ql and " to " in ql):
        args = {}
        m = re.search(r"([\w.+-]+@[\w.-]+)", q)
        if m:
            args["to"] = m.group(1)
        m = re.search(r"subject\s+(.+)", q, re.I)
        if m:
            args["subject"] = m.group(1).strip()[:60]
        return "send_email", args
    if any(w in ql for w in ("calendar", "event", "meeting", "appointment")):
        args = {"title": "Event"}
        m = re.search(r"(?:titled|called|named)\s+(.+)", q, re.I)
        if m:
            args["title"] = m.group(1).strip()[:60]
        return "create_calendar_event", args
    if "wifi" in ql:
        return "open_wifi_settings", {}
    if any(w in ql for w in ("map", "location", "where is")):
        m = re.search(r"show\s+(.+?)\s+on the map", q, re.I)
        args = {"query": m.group(1).strip()[:60] if m else q.strip()[:60]}
        return "show_map", args
    if "contact" in ql or ("save" in ql and "number" in ql):
        return "create_contact", {}
    return None, {}


# ---------------------------------------------------------------------------
# Phone-screen renderer
# ---------------------------------------------------------------------------
def _e(s):
    return html.escape(str(s or ""))


def render_phone(action=None, args=None, query=""):
    args = args or {}
    q = _e(query)

    def field(label, value):
        return '<div class="row"><span class="lbl">{0}</span><span class="val">{1}</span></div>'.format(
            _e(label), _e(value))

    if action == "send_email":
        body = """
        <div class="screen">
          <div class="appbar">✉️ Compose</div>
          <div class="body">{f_to}{f_sub}{f_body}
            <div class="btn">Send</div>
          </div>
        </div>""".format(
            f_to=field("To", args.get("to", "—")),
            f_sub=field("Subject", args.get("subject", "—")),
            f_body=field("Body", args.get("body", "—")),
        )
    elif action in ("turn_on_flashlight", "turn_off_flashlight"):
        on = action == "turn_on_flashlight"
        body = """
        <div class="screen">
          <div class="appbar">🔦 Flashlight</div>
          <div class="body center">
            <div class="flash {cls}">{icon}</div>
            <div class="big">{label}</div>
            <div class="btn">{btn}</div>
          </div>
        </div>""".format(
            cls="on" if on else "off",
            icon="💡" if on else "🌑",
            label="Flashlight is ON" if on else "Flashlight is OFF",
            btn="Turn on" if not on else "Turn off",
        )
    elif action == "open_wifi_settings":
        body = """
        <div class="screen">
          <div class="appbar">📶 WiFi settings</div>
          <div class="body">
            <div class="wifirow"><span class="wifi-ic">📶</span><span>HomeNetwork</span><span class="lbl">Connected</span></div>
            <div class="wifirow"><span class="wifi-ic sec">🔒</span><span>CoffeeShop_Guest</span></div>
            <div class="wifirow"><span class="wifi-ic sec">🔒</span><span>Office_5G</span></div>
            <div class="btn">Open settings</div>
          </div>
        </div>"""
    elif action == "create_calendar_event":
        body = """
        <div class="screen">
          <div class="appbar">📅 New event</div>
          <div class="body">{f_t}{f_d}
            <div class="btn">Save event</div>
          </div>
        </div>""".format(
            f_t=field("Title", args.get("title", "—")),
            f_d=field("When", args.get("datetime", "—")),
        )
    elif action == "show_map":
        body = """
        <div class="screen">
          <div class="appbar">🗺️ Maps</div>
          <div class="body map">
            <div class="mapgrid"></div>
            <div class="pin">📍</div>
            <div class="maplbl">{query}</div>
          </div>
        </div>""".format(query=_e(args.get("query", "")))
    elif action == "create_contact":
        body = """
        <div class="screen">
          <div class="appbar">👤 New contact</div>
          <div class="body center">
            <div class="avatar">👤</div>
            {f_fn}{f_ln}{f_em}{f_ph}
            <div class="btn">Save contact</div>
          </div>
        </div>""".format(
            f_fn=field("First", args.get("first_name", "—")),
            f_ln=field("Last", args.get("last_name", "—")),
            f_em=field("Email", args.get("email", "—")),
            f_ph=field("Phone", args.get("phone_number", "—")),
        )
    else:
        body = """
        <div class="screen">
          <div class="appbar">🤖 Function call</div>
          <div class="body"><div class="raw">{raw}</div></div>
        </div>""".format(raw=q or "No action detected")

    now = _dt.datetime.now().strftime("%H:%M")
    return """
    <div class="phone">
      <div class="notch"></div>
      <div class="status"><span>{time}</span><span>📶 🔋</span></div>
      <div class="hero"><div class="hero-title">Mobile Actions</div>
        <div class="hero-sub">functiongemma-270m-ft-mobile-actions</div></div>
      {body}
    </div>
    """.format(time=now, body=body)


# ---------------------------------------------------------------------------
# Main handler
# ---------------------------------------------------------------------------
def run(user_query):
    if not (user_query or "").strip():
        return (render_phone(), "**Status:** enter a request first.",
                "", "")
    user_query = user_query.strip()
    prompt = build_prompt(user_query)

    if not _model_event.is_set():
        _model_event.wait(timeout=180)  # first run: give the download a chance

    # Warm up the backend once
    if _engine["backend"] is None and not _engine["error"]:
        _init_backend()

    raw_out, parsed = None, None
    if _engine["model_path"] and _engine["fn"]:
        try:
            raw_out = _engine["fn"](prompt)
            action, args = parse_function_call(raw_out)
            parsed = action is not None
        except Exception as e:
            _engine["error"] = "inference error: {0}".format(e)
            raw_out = None

    if parsed:
        status = ("**Status:** ✅ Model running ({backend}) — emitted "
                  "`call:{name}`").format(backend=_engine["backend"], name=action)
        return render_phone(action, args, user_query), status, prompt, raw_out

    # Fallback: prompt-format demo (explicitly labelled, never silent)
    action, args = demo_intent(user_query)
    if _engine["error"]:
        err = _engine["error"].replace("\n", " ")
        note = "⚠️ Model unavailable ({0}). ".format(err)
    elif not _engine["model_path"]:
        note = ("⚠️ Model not downloaded — add the **HF_TOKEN** secret "
                "(Settings → Variables and secrets) to run real inference. ")
    else:
        note = "⚠️ Backend not ready. "
    status = (note + "Showing the **prompt-format demo** (exact FunctionGemma "
                     "prompt built; action inferred by rules).")
    return (render_phone(action, args, user_query), status, prompt, raw_out or "")


# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------
CSS = """
.phone{width:320px;margin:0 auto;background:#0f1220;border-radius:2.2rem;padding:14px 14px 22px;
  color:#e8eaf6;font-family:ui-sans-serif,system-ui,-apple-system,Segoe UI,Roboto,sans-serif;
  box-shadow:0 20px 50px rgba(0,0,0,.45);border:1px solid #2a2f4a}
.notch{width:90px;height:18px;background:#000;border-radius:12px;margin:4px auto 8px}
.status{display:flex;justify-content:space-between;font-size:11px;color:#9aa3c7;padding:0 6px}
.hero{background:linear-gradient(135deg,#2563eb,#7c3aed);border-radius:18px;padding:14px 16px;margin:10px 0 12px}
.hero-title{font-weight:800;font-size:17px;color:#fff}
.hero-sub{font-size:10px;color:#dbe3ff;margin-top:2px}
.screen{background:#171a2e;border-radius:16px;padding:6px 0 14px;overflow:hidden}
.appbar{background:#1f2440;padding:10px 14px;font-size:13px;font-weight:700;color:#c7d2fe}
.body{padding:12px 14px;font-size:13px}
.row{display:flex;padding:7px 0;border-bottom:1px solid #262c4d}
.lbl{color:#8b93bd;width:72px;flex:none}
.val{color:#e8eaf6;word-break:break-word}
.btn{margin-top:14px;background:#4f46e5;color:#fff;text-align:center;padding:9px;border-radius:10px;
  font-weight:700;font-size:13px}
.center{text-align:center}
.big{margin:6px 0 2px;font-weight:700}
.flash{font-size:44px}
.flash.on{filter:drop-shadow(0 0 14px #facc15)}
.flash.off{filter:grayscale(1);opacity:.55}
.map{position:relative;height:170px}
.mapgrid{position:absolute;inset:0;background:
  repeating-linear-gradient(0deg,transparent 0 28px,#243050 28px 29px),
  repeating-linear-gradient(90deg,transparent 0 28px,#243050 28px 29px),#1b2040}
.pin{position:absolute;top:42%;left:50%;transform:translate(-50%,-100%);font-size:30px}
.maplbl{position:absolute;bottom:10px;left:0;right:0;text-align:center;font-weight:700;color:#fff}
.avatar{font-size:38px;margin:4px auto 8px}
.wifirow{display:flex;gap:10px;align-items:center;padding:9px 0;border-bottom:1px solid #262c4d}
.wifi-ic.sec{font-size:12px}
.raw{background:#101327;border-radius:10px;padding:10px;font-family:ui-monospace,Menlo,Consolas,monospace;
  font-size:11px;color:#a5b4fc;word-break:break-all;white-space:pre-wrap}
"""

EXAMPLES = [
    "Turn on the flashlight",
    "Turn off the flashlight",
    "Open the WiFi settings",
    "Send an email to jane@example.com with the subject Project Update and body Please review the draft.",
    "Create a calendar event titled Standup tomorrow at 3pm",
    "Show the Eiffel Tower on the map",
    "Create a contact named John Smith with phone 555-1234",
]

with gr.Blocks(title="FunctionGemma Mobile Actions", css=CSS,
               theme=gr.themes.Soft(primary_hue="indigo")) as demo:
    gr.Markdown(
        "# 📱 FunctionGemma Mobile Actions\n"
        "A demo of **[litert-community/functiongemma-270m-ft-mobile-actions]"
        "(https://huggingface.co/litert-community/functiongemma-270m-ft-mobile-actions)** — "
        "Google's 270M function-calling Gemma compiled for on-device "
        "[LiteRT-LM](https://ai.google.dev/edge/litert-lm) inference. The model runs "
        "**inside this Space on CPU** (`litert-lm` Python runtime, "
        "`mobile_actions_q8_ekv1024.litertlm`, 1024-token context). "
        "Describe a phone action in natural language and watch it become a "
        "structured function call."
    )
    with gr.Row():
        with gr.Column(scale=3):
            inp = gr.Textbox(
                label="What do you want on your phone?",
                placeholder='e.g. "Send an email to jane@example.com with the subject Hello"',
                lines=2,
            )
            btn = gr.Button("Run agent", variant="primary")
            gr.Examples(examples=EXAMPLES, inputs=inp)
            status = gr.Markdown("**Status:** booting…")
        with gr.Column(scale=2):
            phone = gr.HTML(render_phone())
    with gr.Accordion("Raw FunctionGemma prompt", open=False):
        raw_prompt = gr.Textbox(label="Raw FunctionGemma prompt", lines=8,
                                interactive=False)
    with gr.Accordion("Raw model output", open=False):
        raw_out = gr.Textbox(label="Raw model output", lines=8,
                             interactive=False)

    outputs = [phone, status, raw_prompt, raw_out]
    btn.click(run, inputs=inp, outputs=outputs)
    inp.submit(run, inputs=inp, outputs=outputs)


if __name__ == "__main__":
    demo.queue(default_concurrency_limit=2).launch()