package wire

import (
	"bytes"
	"context"
	"testing"

	"github.com/vmihailenco/msgpack/v5"
	"nhooyr.io/websocket"
)

type testWriter struct {
	msgs [][]byte
}

func (w *testWriter) Write(ctx context.Context, messageType websocket.MessageType, p []byte) error {
	w.msgs = append(w.msgs, append([]byte(nil), p...))
	return nil
}

func TestHelloMarshaling(t *testing.T) {
	hello := Hello{
		Type:     "hello",
		ID:       "test-id",
		HWID:     "test-hwid",
		Host:     "test-host",
		OS:       "test-os",
		Arch:     "test-arch",
		Version:  "1.0.0",
		User:     "test-user",
		Monitors: 2,
		MonitorInfo: []MonitorInfo{
			{Width: 1920, Height: 1080},
			{Width: 1280, Height: 720},
		},
		Country: "US",
	}

	data, err := msgpack.Marshal(hello)
	if err != nil {
		t.Fatalf("Failed to marshal Hello: %v", err)
	}

	var decoded Hello
	if err := msgpack.Unmarshal(data, &decoded); err != nil {
		t.Fatalf("Failed to unmarshal Hello: %v", err)
	}

	if decoded.ID != hello.ID {
		t.Errorf("ID mismatch: got %s, want %s", decoded.ID, hello.ID)
	}
	if decoded.Monitors != hello.Monitors {
		t.Errorf("Monitors mismatch: got %d, want %d", decoded.Monitors, hello.Monitors)
	}
	if len(decoded.MonitorInfo) != len(hello.MonitorInfo) {
		t.Errorf("MonitorInfo length mismatch: got %d, want %d", len(decoded.MonitorInfo), len(hello.MonitorInfo))
	}
}

func TestHelloBuildTagMarshaling(t *testing.T) {
	t.Run("with build tag", func(t *testing.T) {
		hello := Hello{
			Type:     "hello",
			ID:       "test-id",
			HWID:     "test-hwid",
			Host:     "test-host",
			OS:       "test-os",
			Arch:     "test-arch",
			Version:  "1.0.0",
			User:     "test-user",
			BuildTag: "abc-123-def-456",
		}

		data, err := msgpack.Marshal(hello)
		if err != nil {
			t.Fatalf("Failed to marshal Hello with BuildTag: %v", err)
		}

		var decoded Hello
		if err := msgpack.Unmarshal(data, &decoded); err != nil {
			t.Fatalf("Failed to unmarshal Hello with BuildTag: %v", err)
		}

		if decoded.BuildTag != "abc-123-def-456" {
			t.Errorf("BuildTag mismatch: got %q, want %q", decoded.BuildTag, "abc-123-def-456")
		}
		if decoded.ID != hello.ID {
			t.Errorf("ID mismatch: got %s, want %s", decoded.ID, hello.ID)
		}
	})

	t.Run("without build tag (omitempty)", func(t *testing.T) {
		hello := Hello{
			Type:    "hello",
			ID:      "test-id",
			HWID:    "test-hwid",
			Host:    "test-host",
			OS:      "test-os",
			Arch:    "test-arch",
			Version: "1.0.0",
			User:    "test-user",
		}

		data, err := msgpack.Marshal(hello)
		if err != nil {
			t.Fatalf("Failed to marshal Hello without BuildTag: %v", err)
		}

		// Verify buildTag key is not present in the wire data (omitempty)
		if bytes.Contains(data, []byte("buildTag")) {
			t.Error("Expected buildTag to be omitted from wire data when empty")
		}

		var decoded Hello
		if err := msgpack.Unmarshal(data, &decoded); err != nil {
			t.Fatalf("Failed to unmarshal Hello without BuildTag: %v", err)
		}

		if decoded.BuildTag != "" {
			t.Errorf("BuildTag should be empty, got %q", decoded.BuildTag)
		}
	})

	t.Run("server-side decode without build tag field", func(t *testing.T) {
		// Simulate an older client that doesn't send buildTag at all
		type OldHello struct {
			Type    string `msgpack:"type"`
			ID      string `msgpack:"id"`
			HWID    string `msgpack:"hwid"`
			Host    string `msgpack:"host"`
			OS      string `msgpack:"os"`
			Arch    string `msgpack:"arch"`
			Version string `msgpack:"version"`
			User    string `msgpack:"user"`
		}

		old := OldHello{
			Type:    "hello",
			ID:      "old-client",
			HWID:    "old-hwid",
			Host:    "old-host",
			OS:      "windows",
			Arch:    "amd64",
			Version: "0.9.0",
			User:    "old-user",
		}

		data, err := msgpack.Marshal(old)
		if err != nil {
			t.Fatalf("Failed to marshal OldHello: %v", err)
		}

		var decoded Hello
		if err := msgpack.Unmarshal(data, &decoded); err != nil {
			t.Fatalf("Failed to unmarshal OldHello into Hello: %v", err)
		}

		if decoded.BuildTag != "" {
			t.Errorf("BuildTag should default to empty for old client, got %q", decoded.BuildTag)
		}
		if decoded.ID != "old-client" {
			t.Errorf("ID mismatch: got %s, want old-client", decoded.ID)
		}
	})
}

