package runnerregistry

import (
	"testing"
	"time"

	otopb "github.com/toki/oto/services/core/oto"
)

func TestRegistryAcceptsRunnerRegistration(t *testing.T) {
	acceptedAt := time.Date(2026, 6, 5, 12, 0, 0, 0, time.UTC)
	registry := NewWithClock(func() time.Time { return acceptedAt })

	response := registry.Register(&otopb.RegisterRunnerRequest{
		EnrollmentToken: "token-123",
		RunnerId:        "runner-123",
		Alias:           "build-linux",
		ProtocolVersion: "oto.runner.v1",
		Capability: &otopb.RunnerCapability{
			Name:    "oto-runner",
			Version: "1.0.0",
		},
		CommandCatalog: &otopb.CommandCatalogSummary{
			CommandTypes: []string{"Shell", "Git"},
		},
	})

	if !response.GetAccepted() {
		t.Fatalf("Register accepted=false, reason=%q", response.GetRejectReason())
	}
	if response.GetRunnerId() != "runner-123" {
		t.Fatalf("response runner id = %q, want runner-123", response.GetRunnerId())
	}

	record, ok := registry.Snapshot("runner-123")
	if !ok {
		t.Fatal("registered runner was not stored")
	}
	if record.Alias != "build-linux" {
		t.Fatalf("record alias = %q, want build-linux", record.Alias)
	}
	if record.ProtocolVersion != "oto.runner.v1" {
		t.Fatalf("record protocol = %q, want oto.runner.v1", record.ProtocolVersion)
	}
	if record.Capability.Name != "oto-runner" || record.Capability.Version != "1.0.0" {
		t.Fatalf("record capability = %#v", record.Capability)
	}
	if record.Status != StatusAccepted {
		t.Fatalf("record status = %q, want %q", record.Status, StatusAccepted)
	}
	if !record.AcceptedAt.Equal(acceptedAt) {
		t.Fatalf("record accepted at = %v, want %v", record.AcceptedAt, acceptedAt)
	}
	if got := record.CommandTypes; len(got) != 2 || got[0] != "Shell" || got[1] != "Git" {
		t.Fatalf("record command types = %#v, want [Shell Git]", got)
	}
}

func TestRegistryRejectsIncompleteRegistration(t *testing.T) {
	tests := []struct {
		name string
		req  *otopb.RegisterRunnerRequest
		want string
	}{
		{
			name: "missing token",
			req: &otopb.RegisterRunnerRequest{
				RunnerId:        "runner-123",
				ProtocolVersion: "oto.runner.v1",
			},
			want: "missing enrollment token",
		},
		{
			name: "missing runner id",
			req: &otopb.RegisterRunnerRequest{
				EnrollmentToken: "token-123",
				ProtocolVersion: "oto.runner.v1",
			},
			want: "missing runner id",
		},
		{
			name: "missing protocol",
			req: &otopb.RegisterRunnerRequest{
				EnrollmentToken: "token-123",
				RunnerId:        "runner-123",
			},
			want: "missing protocol version",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			response := New().Register(tt.req)
			if response.GetAccepted() {
				t.Fatal("Register accepted incomplete request")
			}
			if response.GetRejectReason() != tt.want {
				t.Fatalf("reject reason = %q, want %q", response.GetRejectReason(), tt.want)
			}
		})
	}
}

