from __future__ import annotations

from pathlib import Path

import gradio as gr
import torch
from model import NeuralModelMachine
from safetensors.torch import load_file
from schema import ACTIONS, CAPABILITY_NAMES, STATE_NAMES, reference_transition

ARTIFACT_DIR = (
    Path(__file__).resolve().parent / "artifacts" / "kernelmind-model-machine"
)
MODEL = NeuralModelMachine()
MODEL.load_state_dict(load_file(ARTIFACT_DIR / "model_machine.safetensors"))
MODEL.eval()


@torch.inference_mode()
def run_transition(
    action: str,
    active_state: list[str],
    active_capabilities: list[str],
) -> tuple[dict, dict]:
    state = [int(name in active_state) for name in STATE_NAMES]
    capabilities = [int(name in active_capabilities) for name in CAPABILITY_NAMES]
    action_index = ACTIONS.index(action)
    predicted, blocked_probability = MODEL.transition(
        torch.tensor([state], dtype=torch.float32),
        torch.tensor([capabilities], dtype=torch.float32),
        torch.tensor([action_index]),
    )
    reference_state, reference_blocked = reference_transition(
        state, capabilities, action
    )
    learned_state = predicted[0].int().tolist()
    learned = {
        "next_state": dict(zip(STATE_NAMES, learned_state, strict=True)),
        "blocked_probability": float(blocked_probability[0]),
    }
    reference = {
        "next_state": dict(zip(STATE_NAMES, reference_state, strict=True)),
        "blocked": bool(reference_blocked),
        "exact_match": learned_state == reference_state
        and int(blocked_probability[0] >= 0.5) == reference_blocked,
    }
    return learned, reference


with gr.Blocks(title="KernelMind Neural Model Machine") as demo:
    gr.Markdown(
        "# KernelMind Neural Model Machine\n"
        "A trained world model predicts how an OS action changes virtual machine "
        "state. Compare its transition directly with the deterministic reference."
    )
    action = gr.Dropdown(ACTIONS, value="WRITE_BACKUP", label="Action")
    state = gr.CheckboxGroup(
        STATE_NAMES,
        value=["file_exists", "service_running"],
        label="Current state bits",
    )
    capabilities = gr.CheckboxGroup(
        CAPABILITY_NAMES,
        value=["network"],
        label="Capabilities",
    )
    initial = run_transition(
        "WRITE_BACKUP", ["file_exists", "service_running"], ["network"]
    )
    with gr.Row():
        learned = gr.JSON(value=initial[0], label="Learned transition")
        reference = gr.JSON(value=initial[1], label="Reference transition")
    button = gr.Button("Advance the model machine", variant="primary")
    button.click(
        run_transition,
        inputs=[action, state, capabilities],
        outputs=[learned, reference],
    )


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