package millproto

import (
	"bytes"
	"encoding/binary"
	"strings"
	"testing"

	millv1 "tangled.org/core/spindle/mill/proto/gen"
	"tangled.org/core/spindle/quota"
)

func TestEncodeDecodeRoundTrip(t *testing.T) {
	var buf bytes.Buffer
	enc := NewEncoder(&buf)

	want := &Message{
		ReserveSeat: &millv1.ReserveSeat{
			LeaseId:         "lease-1",
			TargetEngine:    "microvm",
			RawWorkflowJson: `{"name":"build"}`,
			Knot:            "knot.example",
			Rkey:            "abc123",
			TtlSeconds:      30,
			RepoDid:         "did:web:example.com",
		},
	}
	if err := enc.Encode(want); err != nil {
		t.Fatalf("Encode() error = %v", err)
	}

	got, err := NewDecoder(&buf).Decode()
	if err != nil {
		t.Fatalf("Decode() error = %v", err)
	}
	rs := got.GetReserveSeat()
	if rs == nil {
		t.Fatal("decoded message missing reserve_seat")
	}
	if rs.LeaseId != "lease-1" || rs.TargetEngine != "microvm" || rs.TtlSeconds != 30 || rs.RepoDid != "did:web:example.com" {
		t.Fatalf("round-trip mismatch: %+v", rs)
	}
}

func TestDecoderRejectsOversizedMessage(t *testing.T) {
	var tooLarge bytes.Buffer
	var header [4]byte
	binary.BigEndian.PutUint32(header[:], MaxMessageBytes+1)
	tooLarge.Write(header[:])

	if _, err := NewDecoder(&tooLarge).Decode(); err == nil {
		t.Fatal("expected oversized message error")
	}
}

