package xturbo

import (
	"encoding/binary"
	"io"
	"os"
	"sync"
	"sync/atomic"
	"testing"
)

// scriptExt is an extension the scripted server advertises.
type scriptExt struct{ major, firstEvent, firstError uint8 }

// script is the behaviour of the scripted X server behind a socketpair. It
// answers QueryExtension, InternAtom, GetInputFocus, BIG-REQUESTS Enable and
// XC-MISC GetXIDRange like an X server; other requests get no reply unless
// hook handles them.
type script struct {
	exts      map[string]scriptExt
	bigreqMax uint32            // Enable reply maximum_request_length
	atoms     map[string]uint32 // known atoms; new names get nextAtom++
	nextAtom  uint32
	badAtoms  map[string]bool // InternAtom answers BadAlloc
	xidRanges [][2]uint32     // successive GetXIDRange replies {start, count}
	// hook sees every request first; returning true skips the default.
	hook func(s *scriptConn, major, minor uint8, seq uint16, body []byte) bool
	// silent, once set, makes the server read requests and never answer
	// (alive but hung).
	silent atomic.Bool
}

// scriptConn is the server side of one scripted connection.
type scriptConn struct {
	f   *os.File
	sc  *script
	seq uint16

	mu   sync.Mutex
	reqs []scriptReq // every request after setup, in order
}

type scriptReq struct {
	seq   uint16
	major uint8
	minor uint8
	raw   []byte // whole request, header included
}

func (s *scriptConn) requests() []scriptReq {
	s.mu.Lock()
	defer s.mu.Unlock()
	return append([]scriptReq(nil), s.reqs...)
}

func (s *scriptConn) write(p []byte) { _, _ = s.f.Write(p) }

func (s *scriptConn) reply(data uint8, body []byte) {
	words := (max(len(body), 24) - 24 + 3) / 4
	p := make([]byte, 32+4*words)
	p[0], p[1] = 1, data
	binary.LittleEndian.PutUint16(p[2:], s.seq)
	binary.LittleEndian.PutUint32(p[4:], uint32(words))
	copy(p[8:], body)
	s.write(p)
}

func (s *scriptConn) xerror(code, major uint8, bad uint32) {
	p := make([]byte, 32)
	p[1] = code
	binary.LittleEndian.PutUint16(p[2:], s.seq)
	binary.LittleEndian.PutUint32(p[4:], bad)
	p[10] = major
	s.write(p)
}

// event writes a 32-byte core event carrying the current sequence number.
func (s *scriptConn) event(code uint8, fill byte) {
	p := make([]byte, 32)
	for i := range p {
		p[i] = fill
	}
	p[0] = code
	binary.LittleEndian.PutUint16(p[2:], s.seq)
	s.write(p)
}

func (sc *script) extByMajor(major uint8) string {
	for name, e := range sc.exts {
		if e.major == major {
			return name
		}
	}
	return ""
}

// serve reads requests until EOF and answers them.
func (s *scriptConn) serve() {
	sc := s.sc
	for {
		var hdr [4]byte
		if _, err := io.ReadFull(s.f, hdr[:]); err != nil {
			return
		}
		raw := hdr[:]
		n := int(binary.LittleEndian.Uint16(hdr[2:])) * 4
		if n == 0 { // BIG-REQUESTS header
			var l [4]byte
			if _, err := io.ReadFull(s.f, l[:]); err != nil {
				return
			}
			raw = append(raw, l[:]...)
			n = int(binary.LittleEndian.Uint32(l[:])) * 4
		}
		rest := make([]byte, n-len(raw))
		if _, err := io.ReadFull(s.f, rest); err != nil {
			return
		}
		raw = append(append([]byte(nil), raw...), rest...)
		body := raw[len(raw)-len(rest):]
		s.seq++
		major, minor := hdr[0], hdr[1]
		s.mu.Lock()
		s.reqs = append(s.reqs, scriptReq{s.seq, major, minor, raw})
		s.mu.Unlock()
		if sc.silent.Load() || sc.hook != nil && sc.hook(s, major, minor, s.seq, body) {
			continue
		}
		s.answer(major, minor, body)
	}
}

// answer runs the default behaviour for one request, holding s.mu so tests
// can read sc (atom) concurrently.
func (s *scriptConn) answer(major, minor uint8, body []byte) {
	sc := s.sc
	s.mu.Lock()
	defer s.mu.Unlock()
	switch {
	case major == opQueryExtension:
		e, ok := sc.exts[string(body[4:4+binary.LittleEndian.Uint16(body)])]
		var present uint8
		if ok {
			present = 1
		}
		s.reply(0, []byte{present, e.major, e.firstEvent, e.firstError})
	case major == opInternAtom:
		name := string(body[4 : 4+binary.LittleEndian.Uint16(body)])
		if sc.badAtoms[name] {
			s.xerror(11, opInternAtom, 0) // BadAlloc
			return
		}
		a, ok := sc.atoms[name]
		if !ok && minor == 0 {
			if sc.atoms == nil {
				sc.atoms = map[string]uint32{}
			}
			sc.nextAtom++
			a = sc.nextAtom + 1000
			sc.atoms[name] = a
		}
		s.reply(0, binary.LittleEndian.AppendUint32(nil, a))
	case major == opGetInputFocus:
		s.reply(1, binary.LittleEndian.AppendUint32(nil, 0x1234))
	case major >= 128 && sc.extByMajor(major) == "BIG-REQUESTS" && minor == 0:
		s.reply(0, binary.LittleEndian.AppendUint32(nil, sc.bigreqMax))
	case major >= 128 && sc.extByMajor(major) == "XC-MISC" && minor == minorGetXIDRange:
		r := [2]uint32{0, 1}
		if len(sc.xidRanges) > 0 {
			r, sc.xidRanges = sc.xidRanges[0], sc.xidRanges[1:]
		}
		s.reply(0, binary.LittleEndian.AppendUint32(binary.LittleEndian.AppendUint32(nil, r[0]), r[1]))
	}
}

// fullExts is what a typical Xvfb advertises, at its usual opcodes.
func fullExts() map[string]scriptExt {
	return map[string]scriptExt{
		"BIG-REQUESTS":            {133, 0, 0},
		"Generic Event Extension": {128, 0, 0},
		"XC-MISC":                 {136, 0, 0},
		"Composite":               {142, 0, 0},
		"XFIXES":                  {138, 87, 140},
		"X-Resource":              {144, 0, 0},
		"RENDER":                  {139, 0, 141},
		"SHAPE":                   {129, 64, 0},
	}
}

// startScript serves sc behind a socketpair after answering the setup with
// the recorded Success reply, and returns the client fd (not yet set up) and
// the server connection.
func startScript(t *testing.T, sc *script) (int, *scriptConn) {
	t.Helper()
	fd, srv, _ := serveSetup(t, loadSetup(t, "Success"), 0, false, sc)
	return fd, srv
}

// dialScript runs NewConn against sc.
func dialScript(t *testing.T, sc *script) (*Conn, *scriptConn) {
	t.Helper()
	fd, srv := startScript(t, sc)
	c, err := NewConn(fd)
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = c.Close() })
	return c, srv
}

// atom returns the atom the server assigned to name.
func (s *scriptConn) atom(name string) uint32 {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.sc.atoms[name]
}