from __future__ import annotations

import json
from pathlib import Path

import gradio as gr
import torch
from model import KernelMindPolicy
from runtime import ModelMachine
from safetensors.torch import load_file
from schema import ACTIONS, INTENTS, TARGETS, encode_scenario

ARTIFACT_DIR = Path(__file__).resolve().parent / "artifacts" / "kernelmind-ai-os"
MODEL = KernelMindPolicy()
MODEL.load_state_dict(load_file(ARTIFACT_DIR / "policy.safetensors"))
MODEL.eval()


@torch.inference_mode()
def run_model_machine(
    intent: str,
    target: str,
    administrator: bool,
    network: bool,
    confirmed: bool,
    exists: bool,
    service_running: bool,
) -> tuple[dict, dict]:
    scenario = {
        "intent": intent,
        "target": target,
        "administrator": administrator,
        "network": network,
        "confirmed": confirmed,
        "exists": exists,
        "service_running": service_running,
    }
    tokens = torch.tensor([encode_scenario(scenario)])
    logits = MODEL(tokens)[0]
    indexes = logits.argmax(1).tolist()
    plan = [ACTIONS[index] for index in indexes if ACTIONS[index] != "PAD"]
    confidence = [
        float(torch.softmax(slot, dim=0).max())
        for slot, index in zip(logits, indexes, strict=True)
        if ACTIONS[index] != "PAD"
    ]
    machine = ModelMachine(scenario)
    before = json.loads(json.dumps(machine.state))
    after = machine.execute(plan)
    decision = {
        "learned_plan": plan,
        "action_confidence": confidence,
        "capability_gate": "independent deterministic runtime enforcement",
    }
    execution = {"before": before, "after": after}
    return decision, execution


with gr.Blocks(title="KernelMind AI OS") as demo:
    gr.Markdown(
        "# KernelMind AI OS + Model Machine\n"
        "A trained policy kernel converts intent and machine state into a structured "
        "plan. A separate capability gate executes it only inside a virtual computer."
    )
    with gr.Row():
        intent = gr.Dropdown(INTENTS, value="backup", label="Intent")
        target = gr.Dropdown(TARGETS, value="notes", label="Target")
    with gr.Row():
        administrator = gr.Checkbox(False, label="Administrator capability")
        network = gr.Checkbox(True, label="Network available")
        confirmed = gr.Checkbox(False, label="Destructive action confirmed")
        exists = gr.Checkbox(True, label="Target exists")
        service = gr.Checkbox(True, label="Service running")
    initial = run_model_machine("backup", "notes", False, True, False, True, True)
    with gr.Row():
        decision = gr.JSON(value=initial[0], label="AI OS policy decision")
        execution = gr.JSON(value=initial[1], label="Model Machine transition")
    button = gr.Button("Plan and execute safely", variant="primary")
    button.click(
        run_model_machine,
        inputs=[
            intent,
            target,
            administrator,
            network,
            confirmed,
            exists,
            service,
        ],
        outputs=[decision, execution],
    )


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