package transport

import (
	"testing"

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

	iop "iop/proto/gen/iop"
)

func TestNodeParserMap_RunRequest(t *testing.T) {
	parsers := nodeParserMap()
	original := &iop.RunRequest{
		RunId:      "run-1",
		Adapter:    "mock",
		Target:     "v1",
		SessionId:  "session-a",
		Background: true,
	}
	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.GetTarget() != original.GetTarget() ||
		got.GetSessionId() != original.GetSessionId() ||
		got.GetBackground() != original.GetBackground() {
		t.Fatalf("unexpected run request: %+v", got)
	}
}

func TestNodeParserMapWorkspace(t *testing.T) {
	parsers := nodeParserMap()
	cases := []proto.Message{
		&iop.WorkspaceOpenRequest{
			RequestId: "request-1", WorkspaceRef: "workspace-1", TimeoutMs: 1000,
			Operations:   []iop.WorkspaceOperation{iop.WorkspaceOperation_WORKSPACE_OPERATION_READ, iop.WorkspaceOperation_WORKSPACE_OPERATION_WRITE},
			MaxReadBytes: 64, MaxWriteBytes: 64,
		},
		&iop.WorkspaceToolRequest{
			RequestId: "request-1", StageId: "work", ToolCallId: "tool-1", Operation: iop.WorkspaceOperation_WORKSPACE_OPERATION_WRITE,
			Input: &iop.WorkspaceToolRequest_Write{Write: &iop.WorkspaceWriteInput{RelativePath: "output.txt", Content: []byte("bounded")}},
		},
		&iop.WorkspaceToolRequest{RequestId: "request-1", StageId: "work", ToolCallId: "legacy-write", Operation: iop.WorkspaceOperation_WORKSPACE_OPERATION_WRITE, Input: &iop.WorkspaceToolRequest_WriteContent{WriteContent: []byte("legacy")}},
		&iop.WorkspaceCancelRequest{RequestId: "request-1", StageId: "work", ToolCallId: "tool-1"},
		&iop.WorkspaceCleanupRequest{RequestId: "request-1"},
	}
	for _, original := range cases {
		payload, err := proto.Marshal(original)
		if err != nil {
			t.Fatalf("marshal %T: %v", original, err)
		}
		parser, ok := parsers[toki.TypeNameOf(original)]
		if !ok {
			t.Fatalf("parser not found for %T", original)
		}
		parsed, err := parser(payload)
		if err != nil {
			t.Fatalf("parse %T: %v", original, err)
		}
		if !proto.Equal(parsed, original) {
			t.Fatalf("round trip %T = %v, want %v", original, parsed, original)
		}
	}
	toolFields := (&iop.WorkspaceToolRequest{}).ProtoReflect().Descriptor().Fields()
	for name, number := range map[string]int32{
		"request_id": 1, "stage_id": 2, "tool_call_id": 3, "operation": 4, "timeout_ms": 5,
		"relative_path": 6, "write_content": 7, "command_id": 8, "environment": 9, "write": 10,
	} {
		field := toolFields.ByName(protoreflect.Name(name))
		if field == nil || int32(field.Number()) != number {
			t.Fatalf("WorkspaceToolRequest.%s number = %v, want %d", name, field, number)
		}
	}
}

func TestNodeParserMap_ProviderTunnelRequest(t *testing.T) {
	parsers := nodeParserMap()
	original := &iop.ProviderTunnelRequest{
		RunId:      "run-1",
		TunnelId:   "tunnel-1",
		Adapter:    "openai_compat",
		Target:     "qwen",
		Method:     "POST",
		Path:       "/v1/chat/completions",
		Headers:    map[string]string{"content-type": "application/json"},
		Body:       []byte(`{"stream":true}`),
		Stream:     true,
		TimeoutSec: 30,
		Metadata:   map[string]string{"openai_model": "qwen"},
		SessionId:  "openai",
	}
	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.ProviderTunnelRequest)
	if got.GetRunId() != original.GetRunId() ||
		got.GetTunnelId() != original.GetTunnelId() ||
		got.GetAdapter() != original.GetAdapter() ||
		got.GetTarget() != original.GetTarget() ||
		got.GetMethod() != original.GetMethod() ||
		got.GetPath() != original.GetPath() ||
		got.GetHeaders()["content-type"] != "application/json" ||
		string(got.GetBody()) != string(original.GetBody()) ||
		got.GetStream() != original.GetStream() ||
		got.GetTimeoutSec() != original.GetTimeoutSec() ||
		got.GetMetadata()["openai_model"] != "qwen" ||
		got.GetSessionId() != original.GetSessionId() {
		t.Fatalf("unexpected tunnel request: %+v", got)
	}
}

func TestNodeParserMap_CancelRequest(t *testing.T) {
	parsers := nodeParserMap()
	original := &iop.CancelRequest{RunId: "run-1"}
	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() {
		t.Fatalf("unexpected cancel request: %+v", got)
	}
}

func TestNodeParserMap_EdgeNodeEvent(t *testing.T) {
	parsers := nodeParserMap()
	original := &iop.EdgeNodeEvent{
		Type:   "edge.disconnected",
		Source: "node",
		NodeId: "node-1",
		Alias:  "local-node",
		Reason: "transport_closed",
	}
	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.EdgeNodeEvent)
	if got.GetType() != original.GetType() ||
		got.GetSource() != original.GetSource() ||
		got.GetNodeId() != original.GetNodeId() ||
		got.GetAlias() != original.GetAlias() ||
		got.GetReason() != original.GetReason() {
		t.Fatalf("unexpected edge node event: %+v", got)
	}
}

func TestNodeParserMap_NodeCommandRequest(t *testing.T) {
	parsers := nodeParserMap()
	original := &iop.NodeCommandRequest{
		Type:      iop.NodeCommandType_NODE_COMMAND_TYPE_CAPABILITIES,
		Adapter:   "ollama",
		Target:    "model",
		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.GetTarget() != original.GetTarget() ||
		got.GetSessionId() != original.GetSessionId() {
		t.Fatalf("unexpected node command request: %+v", got)
	}
}

func TestNodeParserMap_NodeCommandRequest_NewTypes(t *testing.T) {
	parsers := nodeParserMap()
	types := []iop.NodeCommandType{
		iop.NodeCommandType_NODE_COMMAND_TYPE_CAPABILITIES,
		iop.NodeCommandType_NODE_COMMAND_TYPE_TRANSPORT_STATUS,
		iop.NodeCommandType_NODE_COMMAND_TYPE_OLLAMA_API,
	}
	for _, cmdType := range types {
		original := &iop.NodeCommandRequest{
			RequestId: "req-1",
			Type:      cmdType,
			Adapter:   "ollama",
			Target:    "model",
			SessionId: "default",
		}
		payload, err := proto.Marshal(original)
		if err != nil {
			t.Fatalf("marshal %v: %v", cmdType, err)
		}
		parser, ok := parsers[toki.TypeNameOf(original)]
		if !ok {
			t.Fatalf("parser not found for NodeCommandRequest")
		}
		parsed, err := parser(payload)
		if err != nil {
			t.Fatalf("parse %v: %v", cmdType, err)
		}
		got := parsed.(*iop.NodeCommandRequest)
		if got.GetType() != cmdType {
			t.Fatalf("type: got %v want %v", got.GetType(), cmdType)
		}
	}
}

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)
	}
}