package x11probe

import (
	"errors"
	"fmt"
	"io"
	"log"
	"net"
	"os"
	"os/exec"
	"strings"
	"syscall"
	"time"
)

type ConnType int

const (
	Unix ConnType = iota
	TCP
)

// X11Status represents the result of probing an X11 socket.
type X11Status int

const (
	X11Dead X11Status = iota
	X11Live
	X11AuthRequired
	X11TemporarilyUnavailable
)

func ProbeX11Socket(connType ConnType, target string, timeout time.Duration) (X11Status, error) {
	var (
		conn net.Conn
		err  error
	)

	switch connType {
	case Unix:
		conn, err = net.DialTimeout("unix", target, timeout)
		if err != nil {
			if isResourceTemporarilyUnavailable(err) { // Check for specific error
				return X11TemporarilyUnavailable, nil
			}
			return X11Dead, fmt.Errorf("unix dial failed: %w", err)
		}
	case TCP:
		addr := parseTCPAddress(target)
		conn, err = net.DialTimeout("tcp", addr, timeout)
		if err != nil {
			if isResourceTemporarilyUnavailable(err) { // Check for specific error
				return X11TemporarilyUnavailable, nil
			}
			return X11Dead, fmt.Errorf("TCP dial failed: %w", err)
		}
	}

	defer conn.Close()
	log.Printf("Successfully connected to %s", target)

	if err := performHandshake(conn); err != nil {
		return X11Dead, err
	}

	status, err := readResponse(conn, timeout)
	if err != nil {
		return X11Dead, err
	}

	return status, nil
}

func isResourceTemporarilyUnavailable(err error) bool {
	// Match exact error string first
	if strings.Contains(err.Error(), "resource temporarily unavailable") {
		return true
	}
	// Fallback to syscall-level check
	if errors.Is(err, syscall.EAGAIN) {
		return true
	}
	return false
}

func parseTCPAddress(target string) string {
	parts := strings.SplitN(target, ":", 2)

	ipPart := parts[0]
	portPart := "0"

	if len(parts) == 2 && portPart != "" {
		portPart = parts[1]
	}

	return fmt.Sprintf("%s:%s", ipPart, portPart)
}

func performHandshake(conn net.Conn) error {
	// Minimal valid handshake with correct padding
	handshake := []byte{
		'l', 0, // Byte order
		0, 11, 0, 0, // Protocol major/minor
		0, 0, // Auth name length
		0, 0, // Auth data length
		0, 0, 0, 0, // Padding
	}

	if _, err := conn.Write(handshake); err != nil {
		return fmt.Errorf("write failed: %w", err)
	}
	return nil
}

func readResponse(conn net.Conn, timeout time.Duration) (X11Status, error) {
	buf := make([]byte, 1)

	deadline := time.Now().Add(timeout)
	if err := conn.SetReadDeadline(deadline); err != nil {
		return X11Dead, fmt.Errorf("SetReadDeadline failed: %w", err)
	}

	n, err := conn.Read(buf)
	if err != nil {
		if err == io.EOF {
			return X11AuthRequired, nil
		}
		return X11Dead, fmt.Errorf("read failed: %w", err)
	}

	if n == 0 {
		return X11Dead, fmt.Errorf("read returned 0 bytes")
	}

	switch buf[0] {
	case 0x01:
		return X11Live, nil
	case 0x02:
		return X11AuthRequired, nil
	case 0x00:
		fmt.Printf("X11 handshake rejected (value %x)\n", buf[0])
		return X11AuthRequired, nil
	default:
		fmt.Printf("unexpected value %x\n", buf[0])
		return X11Dead, nil
	}

}

func CheckDisplay(display string) bool {
	cmd := exec.Command("xset", "q")
	cmd.Env = append(os.Environ(), fmt.Sprintf("DISPLAY=%s", display))

	if err := cmd.Run(); err != nil {
		fmt.Fprintf(os.Stderr, "xset q failed for DISPLAY=%s: %v\n", display, err)
		return false
	}

	return true
}