func TestPingPongMarshaling(t *testing.T) {
	ping := Ping{
		Type: "ping",
		TS:   1234567890,
	}

	data, err := msgpack.Marshal(ping)
	if err != nil {
		t.Fatalf("Failed to marshal Ping: %v", err)
	}

	var decoded Ping
	if err := msgpack.Unmarshal(data, &decoded); err != nil {
		t.Fatalf("Failed to unmarshal Ping: %v", err)
	}

	if decoded.TS != ping.TS {
		t.Errorf("Timestamp mismatch: got %d, want %d", decoded.TS, ping.TS)
	}

	pong := Pong{
		Type: "pong",
		TS:   ping.TS,
	}

	data2, err := msgpack.Marshal(pong)
	if err != nil {
		t.Fatalf("Failed to marshal Pong: %v", err)
	}

	var decodedPong Pong
	if err := msgpack.Unmarshal(data2, &decodedPong); err != nil {
		t.Fatalf("Failed to unmarshal Pong: %v", err)
	}

	if decodedPong.TS != pong.TS {
		t.Errorf("Pong timestamp mismatch: got %d, want %d", decodedPong.TS, pong.TS)
	}
}

func TestCommandMarshaling(t *testing.T) {
	cmd := Command{
		Type:           "command",
		CommandType:    "process_list",
		CommandVersion: 1,
		ID:             "cmd-123",
		Payload:        map[string]interface{}{"key": "value"},
	}

	data, err := msgpack.Marshal(cmd)
	if err != nil {
		t.Fatalf("Failed to marshal Command: %v", err)
	}

	var decoded Command
	if err := msgpack.Unmarshal(data, &decoded); err != nil {
		t.Fatalf("Failed to unmarshal Command: %v", err)
	}

	if decoded.CommandType != cmd.CommandType {
		t.Errorf("CommandType mismatch: got %s, want %s", decoded.CommandType, cmd.CommandType)
	}
	if decoded.ID != cmd.ID {
		t.Errorf("ID mismatch: got %s, want %s", decoded.ID, cmd.ID)
	}
	if decoded.CommandVersion != cmd.CommandVersion {
		t.Errorf("CommandVersion mismatch: got %d, want %d", decoded.CommandVersion, cmd.CommandVersion)
	}
}