func TestRegistryHeartbeatTransitionsStatus(t *testing.T) {
	var mockTime time.Time
	registry := NewWithClock(func() time.Time { return mockTime })

	// 1. Heartbeat on unknown runner should fail
	hbRes := registry.Heartbeat("runner-123", otopb.HeartbeatStatus_HEARTBEAT_STATUS_HEALTHY)
	if hbRes.GetSuccess() {
		t.Fatal("expected heartbeat on unknown runner to fail")
	}

	// 2. Register runner
	mockTime = time.Date(2026, 6, 5, 12, 0, 0, 0, time.UTC)
	registry.Register(&otopb.RegisterRunnerRequest{
		EnrollmentToken: "token-123",
		RunnerId:        "runner-123",
		ProtocolVersion: "oto.runner.v1",
		Capability: &otopb.RunnerCapability{
			Name:    "oto-runner",
			Version: "1.0.0",
		},
	})

	record, _ := registry.Snapshot("runner-123")
	if record.Status != StatusAccepted {
		t.Fatalf("status = %q, want %q", record.Status, StatusAccepted)
	}

	// 3. First heartbeat transitions to Online
	mockTime = mockTime.Add(5 * time.Second)
	hbRes = registry.Heartbeat("runner-123", otopb.HeartbeatStatus_HEARTBEAT_STATUS_HEALTHY)
	if !hbRes.GetSuccess() {
		t.Fatalf("heartbeat failed: %s", hbRes.GetErrorMessage())
	}

	record, _ = registry.Snapshot("runner-123")
	if record.Status != StatusOnline {
		t.Fatalf("status = %q, want %q", record.Status, StatusOnline)
	}
	if !record.FirstHeartbeatAt.Equal(mockTime) {
		t.Fatalf("FirstHeartbeatAt = %v, want %v", record.FirstHeartbeatAt, mockTime)
	}
	if !record.LastHeartbeatAt.Equal(mockTime) {
		t.Fatalf("LastHeartbeatAt = %v, want %v", record.LastHeartbeatAt, mockTime)
	}

	// 4. Subsequent heartbeat updates LastHeartbeatAt but not FirstHeartbeatAt
	mockTime = mockTime.Add(10 * time.Second)
	hbRes = registry.Heartbeat("runner-123", otopb.HeartbeatStatus_HEARTBEAT_STATUS_HEALTHY)
	if !hbRes.GetSuccess() {
		t.Fatalf("heartbeat failed: %s", hbRes.GetErrorMessage())
	}

	record, _ = registry.Snapshot("runner-123")
	if record.Status != StatusOnline {
		t.Fatalf("status = %q, want %q", record.Status, StatusOnline)
	}
	if !record.FirstHeartbeatAt.Equal(mockTime.Add(-10 * time.Second)) {
		t.Fatalf("FirstHeartbeatAt should remain unchanged")
	}
	if !record.LastHeartbeatAt.Equal(mockTime) {
		t.Fatalf("LastHeartbeatAt = %v, want %v", record.LastHeartbeatAt, mockTime)
	}
}

func TestRegistryDisconnect(t *testing.T) {
	registry := New()
	registry.Register(&otopb.RegisterRunnerRequest{
		EnrollmentToken: "token-123",
		RunnerId:        "runner-123",
		ProtocolVersion: "oto.runner.v1",
		Capability: &otopb.RunnerCapability{
			Name:    "oto-runner",
			Version: "1.0.0",
		},
	})

	if ok := registry.Disconnect("unknown-runner"); ok {
		t.Fatal("expected disconnect on unknown runner to fail")
	}

	if ok := registry.Disconnect("runner-123"); !ok {
		t.Fatal("expected disconnect to succeed")
	}

	record, _ := registry.Snapshot("runner-123")
	if record.Status != StatusDisconnected {
		t.Fatalf("status = %q, want %q", record.Status, StatusDisconnected)
	}
}

func TestRegistryUpdatesAliasAndPreservesItOnReregister(t *testing.T) {
	registry := New()
	registry.Register(&otopb.RegisterRunnerRequest{
		EnrollmentToken: "token-123",
		RunnerId:        "runner-123",
		Alias:           "device-default",
		ProtocolVersion: "oto.runner.v1",
		Capability: &otopb.RunnerCapability{
			Name:    "oto-runner",
			Version: "1.0.0",
		},
	})

	record, ok := registry.UpdateAlias("runner-123", "friendly-name")
	if !ok {
		t.Fatal("expected alias update to succeed")
	}
	if record.Alias != "friendly-name" {
		t.Fatalf("alias = %q, want friendly-name", record.Alias)
	}

	registry.Register(&otopb.RegisterRunnerRequest{
		EnrollmentToken: "token-123",
		RunnerId:        "runner-123",
		Alias:           "device-default",
		ProtocolVersion: "oto.runner.v1",
		Capability: &otopb.RunnerCapability{
			Name:    "oto-runner",
			Version: "1.0.0",
		},
	})

	record, ok = registry.Snapshot("runner-123")
	if !ok {
		t.Fatal("expected runner to exist")
	}
	if record.Alias != "friendly-name" {
		t.Fatalf("alias after reregister = %q, want friendly-name", record.Alias)
	}

	if _, ok := registry.UpdateAlias("unknown-runner", "name"); ok {
		t.Fatal("expected alias update on unknown runner to fail")
	}
}

