package xturbo

import (
	"bytes"
	"encoding/binary"
	"errors"
	"io"
	"os"
	"path/filepath"
	"testing"

	"bnema.dev/xturbo/internal/xtest"
	"golang.org/x/sys/unix"
)

func loadSetup(t testing.TB, name string) []byte {
	t.Helper()
	b, err := xtest.LoadHex(filepath.Join("testdata", "setup", name+".setup.hex"))
	if err != nil {
		t.Fatal(err)
	}
	return b
}

// scriptedServer runs the server side of a socketpair: it reads one setup
// request, reports it on the returned channel, writes reply in chunks of
// chunk bytes (0 = all at once) and, if hangup, closes its end. Otherwise it
// then answers requests as a server without extensions (script{}).
func scriptedServer(t *testing.T, reply []byte, chunk int, hangup bool) (client int, req <-chan []byte) {
	t.Helper()
	fd, _, ch := serveSetup(t, reply, chunk, hangup, &script{})
	return fd, ch
}

// serveSetup is scriptedServer serving sc after the setup.
func serveSetup(t *testing.T, reply []byte, chunk int, hangup bool, sc *script) (int, *scriptConn, <-chan []byte) {
	t.Helper()
	fds, err := unix.Socketpair(unix.AF_UNIX, unix.SOCK_STREAM|unix.SOCK_CLOEXEC, 0)
	if err != nil {
		t.Fatal(err)
	}
	// A blocking fd: NewConn switches it to O_NONBLOCK itself.
	if err := unix.SetNonblock(fds[0], false); err != nil {
		t.Fatal(err)
	}
	server := os.NewFile(uintptr(fds[1]), "server")
	srv := &scriptConn{f: server, sc: sc}
	ch := make(chan []byte, 1)
	done := make(chan struct{})
	t.Cleanup(func() { <-done; server.Close() })
	go func() {
		defer close(done)
		var hdr [12]byte
		if _, err := io.ReadFull(server, hdr[:]); err != nil {
			ch <- nil
			return
		}
		n := pad4(int(binary.LittleEndian.Uint16(hdr[6:]))) + pad4(int(binary.LittleEndian.Uint16(hdr[8:])))
		rest := make([]byte, n)
		if _, err := io.ReadFull(server, rest); err != nil {
			ch <- nil
			return
		}
		ch <- append(hdr[:], rest...)
		for p := reply; len(p) > 0; {
			k := len(p)
			if chunk > 0 {
				k = min(k, chunk)
			}
			if _, err := server.Write(p[:k]); err != nil {
				return
			}
			p = p[k:]
		}
		if hangup {
			server.Close()
			return
		}
		srv.serve()
	}()
	return fds[0], srv, ch
}

func TestNewConnSuccess(t *testing.T) {
	setup := loadSetup(t, "Success")
	fd, req := scriptedServer(t, setup, 7, false)
	c, err := NewConn(fd)
	if err != nil {
		t.Fatal(err)
	}
	defer c.Close()
	got := <-req
	want := []byte{'l', 0, 11, 0, 0, 0, 0, 0, 0, 0, 0, 0}
	if !bytes.Equal(got, want) {
		t.Fatalf("setup request %x, want %x", got, want)
	}
	if !bytes.Equal(c.SetupBytes(), setup) {
		t.Fatal("SetupBytes differs from the reply")
	}
	// Expected values are read from the fixture header comment
	// (Success.setup.hex), i.e. what that recorded Xvfb assigned.
	if c.Root() != 0x39f || c.xidBase != 0x00200000 || c.xidMask != 0x001fffff {
		t.Fatalf("root %#x base %#x mask %#x", c.Root(), c.xidBase, c.xidMask)
	}
	// The bootstrap issued one QueryExtension per extension and a fence.
	if c.Fd() != fd || c.Buf().Ext != c.Ext() || c.Buf().Next() != Seq(extCount)+1 {
		t.Fatal("accessors", c.Buf().Next())
	}
	if c.Ext().Major(ExtXFixes) != 0 || !c.Ext().Info(ExtCore).Present {
		t.Fatal("ExtTable not empty")
	}
	flags, err := unix.FcntlInt(uintptr(fd), unix.F_GETFL, 0)
	if err != nil || flags&unix.O_NONBLOCK == 0 {
		t.Fatalf("fd not O_NONBLOCK after setup: %v", err)
	}
}

func TestNewConnAuthRequest(t *testing.T) {
	fd, req := scriptedServer(t, loadSetup(t, "Success"), 0, false)
	cookie := []byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}
	c, err := newConn(fd, authMagicCookie, cookie)
	if err != nil {
		t.Fatal(err)
	}
	defer c.Close()
	want := append([]byte{'l', 0, 11, 0, 0, 0, 18, 0, 16, 0, 0, 0}, authMagicCookie+"\x00\x00"...)
	want = append(want, cookie...)
	if got := <-req; !bytes.Equal(got, want) {
		t.Fatalf("setup request %x, want %x", got, want)
	}
}