func TestCommandResultMarshaling(t *testing.T) {
	result := CommandResult{
		Type:      "command_result",
		CommandID: "cmd-123",
		OK:        true,
		Message:   "Success",
	}

	data, err := msgpack.Marshal(result)
	if err != nil {
		t.Fatalf("Failed to marshal CommandResult: %v", err)
	}

	var decoded CommandResult
	if err := msgpack.Unmarshal(data, &decoded); err != nil {
		t.Fatalf("Failed to unmarshal CommandResult: %v", err)
	}

	if decoded.OK != result.OK {
		t.Errorf("OK mismatch: got %v, want %v", decoded.OK, result.OK)
	}
	if decoded.Message != result.Message {
		t.Errorf("Message mismatch: got %s, want %s", decoded.Message, result.Message)
	}
}

func TestFrameMarshaling(t *testing.T) {
	frame := Frame{
		Type: "frame",
		Header: FrameHeader{
			Monitor: 0,
			FPS:     30,
			Format:  "jpeg",
		},
		Data: []byte{0xFF, 0xD8, 0xFF, 0xE0},
	}

	data, err := msgpack.Marshal(frame)
	if err != nil {
		t.Fatalf("Failed to marshal Frame: %v", err)
	}

	var decoded Frame
	if err := msgpack.Unmarshal(data, &decoded); err != nil {
		t.Fatalf("Failed to unmarshal Frame: %v", err)
	}

	if decoded.Header.FPS != frame.Header.FPS {
		t.Errorf("FPS mismatch: got %d, want %d", decoded.Header.FPS, frame.Header.FPS)
	}
	if !bytes.Equal(decoded.Data, frame.Data) {
		t.Errorf("Data mismatch")
	}
}

func TestFileEntryMarshaling(t *testing.T) {
	entry := FileEntry{
		Name:    "test.txt",
		Path:    "/home/user/test.txt",
		IsDir:   false,
		Size:    1024,
		ModTime: 1609459200,
	}

	data, err := msgpack.Marshal(entry)
	if err != nil {
		t.Fatalf("Failed to marshal FileEntry: %v", err)
	}

	var decoded FileEntry
	if err := msgpack.Unmarshal(data, &decoded); err != nil {
		t.Fatalf("Failed to unmarshal FileEntry: %v", err)
	}

	if decoded.Name != entry.Name {
		t.Errorf("Name mismatch: got %s, want %s", decoded.Name, entry.Name)
	}
	if decoded.Size != entry.Size {
		t.Errorf("Size mismatch: got %d, want %d", decoded.Size, entry.Size)
	}
	if decoded.IsDir != entry.IsDir {
		t.Errorf("IsDir mismatch: got %v, want %v", decoded.IsDir, entry.IsDir)
	}
}

func TestFileListResultMarshaling(t *testing.T) {
	result := FileListResult{
		Type:      "file_list_result",
		CommandID: "cmd-456",
		Path:      "/home/user",
		Entries: []FileEntry{
			{Name: "file1.txt", Path: "/home/user/file1.txt", IsDir: false, Size: 100, ModTime: 1609459200},
			{Name: "dir1", Path: "/home/user/dir1", IsDir: true, Size: 0, ModTime: 1609459200},
		},
		Error: "",
	}

	data, err := msgpack.Marshal(result)
	if err != nil {
		t.Fatalf("Failed to marshal FileListResult: %v", err)
	}

	var decoded FileListResult
	if err := msgpack.Unmarshal(data, &decoded); err != nil {
		t.Fatalf("Failed to unmarshal FileListResult: %v", err)
	}

	if len(decoded.Entries) != len(result.Entries) {
		t.Errorf("Entries count mismatch: got %d, want %d", len(decoded.Entries), len(result.Entries))
	}
	if decoded.Path != result.Path {
		t.Errorf("Path mismatch: got %s, want %s", decoded.Path, result.Path)
	}
}

