"""
LFM2.5-1.2B-Thinking tool-calling demo, packaged as a Hugging Face ZeroGPU Space.

This is a rewrite of `main.py` for HF Spaces. The original talked to a local
Ollama server (`http://ubuntu.local:11434/v1`) and let Ollama parse the
OpenAI-style `tools` field for it. On a HF Space there is no Ollama, so we load
the model in-process with `transformers` on GPU and do the tool-call parsing
ourselves.

LFM2.5's native tool format is *Pythonic*: the model emits

    <|tool_call_start|>[web_search(query="liquid ai lfm")]<|tool_call_end|>

i.e. a Python list of function calls wrapped in special tokens. We parse that
with the `ast` module, execute the matching mock tool, feed the JSON result back
as a `tool`-role message, and let the model produce a final answer.

Reference: https://docs.liquid.ai/lfm/key-concepts/tool-use
           https://huggingface.co/LiquidAI/LFM2.5-1.2B-Thinking
"""

# `import spaces` MUST precede anything that touches CUDA (torch) so the
# ZeroGPU patch can apply. On HF Spaces it provides @spaces.GPU; locally the
# shim below makes it a no-op so the file still imports outside a Space.
try:
    import spaces
except ImportError:
    spaces = None

import ast
import inspect
import json
import re
import threading

import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer


def _gpu(duration: int = 180):
    """@spaces.GPU with a local no-op fallback for non-Space environments."""
    if spaces is not None:
        return spaces.GPU(duration=duration)

    def decorator(fn):
        return fn

    return decorator

MODEL_ID = "LiquidAI/LFM2.5-1.2B-Thinking"
MAX_NEW_TOKENS = 4096
MAX_ITERATIONS = 5

# Native LFM2.5 tool-call delimiters.
TOOL_CALL_START = "<|tool_call_start|>"
TOOL_CALL_END = "<|tool_call_end|>"

# LFM2.5-Thinking wraps its internal reasoning in <think>...</think> tags (the
# chat template splits on </think>). Hide it from the UI and show a placeholder
THINK_OPEN = "<think>"
THINK_CLOSE = "</think>"
THINK_PLACEHOLDER = "_🤔 thinking…_"

# ----------------------------------------------------------------------------
# Tools (mocked, same behaviour as main.py)
# ----------------------------------------------------------------------------


def web_search(query: str) -> list[dict]:
    """Mock web search; returns canned results regardless of the query."""
    return [
        {
            "title": "Top result for: " + query,
            "url": "https://example.com/search?q=" + query.replace(" ", "+"),
            "snippet": f"A plausible-looking excerpt relevant to '{query}'.",
        },
        {
            "title": "Secondary result for: " + query,
            "url": "https://example.org/search?q=" + query.replace(" ", "+"),
            "snippet": f"Another excerpt that touches on '{query}' from a different angle.",
        },
    ]


def send_email(to: str, subject: str, body: str) -> dict:
    """Mock email sender; in real life this would talk to an SMTP server."""
    print(f"\n--- drafting email ---\nTo: {to}\nSubject: {subject}\n{body}\n--- end ---")
    return {
        "status": "sent",
        "to": to,
        "subject": subject,
        "message_id": "mock-0001",
    }


# Registry of tools the model can call. Maps name -> callable.
TOOL_FUNCTIONS = {
    "web_search": web_search,
    "send_email": send_email,
}


def build_tools() -> list[dict]:
    """Tool schema in LFM2.5's native (flat) format — what the model was
    trained on. Dropped the OpenAI `{"type":"function","function":{...}}`
    wrapper that `main.py` used for Ollama."""
    return [
        {
            "name": "web_search",
            "description": "Search the web for up-to-date information on a topic",
            "parameters": {
                "type": "object",
                "properties": {
                    "query": {
                        "type": "string",
                        "description": "The search query",
                    }
                },
                "required": ["query"],
            },
        },
        {
            "name": "send_email",
            "description": "Send an email to a recipient",
            "parameters": {
                "type": "object",
                "properties": {
                    "to": {"type": "string", "description": "Recipient email address"},
                    "subject": {"type": "string", "description": "Email subject line"},
                    "body": {"type": "string", "description": "Email body content"},
                },
                "required": ["to", "subject", "body"],
            },
        },
    ]


def system_prompt() -> str:
    return (
        "You have tools available. Always use web_search to fetch facts; never answer "
        "from memory. When the user asks you to, use send_email to send emails.\n"
        "To call a tool, emit a Python list of function calls between the special "
        "tokens, e.g. "
        f"{TOOL_CALL_START}[web_search(query='liquid ai lfm')]{TOOL_CALL_END}.\n"
        "After receiving tool results, summarize them and answer the user.\n\n"
        f"List of tools: {json.dumps(build_tools())}"
    )


