package xturbo

import (
	"encoding/binary"
	"fmt"
	"io"
	"os"
	"time"

	"golang.org/x/sys/unix"
)

// Fixed layout of the xproto Setup reply (offsets from its status byte).
const (
	setupHeaderSize = 8  // status, pad/reason_len, protocol major, minor, length
	setupFixedSize  = 40 // fixed part of a successful Setup
	screenFixedSize = 40 // fixed part of a SCREEN
	formatSize      = 8  // FORMAT
	depthFixedSize  = 8  // fixed part of a DEPTH
	visualSize      = 24 // VISUALTYPE
	protocolMajor   = 11
	protocolMinor   = 0
)

// setupInfo is what root keeps from a successful setup reply. Everything else
// is parsed by callers from Conn.SetupBytes.
type setupInfo struct {
	xidBase, xidMask uint32
	root             uint32
}

func malformed(what string) error { return fmt.Errorf("%w: %s", ErrMalformed, what) }

// encodeSetupRequest builds the little-endian connection setup request.
func encodeSetupRequest(authName string, authData []byte) []byte {
	n, d := pad4(len(authName)), pad4(len(authData))
	p := make([]byte, 12+n+d)
	p[0] = 'l'
	binary.LittleEndian.PutUint16(p[2:], protocolMajor)
	binary.LittleEndian.PutUint16(p[4:], protocolMinor)
	binary.LittleEndian.PutUint16(p[6:], uint16(len(authName)))
	binary.LittleEndian.PutUint16(p[8:], uint16(len(authData)))
	copy(p[12:], authName)
	copy(p[12+n:], authData)
	return p
}

func pad4(n int) int { return (n + 3) &^ 3 }

// handshake runs the setup exchange on the non-blocking fd, waiting with poll
// until deadline, and returns the full setup reply. A refusal is a
// *SetupError; a broken reply wraps ErrMalformed; a silent server an error
// wrapping os.ErrDeadlineExceeded.
func handshake(fd int, authName string, authData []byte, deadline time.Time) ([]byte, setupInfo, error) {
	if len(authName) > 0xffff || len(authData) > 0xffff {
		return nil, setupInfo{}, ErrBadLength
	}
	if err := writeFull(fd, encodeSetupRequest(authName, authData), deadline); err != nil {
		return nil, setupInfo{}, err
	}
	var hdr [setupHeaderSize]byte
	if err := readFull(fd, hdr[:], deadline); err != nil {
		return nil, setupInfo{}, err
	}
	reply := make([]byte, setupHeaderSize+4*int(binary.LittleEndian.Uint16(hdr[6:])))
	copy(reply, hdr[:])
	if err := readFull(fd, reply[setupHeaderSize:], deadline); err != nil {
		return nil, setupInfo{}, err
	}
	info, err := parseSetup(reply)
	if err != nil {
		return nil, setupInfo{}, err
	}
	return reply, info, nil
}

// parseSetup validates a complete setup reply (header length already
// honoured) and extracts what root needs.
func parseSetup(p []byte) (setupInfo, error) {
	if len(p) < setupHeaderSize {
		return setupInfo{}, malformed("short setup reply")
	}
	switch p[0] {
	case setupStatusFailed:
		// SetupFailed: reason_len at byte 1, reason after the header.
		n := int(p[1])
		if setupHeaderSize+n > len(p) {
			return setupInfo{}, malformed("setup failure reason overflows reply")
		}
		return setupInfo{}, &SetupError{Status: setupStatusFailed, Reason: string(p[setupHeaderSize : setupHeaderSize+n])}
	case setupStatusAuthenticate:
		// SetupAuthenticate: reason fills length*4 bytes, NUL padded.
		reason := p[setupHeaderSize:]
		for len(reason) > 0 && reason[len(reason)-1] == 0 {
			reason = reason[:len(reason)-1]
		}
		return setupInfo{}, &SetupError{Status: setupStatusAuthenticate, Reason: string(reason)}
	case setupStatusSuccess:
	default:
		return setupInfo{}, malformed(fmt.Sprintf("setup status %d", p[0]))
	}
	if len(p) < setupFixedSize {
		return setupInfo{}, malformed("short setup reply")
	}
	if binary.LittleEndian.Uint16(p[2:]) != protocolMajor {
		return setupInfo{}, malformed(fmt.Sprintf("protocol major %d", binary.LittleEndian.Uint16(p[2:])))
	}
	info := setupInfo{
		xidBase: binary.LittleEndian.Uint32(p[12:]),
		xidMask: binary.LittleEndian.Uint32(p[16:]),
	}
	if info.xidMask == 0 || info.xidBase&info.xidMask != 0 {
		return setupInfo{}, malformed("resource id base/mask")
	}
	if p[28] == 0 {
		return setupInfo{}, malformed("setup has no screens")
	}
	off := setupFixedSize + pad4(int(binary.LittleEndian.Uint16(p[24:]))) + formatSize*int(p[29])
	if off > len(p) {
		return setupInfo{}, malformed("vendor or pixmap formats overflow setup reply")
	}
	roots := int(p[28])
	for i := range roots {
		if len(p)-off < screenFixedSize {
			return setupInfo{}, malformed(fmt.Sprintf("screen %d overflows setup reply", i))
		}
		if i == 0 {
			info.root = binary.LittleEndian.Uint32(p[off:])
		}
		depths := int(p[off+39])
		off += screenFixedSize
		for range depths {
			if len(p)-off < depthFixedSize {
				return setupInfo{}, malformed(fmt.Sprintf("screen %d depth overflows setup reply", i))
			}
			visuals := visualSize * int(binary.LittleEndian.Uint16(p[off+2:]))
			off += depthFixedSize
			if len(p)-off < visuals {
				return setupInfo{}, malformed(fmt.Sprintf("screen %d visuals overflow setup reply", i))
			}
			off += visuals
		}
	}
	return info, nil
}

// waitFd polls fd for events until deadline.
func waitFd(fd int, events int16, deadline time.Time) error {
	left := time.Until(deadline)
	if left <= 0 {
		return deadlineError("connection setup")
	}
	return pollFd(fd, events, left)
}

func writeFull(fd int, p []byte, deadline time.Time) error {
	for len(p) > 0 {
		n, err := unix.Write(fd, p)
		switch {
		case err == unix.EINTR:
			continue
		case err == unix.EAGAIN:
			if err := waitFd(fd, unix.POLLOUT, deadline); err != nil {
				return err
			}
			continue
		case err == unix.EPIPE || err == unix.ECONNRESET:
			return ErrClosed
		case err != nil:
			return os.NewSyscallError("write", err)
		case n <= 0:
			return io.ErrShortWrite
		}
		p = p[n:]
	}
	return nil
}

func readFull(fd int, p []byte, deadline time.Time) error {
	for len(p) > 0 {
		n, err := unix.Read(fd, p)
		switch {
		case err == unix.EINTR:
			continue
		case err == unix.EAGAIN:
			if err := waitFd(fd, unix.POLLIN, deadline); err != nil {
				return err
			}
			continue
		case err == unix.ECONNRESET:
			return ErrClosed
		case err != nil:
			return os.NewSyscallError("read", err)
		case n == 0:
			return fmt.Errorf("%w: truncated setup reply: %w", ErrMalformed, io.ErrUnexpectedEOF)
		}
		p = p[n:]
	}
	return nil
}