package rdpdr

import (
	"encoding/binary"
	"testing"
)

type captureSender struct {
	data [][]byte
}

func (s *captureSender) SendChannelPDU(_ uint16, data []byte) error {
	s.data = append(s.data, append([]byte(nil), data...))
	return nil
}

func TestClassifyVirtualPathUploadFile(t *testing.T) {
	kind, name := classifyVirtualPath(`Upload\php.ini`)
	if kind != "upload_file" || name != "php.ini" {
		t.Fatalf("Upload file should be classified as upload_file, got kind=%q name=%q", kind, name)
	}
}

func TestHandleCreateUploadFileOpensAsFile(t *testing.T) {
	sender := &captureSender{}
	handler := NewHandler("goclient", "AIOps")
	handler.Attach(sender, 1004)
	handler.AddUpload("php.ini", []byte("ok"))

	encodedPath := encodeUTF16(`Upload\php.ini`)
	payload := make([]byte, 32+len(encodedPath))
	binary.LittleEndian.PutUint32(payload[20:24], fileOpen)
	binary.LittleEndian.PutUint32(payload[24:28], fileDirectoryFile)
	binary.LittleEndian.PutUint32(payload[28:32], uint32(len(encodedPath)))
	copy(payload[32:], encodedPath)

	handler.handleCreate(ioRequest{
		deviceID:     1,
		completionID: 7,
		majorFn:      irpCreate,
		payload:      payload,
	})
	if len(sender.data) == 0 {
		t.Fatal("expected create completion")
	}
	completion := sender.data[len(sender.data)-1]
	if len(completion) < 21 {
		t.Fatalf("completion too short: %d", len(completion))
	}
	if status := binary.LittleEndian.Uint32(completion[12:16]); status != statusSuccess {
		t.Fatalf("expected success, got status=0x%08x", status)
	}
	fileID := binary.LittleEndian.Uint32(completion[16:20])
	file, ok := handler.getOpenFile(fileID)
	if !ok {
		t.Fatalf("open file %d not stored", fileID)
	}
	if file.isDir {
		t.Fatalf("upload file opened as directory: %#v", file)
	}
}

func TestHandleReadEOFReturnsZeroLengthSuccess(t *testing.T) {
	sender := &captureSender{}
	handler := NewHandler("goclient", "AIOps")
	handler.Attach(sender, 1004)
	fileID := handler.storeOpenFile(&openFile{path: `Upload\x.yaml`, name: "x.yaml", data: []byte("abc")})

	payload := make([]byte, 20)
	binary.LittleEndian.PutUint32(payload[0:4], 4096)
	binary.LittleEndian.PutUint64(payload[4:12], 3)
	handler.handleRead(ioRequest{
		deviceID:     1,
		fileID:       fileID,
		completionID: 9,
		payload:      payload,
	})

	completion := sender.data[len(sender.data)-1]
	if status := binary.LittleEndian.Uint32(completion[12:16]); status != statusSuccess {
		t.Fatalf("expected EOF success, got status=0x%08x", status)
	}
	if got := binary.LittleEndian.Uint32(completion[16:20]); got != 0 {
		t.Fatalf("expected zero read length at EOF, got %d", got)
	}
}

func TestHandleDeviceControlReturnsInvalidDeviceRequest(t *testing.T) {
	sender := &captureSender{}
	handler := NewHandler("goclient", "AIOps")
	handler.Attach(sender, 1004)

	payload := make([]byte, 32)
	binary.LittleEndian.PutUint32(payload[8:12], 0x000900a8)
	handler.handleDeviceControl(ioRequest{
		deviceID:     1,
		fileID:       2,
		completionID: 10,
		payload:      payload,
	})

	completion := sender.data[len(sender.data)-1]
	if status := binary.LittleEndian.Uint32(completion[12:16]); status != statusInvalidDeviceReq {
		t.Fatalf("expected invalid device request, got status=0x%08x", status)
	}
	if got := binary.LittleEndian.Uint32(completion[16:20]); got != 0 {
		t.Fatalf("expected empty device control output, got %d", got)
	}
}