package main

import (
	"fmt"
	"net"
	"os"
	"os/exec"
	"path"
	"strconv"
	"strings"
	"x11proxy/x11probe"
)

func resolveDisplay(display string) (x11probe.ConnType, string) {
	var dispNum int
	var err error

	if strings.HasPrefix(display, ":") {
		dispNum, err = strconv.Atoi(strings.Split(display[1:], ".")[0])
	} else {
		parts := strings.Split(display, ":")
		dispNum, err = strconv.Atoi(strings.Split(parts[1], ".")[0])
	}
	if err != nil {
		dispNum = 0
	}

	unixPath := path.Join("/tmp/.X11-unix", fmt.Sprintf("X%d", dispNum))
	if _, err := os.Stat(unixPath); err == nil {
		return x11probe.Unix, unixPath
	}

	tcpAddr := fmt.Sprintf("127.0.0.1:%d", 6000+dispNum)
	conn, err := net.Dial("tcp", tcpAddr)
	if err == nil {
		conn.Close()
		return x11probe.TCP, tcpAddr
	}

	return x11probe.Unix, unixPath
}

func getXAuthCookie(display string) ([]byte, error) {
	// Extract display number
	var dispNum string
	if strings.HasPrefix(display, ":") {
		dispNum = strings.Split(display[1:], ".")[0]
	} else {
		parts := strings.Split(display, ":")
		dispNum = strings.Split(parts[1], ".")[0]
	}

	// Query only entries for that display
	out, err := exec.Command("xauth", "list", ":"+dispNum).Output()
	if err != nil {
		return nil, fmt.Errorf("xauth list failed: %w", err)
	}

	lines := strings.Split(string(out), "\n")
	var bestMatch string
	var bestCookie string
	var bestScore int

	for _, line := range lines {
		fields := strings.Fields(line)
		if len(fields) < 3 || fields[1] != "MIT-MAGIC-COOKIE-1" {
			continue
		}

		entry := fields[0]
		cookie := fields[2]

		// Score based on specificity
		score := 0
		switch {
		case strings.HasPrefix(entry, os.Getenv("HOSTNAME")+"/unix:"):
			score = 100
		case strings.HasPrefix(entry, "unix:"):
			score = 90
		case strings.Contains(entry, "/unix:"):
			score = 80
		case strings.HasPrefix(entry, "#ffff#"):
			score = 70
		default:
			score = 50
		}

		if score > bestScore {
			bestMatch = entry
			bestCookie = cookie
			bestScore = score
		}
	}

	if bestCookie != "" {
		fmt.Printf("Using XAuth cookie from entry: %s\n", bestMatch)
		return parseHexCookie(bestCookie), nil
	}

	return nil, fmt.Errorf("no matching cookie found for display %s", display)
}

func parseHexCookie(hexStr string) []byte {
	var cookie []byte
	for i := 0; i < len(hexStr); i += 2 {
		var val byte
		fmt.Sscanf(hexStr[i:i+2], "%02x", &val)
		cookie = append(cookie, val)
	}
	return cookie
}

func PatchAuth(data []byte, cookie []byte) []byte {
	isLittleEndian := data[0] == 'l'

	var authProtoLen, authDataLen int
	if isLittleEndian {
		authProtoLen = int(data[7])<<8 | int(data[6])
		authDataLen = int(data[9])<<8 | int(data[8])
	} else {
		authProtoLen = int(data[6])<<8 | int(data[7])
		authDataLen = int(data[8])<<8 | int(data[9])
	}

	headerLen := 12
	authProtoPad := (authProtoLen + 3) & ^3
	authDataPad := (authDataLen + 3) & ^3
	authDataStart := headerLen + authProtoPad

	// Replace cookie and update length
	patched := make([]byte, headerLen+authProtoPad+authDataPad)
	copy(patched, data[:headerLen+authProtoPad])
	copy(patched[authDataStart:], cookie)

	// Update authDataLen to match cookie length
	cookieLen := len(cookie)
	if isLittleEndian {
		patched[8] = byte(cookieLen)
		patched[9] = byte(cookieLen >> 8)
	} else {
		patched[8] = byte(cookieLen >> 8)
		patched[9] = byte(cookieLen)
	}

	return patched
}