func TestRegistryCheckTimeouts(t *testing.T) {
	var mockTime time.Time
	registry := NewWithClock(func() time.Time { return mockTime })

	// Register 2 runners:
	// runner-accepted: remains in accepted state (no heartbeat)
	// runner-online: transitions to online state
	mockTime = time.Date(2026, 6, 5, 12, 0, 0, 0, time.UTC)
	registry.Register(&otopb.RegisterRunnerRequest{
		EnrollmentToken: "token-123",
		RunnerId:        "runner-accepted",
		ProtocolVersion: "oto.runner.v1",
		Capability: &otopb.RunnerCapability{
			Name:    "oto-runner",
			Version: "1.0.0",
		},
	})
	registry.Register(&otopb.RegisterRunnerRequest{
		EnrollmentToken: "token-123",
		RunnerId:        "runner-online",
		ProtocolVersion: "oto.runner.v1",
		Capability: &otopb.RunnerCapability{
			Name:    "oto-runner",
			Version: "1.0.0",
		},
	})

	// transition runner-online to online
	mockTime = mockTime.Add(1 * time.Second)
	registry.Heartbeat("runner-online", otopb.HeartbeatStatus_HEARTBEAT_STATUS_HEALTHY)

	// verify initial states
	recAccepted, _ := registry.Snapshot("runner-accepted")
	recOnline, _ := registry.Snapshot("runner-online")
	if recAccepted.Status != StatusAccepted || recOnline.Status != StatusOnline {
		t.Fatal("initial status setup mismatch")
	}

	// Let 30 seconds pass.
	// CheckTimeouts with 10s threshold. Both should timeout.
	mockTime = mockTime.Add(30 * time.Second)
	registry.CheckTimeouts(10 * time.Second)

	recAccepted, _ = registry.Snapshot("runner-accepted")
	recOnline, _ = registry.Snapshot("runner-online")

	if recAccepted.Status != StatusHeartbeatTimeout {
		t.Fatalf("runner-accepted status = %q, want %q", recAccepted.Status, StatusHeartbeatTimeout)
	}
	if recOnline.Status != StatusHeartbeatTimeout {
		t.Fatalf("runner-online status = %q, want %q", recOnline.Status, StatusHeartbeatTimeout)
	}
	if recAccepted.FailureReason != "heartbeat timeout before first heartbeat" {
		t.Fatalf("runner-accepted failure reason: %q", recAccepted.FailureReason)
	}
	if recOnline.FailureReason != "heartbeat timeout" {
		t.Fatalf("runner-online failure reason: %q", recOnline.FailureReason)
	}
}

func TestRegistryHeartbeatTerminalStateDefense(t *testing.T) {
	registry := New()
	registry.Register(&otopb.RegisterRunnerRequest{
		EnrollmentToken: "token-123",
		RunnerId:        "runner-123",
		ProtocolVersion: "oto.runner.v1",
		Capability: &otopb.RunnerCapability{
			Name:    "oto-runner",
			Version: "1.0.0",
		},
	})

	// Move to Online
	registry.Heartbeat("runner-123", otopb.HeartbeatStatus_HEARTBEAT_STATUS_HEALTHY)

	// Disconnect
	registry.Disconnect("runner-123")
	rec, _ := registry.Snapshot("runner-123")
	if rec.Status != StatusDisconnected {
		t.Fatalf("expected disconnected, got %s", rec.Status)
	}

	// Heartbeat should fail and keep state as disconnected
	hbRes := registry.Heartbeat("runner-123", otopb.HeartbeatStatus_HEARTBEAT_STATUS_HEALTHY)
	if hbRes.GetSuccess() {
		t.Fatal("expected heartbeat to be rejected on disconnected runner")
	}
	rec, _ = registry.Snapshot("runner-123")
	if rec.Status != StatusDisconnected {
		t.Fatalf("expected status to remain disconnected, got %s", rec.Status)
	}

	// Re-register to make it accepted again
	registry.Register(&otopb.RegisterRunnerRequest{
		EnrollmentToken: "token-123",
		RunnerId:        "runner-123",
		ProtocolVersion: "oto.runner.v1",
		Capability: &otopb.RunnerCapability{
			Name:    "oto-runner",
			Version: "1.0.0",
		},
	})
	rec, _ = registry.Snapshot("runner-123")
	if rec.Status != StatusAccepted {
		t.Fatalf("expected status to be accepted, got %s", rec.Status)
	}

	// Revive to online via heartbeat
	registry.Heartbeat("runner-123", otopb.HeartbeatStatus_HEARTBEAT_STATUS_HEALTHY)

	// Force timeout
	var mockTime time.Time
	registry = NewWithClock(func() time.Time { return mockTime })
	mockTime = time.Date(2026, 6, 5, 12, 0, 0, 0, time.UTC)
	registry.Register(&otopb.RegisterRunnerRequest{
		EnrollmentToken: "token-123",
		RunnerId:        "runner-123",
		ProtocolVersion: "oto.runner.v1",
		Capability: &otopb.RunnerCapability{
			Name:    "oto-runner",
			Version: "1.0.0",
		},
	})
	mockTime = mockTime.Add(1 * time.Second)
	registry.Heartbeat("runner-123", otopb.HeartbeatStatus_HEARTBEAT_STATUS_HEALTHY)

	// pass timeout threshold
	mockTime = mockTime.Add(30 * time.Second)
	registry.CheckTimeouts(10 * time.Second)
	rec, _ = registry.Snapshot("runner-123")
	if rec.Status != StatusHeartbeatTimeout {
		t.Fatalf("expected status heartbeat_timeout, got %s", rec.Status)
	}

	// Heartbeat should fail and keep state as heartbeat_timeout
	hbRes = registry.Heartbeat("runner-123", otopb.HeartbeatStatus_HEARTBEAT_STATUS_HEALTHY)
	if hbRes.GetSuccess() {
		t.Fatal("expected heartbeat to be rejected on timed out runner")
	}
	rec, _ = registry.Snapshot("runner-123")
	if rec.Status != StatusHeartbeatTimeout {
		t.Fatalf("expected status to remain heartbeat_timeout, got %s", rec.Status)
	}
}