# ----------------------------------------------------------------------------
# Tool-call parsing (LFM2.5 emits Pythonic calls)
# ----------------------------------------------------------------------------


def parse_tool_calls(text: str) -> list[dict]:
    """Extract tool calls from raw model output.

    The model writes `<|tool_call_start|>[fn(a='1', b='2')]<|tool_call_end|>`,
    possibly with several calls in one list. We parse the list with `ast`
    (not `ast.literal_eval`, since a bare function call isn't a literal) and
    walk the AST for each call's name + keyword arguments.
    """
    calls = []
    pattern = re.escape(TOOL_CALL_START) + r"(.*?)" + re.escape(TOOL_CALL_END)
    for match in re.finditer(pattern, text, re.DOTALL):
        body = match.group(1).strip()
        try:
            tree = ast.parse(body, mode="eval").body
        except SyntaxError:
            continue
        if isinstance(tree, ast.Call):
            call_nodes = [tree]
        elif isinstance(tree, ast.List):
            call_nodes = [e for e in tree.elts if isinstance(e, ast.Call)]
        else:
            continue

        for node in call_nodes:
            if not isinstance(node.func, ast.Name):
                continue
            name = node.func.id
            arguments: dict = {}
            # Keyword arguments, e.g. query="..."
            for kw in node.keywords:
                if kw.arg is None:
                    continue
                try:
                    arguments[kw.arg] = ast.literal_eval(kw.value)
                except (ValueError, SyntaxError):
                    arguments[kw.arg] = ast.unparse(kw.value)
            # Positional arguments -> map onto parameter names by signature.
            fn = TOOL_FUNCTIONS.get(name)
            if fn is not None:
                params = list(inspect.signature(fn).parameters)
                for i, arg in enumerate(node.args):
                    if i < len(params):
                        try:
                            arguments[params[i]] = ast.literal_eval(arg)
                        except (ValueError, SyntaxError):
                            arguments[params[i]] = ast.unparse(arg)
            calls.append({"name": name, "arguments": arguments})
    return calls


def execute_tool(name: str, arguments: dict) -> str:
    """Run one parsed tool call, leniently — mirrors main.py's handling.

    Small models hallucinate extra params; we drop anything the function
    doesn't accept, and ask for a retry if a required param ends up missing.
    """
    fn = TOOL_FUNCTIONS.get(name)
    if fn is None:
        return json.dumps({"error": f"Unknown tool: {name}"})

    accepted = set(inspect.signature(fn).parameters)
    valid = {k: v for k, v in arguments.items() if k in accepted}
    dropped = sorted(set(arguments) - accepted)
    if dropped:
        print(f"  (dropping hallucinated params: {dropped})")

    required = {
        p
        for p, param in inspect.signature(fn).parameters.items()
        if param.default is inspect.Parameter.empty
    }
    missing = sorted(required - set(valid))
    if missing:
        return json.dumps(
            {"error": f"Missing required parameter(s) {missing} for tool '{name}'"}
        )

    result = fn(**valid)
    return json.dumps(result)


# ----------------------------------------------------------------------------
# Model loading (CPU, in-process)
# ----------------------------------------------------------------------------

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Loading {MODEL_ID} on {DEVICE} (bfloat16)…")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(MODEL_ID, dtype=torch.bfloat16).to(DEVICE)
model.eval()
print("Model ready.")


def generate_stream(messages: list[dict]):
    """Run one generation turn, streaming tokens. Returns (streamer, thread)."""
    # Render the chat template to a string, then tokenize explicitly. We do
    # NOT use apply_chat_template(tokenize=True, return_tensors="pt"): with
    # tokenize=True it returns the tokenizer.__call__ result, a BatchEncoding
    # (dict-like), not a plain tensor. model.generate then does
    # inputs_tensor.shape[0] on that BatchEncoding and raises AttributeError
    # (its __getattr__ falls through to self.data['shape']). Tokenizing the
    # rendered string ourselves gives a real tensor we control.
    text = tokenizer.apply_chat_template(
        messages, add_generation_prompt=True, tokenize=False
    )
    input_ids = tokenizer([text], return_tensors="pt").input_ids.to(model.device)
    streamer = TextIteratorStreamer(
        tokenizer, skip_special_tokens=False, skip_prompt=True
    )
    thread = threading.Thread(
        target=model.generate,
        args=(input_ids,),
        kwargs={
            "do_sample": True,
            "temperature": 0.05,
            "top_k": 50,
            "repetition_penalty": 1.05,
            "max_new_tokens": MAX_NEW_TOKENS,
            "streamer": streamer,
        },
    )
    thread.start()
    return streamer, thread