func TestProcessInfoMarshaling(t *testing.T) {
	proc := ProcessInfo{
		PID:      1234,
		PPID:     1,
		Name:     "test.exe",
		CPU:      25.5,
		Memory:   1048576,
		Username: "testuser",
		Type:     "own",
	}

	data, err := msgpack.Marshal(proc)
	if err != nil {
		t.Fatalf("Failed to marshal ProcessInfo: %v", err)
	}

	var decoded ProcessInfo
	if err := msgpack.Unmarshal(data, &decoded); err != nil {
		t.Fatalf("Failed to unmarshal ProcessInfo: %v", err)
	}

	if decoded.PID != proc.PID {
		t.Errorf("PID mismatch: got %d, want %d", decoded.PID, proc.PID)
	}
	if decoded.PPID != proc.PPID {
		t.Errorf("PPID mismatch: got %d, want %d", decoded.PPID, proc.PPID)
	}
	if decoded.Name != proc.Name {
		t.Errorf("Name mismatch: got %s, want %s", decoded.Name, proc.Name)
	}
	if decoded.CPU != proc.CPU {
		t.Errorf("CPU mismatch: got %f, want %f", decoded.CPU, proc.CPU)
	}
	if decoded.Type != proc.Type {
		t.Errorf("Type mismatch: got %s, want %s", decoded.Type, proc.Type)
	}
}

func TestProcessListResultMarshaling(t *testing.T) {
	result := ProcessListResult{
		Type:      "process_list_result",
		CommandID: "cmd-789",
		Processes: []ProcessInfo{
			{PID: 1, PPID: 0, Name: "init", CPU: 0.1, Memory: 1024, Username: "root", Type: "system"},
			{PID: 100, PPID: 1, Name: "service", CPU: 5.0, Memory: 2048, Username: "system", Type: "service"},
		},
		Error: "",
	}

	data, err := msgpack.Marshal(result)
	if err != nil {
		t.Fatalf("Failed to marshal ProcessListResult: %v", err)
	}

	var decoded ProcessListResult
	if err := msgpack.Unmarshal(data, &decoded); err != nil {
		t.Fatalf("Failed to unmarshal ProcessListResult: %v", err)
	}

	if len(decoded.Processes) != len(result.Processes) {
		t.Errorf("Processes count mismatch: got %d, want %d", len(decoded.Processes), len(result.Processes))
	}
}

func TestWriteMsg(t *testing.T) {
	writer := &testWriter{}
	ctx := context.Background()

	msg := Pong{
		Type: "pong",
		TS:   1234567890,
	}

	if err := WriteMsg(ctx, writer, msg); err != nil {
		t.Fatalf("WriteMsg failed: %v", err)
	}

	if len(writer.msgs) != 1 {
		t.Fatalf("Expected 1 message, got %d", len(writer.msgs))
	}

	var decoded Pong
	if err := msgpack.Unmarshal(writer.msgs[0], &decoded); err != nil {
		t.Fatalf("Failed to unmarshal written message: %v", err)
	}

	if decoded.TS != msg.TS {
		t.Errorf("Written message timestamp mismatch: got %d, want %d", decoded.TS, msg.TS)
	}
}

func TestConsoleOutputMarshaling(t *testing.T) {
	exitCode := 0
	output := ConsoleOutput{
		Type:      "console_output",
		SessionID: "session-123",
		Data:      []byte("test output"),
		ExitCode:  &exitCode,
		Error:     "",
	}

	data, err := msgpack.Marshal(output)
	if err != nil {
		t.Fatalf("Failed to marshal ConsoleOutput: %v", err)
	}

	var decoded ConsoleOutput
	if err := msgpack.Unmarshal(data, &decoded); err != nil {
		t.Fatalf("Failed to unmarshal ConsoleOutput: %v", err)
	}

	if decoded.SessionID != output.SessionID {
		t.Errorf("SessionID mismatch: got %s, want %s", decoded.SessionID, output.SessionID)
	}
	if !bytes.Equal(decoded.Data, output.Data) {
		t.Errorf("Data mismatch")
	}
	if decoded.ExitCode == nil || *decoded.ExitCode != *output.ExitCode {
		t.Errorf("ExitCode mismatch")
	}
}