func TestRegistryRejectsIncompatibleRunnerContract(t *testing.T) {
	tests := []struct {
		name     string
		req      *otopb.RegisterRunnerRequest
		wantMsg  string
		wantCode string
	}{
		{
			name: "unsupported protocol version",
			req: &otopb.RegisterRunnerRequest{
				EnrollmentToken: "token-123",
				RunnerId:        "runner-123",
				ProtocolVersion: "oto.runner.v2",
				Capability: &otopb.RunnerCapability{
					Name:    "oto-runner",
					Version: "1.0.0",
				},
			},
			wantMsg:  `unsupported protocol version: "oto.runner.v2"`,
			wantCode: "incompatible_runner",
		},
		{
			name: "missing capability",
			req: &otopb.RegisterRunnerRequest{
				EnrollmentToken: "token-123",
				RunnerId:        "runner-123",
				ProtocolVersion: "oto.runner.v1",
			},
			wantMsg:  "missing runner capability",
			wantCode: "incompatible_runner",
		},
		{
			name: "unsupported capability name",
			req: &otopb.RegisterRunnerRequest{
				EnrollmentToken: "token-123",
				RunnerId:        "runner-123",
				ProtocolVersion: "oto.runner.v1",
				Capability: &otopb.RunnerCapability{
					Name:    "other-runner",
					Version: "1.0.0",
				},
			},
			wantMsg:  `unsupported capability name: "other-runner"`,
			wantCode: "incompatible_runner",
		},
		{
			name: "missing capability version",
			req: &otopb.RegisterRunnerRequest{
				EnrollmentToken: "token-123",
				RunnerId:        "runner-123",
				ProtocolVersion: "oto.runner.v1",
				Capability: &otopb.RunnerCapability{
					Name:    "oto-runner",
					Version: "",
				},
			},
			wantMsg:  "missing capability version",
			wantCode: "incompatible_runner",
		},
		{
			name: "incompatible capability major version",
			req: &otopb.RegisterRunnerRequest{
				EnrollmentToken: "token-123",
				RunnerId:        "runner-123",
				ProtocolVersion: "oto.runner.v1",
				Capability: &otopb.RunnerCapability{
					Name:    "oto-runner",
					Version: "2.0.0",
				},
			},
			wantMsg:  `incompatible capability major version: "2.0.0", want major version 1`,
			wantCode: "incompatible_runner",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			response := New().Register(tt.req)
			if response.GetAccepted() {
				t.Fatal("Register accepted incompatible contract")
			}
			if response.GetRejectReason() != tt.wantMsg {
				t.Fatalf("reject reason = %q, want %q", response.GetRejectReason(), tt.wantMsg)
			}
			if response.GetError() == nil {
				t.Fatal("missing structured error")
			}
			if response.GetError().GetCode() != tt.wantCode {
				t.Fatalf("error code = %q, want %q", response.GetError().GetCode(), tt.wantCode)
			}
			if response.GetError().GetMessage() != tt.wantMsg {
				t.Fatalf("error message = %q, want %q", response.GetError().GetMessage(), tt.wantMsg)
			}
		})
	}
}