# ----------------------------------------------------------------------------
# Display helpers
# ----------------------------------------------------------------------------


def render_display(raw: str) -> str:
    """Turn raw model output (with special tokens) into readable markdown.

    LFM2.5-Thinking emits its chain-of-thought wrapped in the THINK_OPEN and
    THINK_CLOSE tags. Hide it from the UI and just show 'thinking…' while it
    is not finished
    """
    s = raw.replace("<|im_start|>", "").replace("<|im_end|>", "")

    if THINK_OPEN in s:
        pre, _, rest = s.partition(THINK_OPEN)
        if THINK_CLOSE in rest:
            _, _, post = rest.partition(THINK_CLOSE)
            s = f"{pre.strip()}\n\n{THINK_PLACEHOLDER}\n\n{post}"
        else:
            # Still reasoning — never leak the partial thinking text.
            s = (f"{pre.strip()}\n\n" if pre.strip() else "") + THINK_PLACEHOLDER

    s = s.replace(TOOL_CALL_START, "\n\n🔧 **Tool call:**\n```python\n")
    s = s.replace(TOOL_CALL_END, "\n```\n")
    return s.strip()


# ----------------------------------------------------------------------------
# Gradio app
# ----------------------------------------------------------------------------

EXAMPLES = [
    "Find the latest news about LiquidAI's LFM models, then email a short "
    "summary with the URLs to alice@example.com.",
    "Search the web for what C. elegans is and explain it.",
]


@_gpu(duration=120)
def respond(user_msg: str, history: list[dict]):
    """Generator driving the tool-calling loop, streaming into the chatbot.

    Decorated with @spaces.GPU so ZeroGPU attaches a GPU for the entire
    multi-turn loop, including the streamed tokens.
    """
    messages = [{"role": "system", "content": system_prompt()}] + list(history)
    messages.append({"role": "user", "content": user_msg})

    chatbot = [{"role": "user", "content": user_msg}]
    yield chatbot, messages[1:], ""

    for turn in range(1, MAX_ITERATIONS + 1):
        chatbot.append({"role": "assistant", "content": ""})  # streaming placeholder
        streamer, thread = generate_stream(messages)
        raw = ""
        for chunk in streamer:
            raw += chunk
            chatbot[-1] = {"role": "assistant", "content": render_display(raw)}
            yield chatbot, messages[1:], ""
        thread.join()

        # Keep the raw assistant turn (special tokens intact) for the next
        # round — the LFM2.5 1.2B chat template drops a structured `tool_calls`
        # field on re-render (known bug), so we must store the literal text.
        messages.append({"role": "assistant", "content": raw.replace("<|im_end|>", "").rstrip()})

        tool_calls = parse_tool_calls(raw)
        if not tool_calls:
            # No tool call => final answer; show it cleaned up and stop.
            chatbot[-1] = {"role": "assistant", "content": render_display(raw)}
            yield chatbot, messages[1:], ""
            return

        # Execute every requested tool and feed results back as tool messages.
        for call in tool_calls:
            result = execute_tool(call["name"], call["arguments"])
            messages.append({"role": "tool", "content": result})
            chatbot.append(
                {
                    "role": "assistant",
                    "content": f"🔧 **{call['name']}** result:\n```json\n{result}\n```",
                }
            )
            yield chatbot, messages[1:], ""

    chatbot.append(
        {
            "role": "assistant",
            "content": f"_Reached the {MAX_ITERATIONS}-iteration cap without a final answer._",
        }
    )
    yield chatbot, messages[1:], ""


with gr.Blocks(title="LFM2.5 Tool Use", theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        "# 🛠️ LFM2.5-1.2B-Thinking — Tool Calling (ZeroGPU)\n"
        "Runs **in-process on GPU** via ZeroGPU. The model can call "
        "`web_search` and `send_email` (both mocked). Watch it emit tool "
        "calls, execute them, and produce a final answer."
    )
    chatbot = gr.Chatbot(type="messages", height=520, label="Conversation")
    with gr.Row():
        txt = gr.Textbox(
            placeholder="Ask me to search the web or send an email…",
            scale=8,
            show_label=False,
            autofocus=True,
        )
        btn = gr.Button("Send", variant="primary")
    clr = gr.Button("Clear")

    history_state = gr.State([])

    btn.click(respond, [txt, history_state], [chatbot, history_state, txt])
    txt.submit(respond, [txt, history_state], [chatbot, history_state, txt])
    clr.click(
        lambda: ([], [], ""),
        outputs=[chatbot, history_state, txt],
    )

    gr.Examples(examples=EXAMPLES, inputs=txt)


if __name__ == "__main__":
    demo.launch()