func TestValidationRules(t *testing.T) {
	tests := []struct {
		name    string
		msg     *Message
		wantErr bool
	}{
		{
			name: "valid ack message",
			msg: &Message{
				Ack: &millv1.Ack{
					Epoch:     "inc-1",
					UpToSeqno: 5,
				},
			},
			wantErr: false,
		},
		{
			name: "valid hello message",
			msg: &Message{
				Hello: &millv1.Hello{ProtocolVersion: 1,
					Arch:           "amd64",
					Labels:         []string{"linux"},
					Epoch:          "inc-1",
					CacheNamespace: "tar-zstd-v1/linux/amd64",
				},
			},
			wantErr: false,
		},
		{
			name:    "invalid message with zero payloads",
			msg:     &Message{},
			wantErr: true,
		},
		{
			name: "invalid message with multiple payloads",
			msg: &Message{
				Ack:       &millv1.Ack{Epoch: "inc-1", UpToSeqno: 5},
				Committed: &millv1.Committed{LeaseId: "x"},
			},
			wantErr: true,
		},
		{
			name: "invalid ack message missing epoch",
			msg: &Message{
				Ack: &millv1.Ack{
					UpToSeqno: 5,
				},
			},
			wantErr: true,
		},
		{
			name: "invalid node snapshot with zero seqno",
			msg: &Message{
				NodeSnapshot: &millv1.NodeSnapshot{
					Seqno: 0,
				},
			},
			wantErr: true,
		},
		{
			name: "valid node snapshot with positive seqno",
			msg: &Message{
				NodeSnapshot: &millv1.NodeSnapshot{
					Seqno: 1,
				},
			},
			wantErr: false,
		},
		{
			name: "invalid stream batch with zero seqno entry",
			msg: &Message{
				EventBatch: &millv1.EventBatch{
					Epoch: "inc-1",
					Events: []*millv1.Event{
						{
							Seqno:   0,
							LeaseId: "lease-1",
							Payload: &millv1.Event_StatusEvent{
								StatusEvent: &millv1.StatusEvent{
									Status: millv1.NonterminalStatus_NONTERMINAL_STATUS_RUNNING,
								},
							},
						},
					},
				},
			},
			wantErr: true,
		},
		{
			name: "invalid stream batch with empty entries",
			msg: &Message{
				EventBatch: &millv1.EventBatch{
					Epoch:  "inc-1",
					Events: []*millv1.Event{},
				},
			},
			wantErr: true,
		},
		{
			name: "invalid reserve result with unknown enum",
			msg: &Message{
				ReserveResult: &millv1.ReserveResult{
					LeaseId:     "lease-1",
					RejectClass: millv1.RejectClass(99),
				},
			},
			wantErr: true,
		},
		{
			name: "invalid stream entry with malformed oneof (empty payload)",
			msg: &Message{
				EventBatch: &millv1.EventBatch{
					Epoch: "inc-1",
					Events: []*millv1.Event{
						{
							Seqno:   1,
							LeaseId: "lease-1",
							Payload: nil,
						},
					},
				},
			},
			wantErr: true,
		},
		{
			name: "valid stream batch status event",
			msg: &Message{
				EventBatch: &millv1.EventBatch{
					Epoch: "inc-1",
					Events: []*millv1.Event{
						{
							Seqno:   1,
							LeaseId: "lease-1",
							Payload: &millv1.Event_StatusEvent{
								StatusEvent: &millv1.StatusEvent{
									Status: millv1.NonterminalStatus_NONTERMINAL_STATUS_RUNNING,
								},
							},
						},
					},
				},
			},
			wantErr: false,
		},
		{
			name: "valid stream batch attempt result",
			msg: &Message{
				EventBatch: &millv1.EventBatch{
					Epoch: "inc-1",
					Events: []*millv1.Event{
						{
							Seqno:   1,
							LeaseId: "lease-1",
							Payload: &millv1.Event_AttemptResult{
								AttemptResult: &millv1.AttemptResult{
									Status: millv1.TerminalStatus_TERMINAL_STATUS_SUCCESS,
								},
							},
						},
					},
				},
			},
			wantErr: false,
		},
		{
			name: "valid quota reserve request",
			msg: &Message{
				QuotaReq: &millv1.QuotaRequest{
					RequestId: "req-1",
					Operation: millv1.QuotaOperation_QUOTA_OPERATION_RESERVE,
					LeaseId:   "lease-1",
					Kind:      "nix_cache",
					Key:       "key-1",
					Resources: map[string]int64{"cache_storage_bytes": 100},
				},
			},
			wantErr: false,
		},
		{
			name: "invalid quota request - empty request_id",
			msg: &Message{
				QuotaReq: &millv1.QuotaRequest{
					RequestId: "",
					Operation: millv1.QuotaOperation_QUOTA_OPERATION_RESERVE,
					LeaseId:   "lease-1",
				},
			},
			wantErr: true,
		},
		{
			name: "invalid quota request - unspecified operation",
			msg: &Message{
				QuotaReq: &millv1.QuotaRequest{
					RequestId: "req-1",
					Operation: millv1.QuotaOperation_QUOTA_OPERATION_UNSPECIFIED,
					LeaseId:   "lease-1",
				},
			},
			wantErr: true,
		},
		{
			name: "invalid quota request - empty lease_id",
			msg: &Message{
				QuotaReq: &millv1.QuotaRequest{
					RequestId: "req-1",
					Operation: millv1.QuotaOperation_QUOTA_OPERATION_RESERVE,
					LeaseId:   "",
				},
			},
			wantErr: true,
		},
		{
			name: "invalid quota request - negative resource amount",
			msg: &Message{
				QuotaReq: &millv1.QuotaRequest{
					RequestId: "req-1",
					Operation: millv1.QuotaOperation_QUOTA_OPERATION_RESERVE,
					LeaseId:   "lease-1",
					Resources: map[string]int64{"cache_storage_bytes": -1},
				},
			},
			wantErr: true,
		},
		{
			name: "valid quota transition request with reservation_id",
			msg: &Message{
				QuotaReq: &millv1.QuotaRequest{
					RequestId:     "req-1",
					Operation:     millv1.QuotaOperation_QUOTA_OPERATION_BEGIN_COMMIT,
					LeaseId:       "lease-1",
					ReservationId: "res-1",
				},
			},
			wantErr: false,
		},
		{
			name: "valid reserve result reporting workflow resources",
			msg: &Message{
				ReserveResult: &millv1.ReserveResult{
					LeaseId:  "lease-1",
					Accepted: true,
					QuotaResources: map[string]int64{
						"workflows":  1,
						"memory_mib": 2048,
						"vcpus":      2,
						"disk_mib":   8192,
					},
				},
			},
			wantErr: false,
		},
		{
			name: "valid reserve result reporting custom resources",
			msg: &Message{
				ReserveResult: &millv1.ReserveResult{
					LeaseId:  "lease-1",
					Accepted: true,
					QuotaResources: map[string]int64{
						"gpu_count": 4,
					},
				},
			},
			wantErr: false,
		},
		{
			name: "invalid reserve result - negative resource amount",
			msg: &Message{
				ReserveResult: &millv1.ReserveResult{
					LeaseId:        "lease-1",
					Accepted:       true,
					QuotaResources: map[string]int64{"memory_mib": -1},
				},
			},
			wantErr: true,
		},
		{
			name: "valid quota response with resource field shape",
			msg: &Message{
				QuotaResp: &millv1.QuotaResponse{
					RequestId:     "req-1",
					ReservationId: "res-1",
					Allowed:       true,
					Resource:      "cache_storage_bytes",
				},
			},
			wantErr: false,
		},
		{
			name: "invalid quota response - empty request_id",
			msg: &Message{
				QuotaResp: &millv1.QuotaResponse{
					RequestId: "",
				},
			},
			wantErr: true,
		},
	}

	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			err := validator.Validate(tc.msg)
			if (err != nil) != tc.wantErr {
				t.Fatalf("Validate() error = %v, wantErr = %v", err, tc.wantErr)
			}
		})
	}
}

func TestQuotaResourceMapBounds(t *testing.T) {
	tooMany := make(map[string]int64)
	for i := range quota.MaxResourcePairs + 1 {
		tooMany[strings.Repeat("x", i+1)] = 1
	}
	for _, resources := range []map[string]int64{
		{"": 1},
		{strings.Repeat("é", 33): 1},
		{"large": quota.MaxResourceAmount + 1},
		tooMany,
	} {
		msg := &Message{QuotaReq: &millv1.QuotaRequest{
			RequestId: "req-1",
			Operation: millv1.QuotaOperation_QUOTA_OPERATION_RESERVE,
			LeaseId:   "lease-1",
			Kind:      "custom",
			Key:       "key",
			Resources: resources,
		}}
		if err := validator.Validate(msg); err == nil {
			t.Fatalf("expected invalid resources: %v", resources)
		}
	}
}