"""Minimal example for running Sakura Needle 3 ONNX inference."""

import json
import os
import sys

# Ensure parent directory is on sys.path when running directly
HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.dirname(HERE)
sys.path.insert(0, ROOT)

from runtime.needle_onnx import NeedleONNX

def main():
    model_path = os.path.join(ROOT, "models", "model_int8_compact.onnx")
    print(f"Loading Needle 3 ONNX engine from {model_path}...")
    engine = NeedleONNX(model_path=model_path, num_threads=4)
    print(f"Engine loaded in {engine.load_time_ms:.1f} ms.\n")

    tools = [
        {
            "name": "control_lights",
            "description": "Adjust lights in a specified room",
            "parameters": {
                "type": "object",
                "properties": {
                    "room": {"type": "string", "description": "Room name"},
                    "action": {"type": "string", "enum": ["on", "off", "dim"]},
                    "brightness": {"type": "integer", "description": "Brightness percentage 0-100"},
                },
                "required": ["room", "action"],
            },
        },
        {
            "name": "set_timer",
            "description": "Set a countdown timer",
            "parameters": {
                "type": "object",
                "properties": {
                    "duration_minutes": {"type": "integer"},
                    "label": {"type": "string"},
                },
                "required": ["duration_minutes"],
            },
        },
    ]

    query = "Turn off the kitchen lights and dim the living room lights to 35%."
    print(f"Query: {query}\n")

    system = "date: 2026-10-04 Sun 12:00; You control smart home devices."
    print("Running completion (compatibility='cactus')...")
    res = engine.complete(query, tools=tools, system=system, compatibility="cactus")

    print(f"Prefill: {res['prefill_tps']:.1f} tok/s | Decode: {res['decode_tps']:.1f} tok/s | TTFT: {res['first_token_latency_ms']:.1f} ms")
    print(f"Reasoning: {res['reasoning'].strip()}")
    print("Function calls:")
    print(json.dumps(res["function_calls"], indent=2))

if __name__ == "__main__":
    main()