package transport

import (
	"testing"

	toki "git.toki-labs.com/toki/common-proto-socket/go"
	"google.golang.org/protobuf/proto"

	iop "iop/proto/gen/iop"
)

func TestNodeParserMap_RunRequest(t *testing.T) {
	parsers := nodeParserMap()
	original := &iop.RunRequest{
		RunId:       "run-1",
		Adapter:     "mock",
		Model:       "v1",
		SessionId:   "session-a",
		Background:  true,
		SessionMode: iop.RunSessionMode_RUN_SESSION_MODE_REQUIRE_EXISTING,
	}
	payload, err := proto.Marshal(original)
	if err != nil {
		t.Fatalf("marshal: %v", err)
	}

	key := toki.TypeNameOf(original)
	parser, ok := parsers[key]
	if !ok {
		t.Fatalf("parser not found for key: %s", key)
	}

	parsed, err := parser(payload)
	if err != nil {
		t.Fatalf("parse: %v", err)
	}
	got := parsed.(*iop.RunRequest)
	if got.GetRunId() != original.GetRunId() ||
		got.GetAdapter() != original.GetAdapter() ||
		got.GetModel() != original.GetModel() ||
		got.GetSessionId() != original.GetSessionId() ||
		got.GetBackground() != original.GetBackground() ||
		got.GetSessionMode() != original.GetSessionMode() {
		t.Fatalf("unexpected run request: %+v", got)
	}
}

func TestNodeParserMap_CancelRequest(t *testing.T) {
	parsers := nodeParserMap()
	original := &iop.CancelRequest{
		RunId:     "run-1",
		Adapter:   "cli",
		Model:     "codex",
		SessionId: "session-a",
		Action:    iop.CancelAction_CANCEL_ACTION_TERMINATE_SESSION,
	}
	payload, err := proto.Marshal(original)
	if err != nil {
		t.Fatalf("marshal: %v", err)
	}

	key := toki.TypeNameOf(original)
	parser, ok := parsers[key]
	if !ok {
		t.Fatalf("parser not found for key: %s", key)
	}

	parsed, err := parser(payload)
	if err != nil {
		t.Fatalf("parse: %v", err)
	}
	got := parsed.(*iop.CancelRequest)
	if got.GetRunId() != original.GetRunId() ||
		got.GetAdapter() != original.GetAdapter() ||
		got.GetModel() != original.GetModel() ||
		got.GetSessionId() != original.GetSessionId() ||
		got.GetAction() != original.GetAction() {
		t.Fatalf("unexpected cancel request: %+v", got)
	}
}

func TestNodeParserMap_NodeCommandRequest(t *testing.T) {
	parsers := nodeParserMap()
	original := &iop.NodeCommandRequest{
		Type:      iop.NodeCommandType_NODE_COMMAND_TYPE_USAGE_STATUS,
		Adapter:   "cli",
		Model:     "codex",
		SessionId: "default",
	}
	payload, err := proto.Marshal(original)
	if err != nil {
		t.Fatalf("marshal: %v", err)
	}

	key := toki.TypeNameOf(original)
	parser, ok := parsers[key]
	if !ok {
		t.Fatalf("parser not found for key: %s", key)
	}

	parsed, err := parser(payload)
	if err != nil {
		t.Fatalf("parse: %v", err)
	}
	got := parsed.(*iop.NodeCommandRequest)
	if got.GetType() != original.GetType() ||
		got.GetAdapter() != original.GetAdapter() ||
		got.GetModel() != original.GetModel() ||
		got.GetSessionId() != original.GetSessionId() {
		t.Fatalf("unexpected node command request: %+v", got)
	}
}

func TestNodeParserMap_RegisterResponse(t *testing.T) {
	parsers := nodeParserMap()
	original := &iop.RegisterResponse{Accepted: true, NodeId: "node-1", Alias: "local-node"}
	payload, err := proto.Marshal(original)
	if err != nil {
		t.Fatalf("marshal: %v", err)
	}

	key := toki.TypeNameOf(original)
	parser, ok := parsers[key]
	if !ok {
		t.Fatalf("parser not found for key: %s", key)
	}

	parsed, err := parser(payload)
	if err != nil {
		t.Fatalf("parse: %v", err)
	}
	got, ok := parsed.(*iop.RegisterResponse)
	if !ok {
		t.Fatalf("unexpected register response type: %T", parsed)
	}
	if !got.GetAccepted() || got.GetNodeId() != original.GetNodeId() || got.GetAlias() != original.GetAlias() {
		t.Fatalf("unexpected register response: %+v", got)
	}
}