func TestNewConnRefused(t *testing.T) {
	authenticate := []byte{2, 0, 0, 0, 0, 0, 2, 0, 'n', 'o', 'p', 'e', '!', 0, 0, 0}
	for _, tc := range []struct {
		name   string
		reply  []byte
		status uint8
		reason string
	}{
		{"failed", loadSetup(t, "Failed"), 0, "Authorization required, but no authorization protocol specified\n"},
		{"authenticate", authenticate, 2, "nope!"},
	} {
		t.Run(tc.name, func(t *testing.T) {
			fd, _ := scriptedServer(t, tc.reply, 5, true)
			c, err := NewConn(fd)
			var se *SetupError
			if c != nil || !errors.As(err, &se) || !errors.Is(err, ErrSetupFailed) {
				t.Fatalf("conn %t, err %v", c != nil, err)
			}
			if se.Status != tc.status || se.Reason != tc.reason {
				t.Fatalf("status %d reason %q", se.Status, se.Reason)
			}
			if _, err := unix.FcntlInt(uintptr(fd), unix.F_GETFD, 0); err != unix.EBADF {
				t.Fatalf("fd not closed on failure: %v", err)
			}
		})
	}
}

func TestNewConnMalformed(t *testing.T) {
	setup := loadSetup(t, "Success")
	noScreens := bytes.Clone(setup)
	noScreens[28] = 0
	badMask := bytes.Clone(setup)
	binary.LittleEndian.PutUint32(badMask[16:], 0)
	badVersion := bytes.Clone(setup)
	badVersion[2] = 10
	badStatus := bytes.Clone(setup)
	badStatus[0] = 3
	// Fixed Setup only (length 8 words): one root declared, no SCREEN bytes.
	screenPast := bytes.Clone(setup[:40])
	binary.LittleEndian.PutUint16(screenPast[6:], 8)
	binary.LittleEndian.PutUint16(screenPast[24:], 0) // vendor_len
	screenPast[29] = 0                                // pixmap_formats_len
	// Walk the recorded reply to find the first screen and its first depth
	// that carries visuals; every extent is a multiple of 4 bytes.
	screen := setupFixedSize + pad4(int(binary.LittleEndian.Uint16(setup[24:]))) + formatSize*int(setup[29])
	visualDepth := screen + screenFixedSize
	for binary.LittleEndian.Uint16(setup[visualDepth+2:]) == 0 {
		visualDepth += depthFixedSize
	}
	// Screen header present, its allowed depths cut off.
	missingDepths := truncateSetup(setup, screen+screenFixedSize)
	// A depth header announcing visuals, the visuals cut off.
	missingVisuals := truncateSetup(setup, visualDepth+depthFixedSize)
	// One visual of the depth present, the rest cut off.
	partialVisuals := truncateSetup(setup, visualDepth+depthFixedSize+visualSize)
	// Complete first screen, a second screen declared but absent.
	missingScreen := bytes.Clone(setup)
	missingScreen[28] = 2
	// Header claims 8 bytes of body, failure reason claims 20.
	reasonPast := []byte{0, 20, 11, 0, 0, 0, 2, 0, 'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h'}
	// Success header whose length is shorter than the fixed Setup.
	shortSuccess := []byte{1, 0, 11, 0, 0, 0, 1, 0, 0, 0, 0, 0}
	for _, tc := range []struct {
		name  string
		reply []byte
	}{
		{"empty", nil},
		{"truncated header", setup[:5]},
		{"truncated body", setup[:len(setup)-3]},
		{"no screens", noScreens},
		{"bad xid mask", badMask},
		{"protocol major", badVersion},
		{"unknown status", badStatus},
		{"screen past end", screenPast},
		{"missing depths", missingDepths},
		{"missing visuals", missingVisuals},
		{"partial visuals", partialVisuals},
		{"missing subsequent screen", missingScreen},
		{"reason past end", reasonPast},
		{"short success", shortSuccess},
	} {
		t.Run(tc.name, func(t *testing.T) {
			fd, _ := scriptedServer(t, tc.reply, 3, true)
			c, err := NewConn(fd)
			if c != nil || !errors.Is(err, ErrMalformed) || errors.Is(err, ErrSetupFailed) {
				t.Fatalf("conn %t, err %v", c != nil, err)
			}
		})
	}
}

// truncateSetup keeps the first n bytes (n ≡ 0 mod 4) of a setup reply and
// rewrites its additional-data length to match.
func truncateSetup(setup []byte, n int) []byte {
	p := bytes.Clone(setup[:n])
	binary.LittleEndian.PutUint16(p[6:], uint16((n-setupHeaderSize)/4))
	return p
}

func TestConnFlushClose(t *testing.T) {
	c, srv := dialScript(t, &script{})
	first := c.Buf().Next()
	p, seq := c.Buf().Request(43, 0, 0) // GetInputFocus
	if p == nil || seq != first {
		t.Fatal("request")
	}
	if err := c.Flush(); err != nil {
		t.Fatal(err)
	}
	waitRequests(t, srv, int(seq))
	if got := srv.requests()[seq-1].raw; !bytes.Equal(got, []byte{43, 0, 1, 0}) {
		t.Fatalf("%x", got)
	}
	if err := c.Close(); err != nil {
		t.Fatal(err)
	}
	if c.Fd() != -1 || !errors.Is(c.Close(), ErrClosed) || !errors.Is(c.Flush(), ErrClosed) {
		t.Fatal("after Close")
	}
	if p, _ := c.Buf().Request(43, 0, 0); p != nil {
		t.Fatal("Buf usable after Close")
	}
}

func TestSetupErrorMessage(t *testing.T) {
	for _, tc := range []struct {
		e    SetupError
		want string
	}{
		{SetupError{0, "no"}, "xturbo: connection setup failed: no"},
		{SetupError{2, "why"}, "xturbo: connection setup authenticate: why"},
	} {
		if got := tc.e.Error(); got != tc.want {
			t.Errorf("%q, want %q", got, tc.want)
		}
	}
}