import gradio as gr
import json
import spaces
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

MODEL_ID = "dispatchAI/Llama-3.2-1B-FunctionCall-mobile"

tokenizer = None
model = None

def load_model():
    global tokenizer, model
    if tokenizer is None:
        tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
        model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype=torch.float16, device_map="auto")
    return tokenizer, model

# Available functions the agent can call
AVAILABLE_FUNCTIONS = {
    "set_alarm": {"description": "Set an alarm", "params": {"time": "string", "label": "string"}},
    "send_message": {"description": "Send a message to a contact", "params": {"to": "string", "message": "string"}},
    "call_contact": {"description": "Call a contact", "params": {"contact": "string"}},
    "search_web": {"description": "Search the web", "params": {"query": "string"}},
    "open_app": {"description": "Open an application", "params": {"app_name": "string"}},
    "set_timer": {"description": "Set a timer", "params": {"duration_minutes": "integer"}},
    "get_weather": {"description": "Get weather for a location", "params": {"location": "string"}},
}

@spaces.GPU
def parse_function_call(user_input: str) -> str:
    """Parse user input into a function call using the mobile function-calling model.
    
    Args:
        user_input: Natural language user input (e.g., "Set an alarm for 7am")
    
    Returns:
        JSON with the parsed function call
    """
    tokenizer, model = load_model()
    
    functions_text = json.dumps(AVAILABLE_FUNCTIONS, indent=2)
    prompt = f"""You are a function-calling assistant. Parse the user's request into a function call.

Available functions:
{functions_text}

User request: "{user_input}"

Respond with ONLY a JSON object, no other text:
{{"function": "function_name", "parameters": {{...}}}}"""
    
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    with torch.no_grad():
        outputs = model.generate(**inputs, max_new_tokens=100, temperature=0.1, do_sample=True, pad_token_id=tokenizer.eos_token_id)
    
    response = tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True).strip()
    
    # Try to parse as JSON
    try:
        # Find JSON in response
        start = response.find("{")
        end = response.rfind("}") + 1
        if start >= 0 and end > start:
            parsed = json.loads(response[start:end])
            func_name = parsed.get("function", "unknown")
            params = parsed.get("parameters", {})
            
            # Check if function exists
            if func_name in AVAILABLE_FUNCTIONS:
                return json.dumps({
                    "status": "success",
                    "function": func_name,
                    "parameters": params,
                    "description": AVAILABLE_FUNCTIONS[func_name]["description"],
                    "user_input": user_input,
                }, indent=2)
            else:
                return json.dumps({
                    "status": "unknown_function",
                    "function": func_name,
                    "parameters": params,
                    "available_functions": list(AVAILABLE_FUNCTIONS.keys()),
                }, indent=2)
    except json.JSONDecodeError:
        pass
    
    # Fallback: keyword-based parsing
    lower = user_input.lower()
    if "alarm" in lower:
        return json.dumps({"status": "fallback", "function": "set_alarm", "parameters": {"time": "07:00", "label": "alarm"}, "source": "keyword_fallback"}, indent=2)
    elif "call" in lower:
        contact = user_input.replace("call", "").strip()
        return json.dumps({"status": "fallback", "function": "call_contact", "parameters": {"contact": contact}, "source": "keyword_fallback"}, indent=2)
    elif "message" in lower or "text" in lower or "send" in lower:
        return json.dumps({"status": "fallback", "function": "send_message", "parameters": {"to": "unknown", "message": ""}, "source": "keyword_fallback"}, indent=2)
    elif "timer" in lower:
        return json.dumps({"status": "fallback", "function": "set_timer", "parameters": {"duration_minutes": 5}, "source": "keyword_fallback"}, indent=2)
    elif "weather" in lower:
        return json.dumps({"status": "fallback", "function": "get_weather", "parameters": {"location": "current"}, "source": "keyword_fallback"}, indent=2)
    
    return json.dumps({"status": "no_match", "message": "Could not parse function call", "available_functions": list(AVAILABLE_FUNCTIONS.keys())}, indent=2)

def list_functions() -> str:
    """List all available functions the agent can call."""
    return json.dumps(AVAILABLE_FUNCTIONS, indent=2)

with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue"), title="dispatchAI Function Calling Agent") as demo:
    gr.Markdown("""
    # 🔧 dispatchAI Function Calling Agent
    
    An on-device agent that parses natural language into structured function calls.
    
    **Model:** [Llama-3.2-1B-FunctionCall-mobile](https://huggingface.co/dispatchAI/Llama-3.2-1B-FunctionCall-mobile) (1B params, Q4 quantized)
    
    Try: "Set an alarm for 7am", "Call mom", "Send a message to John saying I'll be late"
    """)
    
    with gr.Row():
        inp = gr.Textbox(label="User Input", placeholder="Set an alarm for 7am tomorrow", scale=3)
        btn = gr.Button("Parse Function Call", variant="primary", scale=1)
    
    out = gr.Textbox(label="Parsed Function Call (JSON)", lines=12)
    btn.click(fn=parse_function_call, inputs=inp, outputs=out)
    
    with gr.Accordion("Available Functions", open=False):
        fn_btn = gr.Button("List Functions")
        fn_out = gr.Textbox(label="Functions (JSON)", lines=15)
        fn_btn.click(fn=list_functions, outputs=fn_out)
    
    gr.Markdown("---\n🚀 [dispatchAI](https://huggingface.co/dispatchAI) — Small. Mobile. Free. UAE-built.")

demo.launch(mcp_server=True)