//go:build (linux && !android) || freebsd

package server

import (
	"fmt"
	"os"
	"os/exec"
	"os/user"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"

	log "github.com/sirupsen/logrus"
)

// VirtualSession manages a virtual X11 display (Xvfb) with a desktop session
// running as a target user. It implements ScreenCapturer and InputInjector by
// delegating to an X11Capturer/X11InputInjector pointed at the virtual display.
const sessionIdleTimeout = 5 * time.Minute

type VirtualSession struct {
	mu       sync.Mutex
	display  string
	user     *user.User
	uid      uint32
	gid      uint32
	groups   []uint32
	xvfb     *exec.Cmd
	desktop  *exec.Cmd
	poller   *X11Poller
	injector *X11InputInjector
	log      *log.Entry
	stopped  bool
	clients  int
	idleTimer *time.Timer
	onIdle   func() // called when idle timeout fires or Xvfb dies
}

// StartVirtualSession creates and starts a virtual X11 session for the given user.
// Requires root privileges to create sessions as other users.
func StartVirtualSession(username string, logger *log.Entry) (*VirtualSession, error) {
	if os.Getuid() != 0 {
		return nil, fmt.Errorf("virtual sessions require root privileges")
	}

	if _, err := exec.LookPath("Xvfb"); err != nil {
		if _, err := exec.LookPath("Xorg"); err != nil {
			return nil, fmt.Errorf("neither Xvfb nor Xorg found (install xvfb or xserver-xorg)")
		}
		if !hasDummyDriver() {
			return nil, fmt.Errorf("Xvfb not found and Xorg dummy driver not installed (install xvfb or xf86-video-dummy)")
		}
	}

	u, err := user.Lookup(username)
	if err != nil {
		return nil, fmt.Errorf("lookup user %s: %w", username, err)
	}

	uid, err := strconv.ParseUint(u.Uid, 10, 32)
	if err != nil {
		return nil, fmt.Errorf("parse uid: %w", err)
	}
	gid, err := strconv.ParseUint(u.Gid, 10, 32)
	if err != nil {
		return nil, fmt.Errorf("parse gid: %w", err)
	}

	groups, err := supplementaryGroups(u)
	if err != nil {
		logger.Debugf("supplementary groups for %s: %v", username, err)
	}

	vs := &VirtualSession{
		user:   u,
		uid:    uint32(uid),
		gid:    uint32(gid),
		groups: groups,
		log:    logger.WithField("vnc_user", username),
	}

	if err := vs.start(); err != nil {
		return nil, err
	}
	return vs, nil
}

func (vs *VirtualSession) start() error {
	display, err := findFreeDisplay()
	if err != nil {
		return fmt.Errorf("find free display: %w", err)
	}
	vs.display = display

	if err := vs.startXvfb(); err != nil {
		return err
	}

	socketPath := fmt.Sprintf("/tmp/.X11-unix/X%s", vs.display[1:])
	if err := waitForPath(socketPath, 5*time.Second); err != nil {
		vs.stopXvfb()
		return fmt.Errorf("wait for X11 socket %s: %w", socketPath, err)
	}

	// Grant the target user access to the display via xhost.
	xhostCmd := exec.Command("xhost", "+SI:localuser:"+vs.user.Username)
	xhostCmd.Env = []string{"DISPLAY=" + vs.display}
	if out, err := xhostCmd.CombinedOutput(); err != nil {
		vs.log.Debugf("xhost: %s (%v)", strings.TrimSpace(string(out)), err)
	}

	vs.poller = NewX11Poller(vs.display)

	injector, err := NewX11InputInjector(vs.display)
	if err != nil {
		vs.stopXvfb()
		return fmt.Errorf("create X11 injector for %s: %w", vs.display, err)
	}
	vs.injector = injector

	if err := vs.startDesktop(); err != nil {
		vs.injector.Close()
		vs.stopXvfb()
		return fmt.Errorf("start desktop: %w", err)
	}

	vs.log.Infof("virtual session started: display=%s user=%s", vs.display, vs.user.Username)
	return nil
}

// ClientConnect increments the client count and cancels any idle timer.
func (vs *VirtualSession) ClientConnect() {
	vs.mu.Lock()
	defer vs.mu.Unlock()
	vs.clients++
	if vs.idleTimer != nil {
		vs.idleTimer.Stop()
		vs.idleTimer = nil
	}
}

// ClientDisconnect decrements the client count. When the last client
// disconnects, starts an idle timer that destroys the session.
func (vs *VirtualSession) ClientDisconnect() {
	vs.mu.Lock()
	defer vs.mu.Unlock()
	vs.clients--
	if vs.clients <= 0 {
		vs.clients = 0
		vs.log.Infof("no VNC clients connected, session will be destroyed in %s", sessionIdleTimeout)
		vs.idleTimer = time.AfterFunc(sessionIdleTimeout, vs.idleExpired)
	}
}

// idleExpired is called by the idle timer. It stops the session and
// notifies the session manager via onIdle so it removes us from the map.
func (vs *VirtualSession) idleExpired() {
	vs.log.Info("idle timeout reached, destroying virtual session")
	vs.Stop()
	// onIdle acquires sessionManager.mu; safe because Stop() has released vs.mu.
	if vs.onIdle != nil {
		vs.onIdle()
	}
}

// isAlive returns true if the session is running and its X server socket exists.
func (vs *VirtualSession) isAlive() bool {
	vs.mu.Lock()
	stopped := vs.stopped
	display := vs.display
	vs.mu.Unlock()

	if stopped {
		return false
	}
	// Verify the X socket still exists on disk.
	socketPath := fmt.Sprintf("/tmp/.X11-unix/X%s", display[1:])
	if _, err := os.Stat(socketPath); err != nil {
		return false
	}
	return true
}

// Capturer returns the screen capturer for this virtual session.
func (vs *VirtualSession) Capturer() ScreenCapturer {
	return vs.poller
}

// Injector returns the input injector for this virtual session.
func (vs *VirtualSession) Injector() InputInjector {
	return vs.injector
}

// Display returns the X11 display string (e.g., ":99").
func (vs *VirtualSession) Display() string {
	return vs.display
}

// Stop terminates the virtual session, killing the desktop and Xvfb.
func (vs *VirtualSession) Stop() {
	vs.mu.Lock()
	defer vs.mu.Unlock()

	if vs.stopped {
		return
	}
	vs.stopped = true

	if vs.injector != nil {
		vs.injector.Close()
	}

	vs.stopDesktop()
	vs.stopXvfb()

	vs.log.Info("virtual session stopped")
}

func (vs *VirtualSession) startXvfb() error {
	if _, err := exec.LookPath("Xvfb"); err == nil {
		return vs.startXvfbDirect()
	}
	return vs.startXorgDummy()
}

func (vs *VirtualSession) startXvfbDirect() error {
	vs.xvfb = exec.Command("Xvfb", vs.display,
		"-screen", "0", "1280x800x24",
		"-ac",
		"-nolisten", "tcp",
	)
	vs.xvfb.SysProcAttr = &syscall.SysProcAttr{Setsid: true, Pdeathsig: syscall.SIGTERM}

	if err := vs.xvfb.Start(); err != nil {
		return fmt.Errorf("start Xvfb on %s: %w", vs.display, err)
	}
	vs.log.Infof("Xvfb started on %s (pid=%d)", vs.display, vs.xvfb.Process.Pid)

	go vs.monitorXvfb()

	return nil
}

// startXorgDummy starts Xorg with the dummy video driver as a fallback when
// Xvfb is not installed. Most systems with a desktop have Xorg available.
func (vs *VirtualSession) startXorgDummy() error {
	confPath := fmt.Sprintf("/tmp/nbvnc-dummy-%s.conf", vs.display[1:])
	conf := `Section "Device"
    Identifier "dummy"
    Driver "dummy"
    VideoRam 256000
EndSection
Section "Screen"
    Identifier "screen"
    Device "dummy"
    DefaultDepth 24
    SubSection "Display"
        Depth 24
        Modes "1280x800"
    EndSubSection
EndSection
`
	if err := os.WriteFile(confPath, []byte(conf), 0644); err != nil {
		return fmt.Errorf("write Xorg dummy config: %w", err)
	}

	vs.xvfb = exec.Command("Xorg", vs.display,
		"-config", confPath,
		"-noreset",
		"-nolisten", "tcp",
		"-ac",
	)
	vs.xvfb.SysProcAttr = &syscall.SysProcAttr{Setsid: true, Pdeathsig: syscall.SIGTERM}

	if err := vs.xvfb.Start(); err != nil {
		os.Remove(confPath)
		return fmt.Errorf("start Xorg dummy on %s: %w", vs.display, err)
	}
	vs.log.Infof("Xorg (dummy driver) started on %s (pid=%d)", vs.display, vs.xvfb.Process.Pid)

	go func() {
		vs.monitorXvfb()
		os.Remove(confPath)
	}()

	return nil
}

// monitorXvfb waits for the Xvfb/Xorg process to exit. If it exits
// unexpectedly (not via Stop), the session is marked as dead and the
// onIdle callback fires so the session manager removes it from the map.
// The next GetOrCreate call for this user will create a fresh session.
func (vs *VirtualSession) monitorXvfb() {
	if err := vs.xvfb.Wait(); err != nil {
		vs.log.Debugf("X server exited: %v", err)
	}

	vs.mu.Lock()
	alreadyStopped := vs.stopped
	if !alreadyStopped {
		vs.log.Warn("X server exited unexpectedly, marking session as dead")
		vs.stopped = true
		if vs.idleTimer != nil {
			vs.idleTimer.Stop()
			vs.idleTimer = nil
		}
		if vs.injector != nil {
			vs.injector.Close()
		}
		vs.stopDesktop()
	}
	onIdle := vs.onIdle
	vs.mu.Unlock()

	if !alreadyStopped && onIdle != nil {
		onIdle()
	}
}

func (vs *VirtualSession) stopXvfb() {
	if vs.xvfb != nil && vs.xvfb.Process != nil {
		syscall.Kill(-vs.xvfb.Process.Pid, syscall.SIGTERM)
		time.Sleep(200 * time.Millisecond)
		syscall.Kill(-vs.xvfb.Process.Pid, syscall.SIGKILL)
	}
}

func (vs *VirtualSession) startDesktop() error {
	session := detectDesktopSession()

	// Wrap the desktop command with dbus-launch to provide a session bus.
	// Without this, most desktop environments (XFCE, MATE, etc.) fail immediately.
	var args []string
	if _, err := exec.LookPath("dbus-launch"); err == nil {
		args = append([]string{"dbus-launch", "--exit-with-session"}, session...)
	} else {
		args = session
	}

	vs.desktop = exec.Command(args[0], args[1:]...)
	vs.desktop.Dir = vs.user.HomeDir
	vs.desktop.Env = vs.buildUserEnv()
	vs.desktop.SysProcAttr = &syscall.SysProcAttr{
		Credential: &syscall.Credential{
			Uid:    vs.uid,
			Gid:    vs.gid,
			Groups: vs.groups,
		},
		Setsid:    true,
		Pdeathsig: syscall.SIGTERM,
	}

	if err := vs.desktop.Start(); err != nil {
		return fmt.Errorf("start desktop session (%v): %w", args, err)
	}
	vs.log.Infof("desktop session started: %v (pid=%d)", args, vs.desktop.Process.Pid)

	go func() {
		if err := vs.desktop.Wait(); err != nil {
			vs.log.Debugf("desktop session exited: %v", err)
		}
	}()

	return nil
}

func (vs *VirtualSession) stopDesktop() {
	if vs.desktop != nil && vs.desktop.Process != nil {
		syscall.Kill(-vs.desktop.Process.Pid, syscall.SIGTERM)
		time.Sleep(200 * time.Millisecond)
		syscall.Kill(-vs.desktop.Process.Pid, syscall.SIGKILL)
	}
}

func (vs *VirtualSession) buildUserEnv() []string {
	return []string{
		"DISPLAY=" + vs.display,
		"HOME=" + vs.user.HomeDir,
		"USER=" + vs.user.Username,
		"LOGNAME=" + vs.user.Username,
		"SHELL=" + getUserShell(vs.user.Uid),
		"PATH=/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin",
		"XDG_RUNTIME_DIR=/run/user/" + vs.user.Uid,
		"DBUS_SESSION_BUS_ADDRESS=unix:path=/run/user/" + vs.user.Uid + "/bus",
	}
}

// detectDesktopSession discovers available desktop sessions from the standard
// /usr/share/xsessions/*.desktop files (FreeDesktop standard, used by all
// display managers). Falls back to a hardcoded list if no .desktop files found.
func detectDesktopSession() []string {
	// Scan xsessions directories (Linux: /usr/share, FreeBSD: /usr/local/share).
	for _, dir := range []string{"/usr/share/xsessions", "/usr/local/share/xsessions"} {
		if cmd := findXSession(dir); cmd != nil {
			return cmd
		}
	}

	// Fallback: try common session commands directly.
	fallbacks := [][]string{
		{"startplasma-x11"},
		{"gnome-session"},
		{"xfce4-session"},
		{"mate-session"},
		{"cinnamon-session"},
		{"openbox-session"},
		{"xterm"},
	}
	for _, s := range fallbacks {
		if _, err := exec.LookPath(s[0]); err == nil {
			return s
		}
	}
	return []string{"xterm"}
}

// sessionPriority defines preference order for desktop environments.
// Lower number = higher priority. Unknown sessions get 100.
var sessionPriority = map[string]int{
	"plasma":     1, // KDE
	"gnome":      2,
	"xfce":       3,
	"mate":       4,
	"cinnamon":   5,
	"lxqt":       6,
	"lxde":       7,
	"budgie":     8,
	"openbox":    20,
	"fluxbox":    21,
	"i3":         22,
	"xinit":      50, // generic user session
	"lightdm":    50,
	"default":    50,
}

func findXSession(dir string) []string {
	entries, err := os.ReadDir(dir)
	if err != nil {
		return nil
	}

	type candidate struct {
		cmd      string
		priority int
	}
	var candidates []candidate

	for _, e := range entries {
		if !strings.HasSuffix(e.Name(), ".desktop") {
			continue
		}
		data, err := os.ReadFile(filepath.Join(dir, e.Name()))
		if err != nil {
			continue
		}
		execCmd := ""
		for _, line := range strings.Split(string(data), "\n") {
			if strings.HasPrefix(line, "Exec=") {
				execCmd = strings.TrimSpace(strings.TrimPrefix(line, "Exec="))
				break
			}
		}
		if execCmd == "" || execCmd == "default" {
			continue
		}

		// Determine priority from the filename or exec command.
		pri := 100
		lower := strings.ToLower(e.Name() + " " + execCmd)
		for keyword, p := range sessionPriority {
			if strings.Contains(lower, keyword) && p < pri {
				pri = p
			}
		}
		candidates = append(candidates, candidate{cmd: execCmd, priority: pri})
	}

	if len(candidates) == 0 {
		return nil
	}

	// Pick the highest priority (lowest number).
	best := candidates[0]
	for _, c := range candidates[1:] {
		if c.priority < best.priority {
			best = c
		}
	}

	// Verify the binary exists.
	parts := strings.Fields(best.cmd)
	if _, err := exec.LookPath(parts[0]); err != nil {
		return nil
	}
	return parts
}

// findFreeDisplay scans for an unused X11 display number.
func findFreeDisplay() (string, error) {
	for n := 50; n < 200; n++ {
		lockFile := fmt.Sprintf("/tmp/.X%d-lock", n)
		socketFile := fmt.Sprintf("/tmp/.X11-unix/X%d", n)
		if _, err := os.Stat(lockFile); err == nil {
			continue
		}
		if _, err := os.Stat(socketFile); err == nil {
			continue
		}
		return fmt.Sprintf(":%d", n), nil
	}
	return "", fmt.Errorf("no free X11 display found (checked :50-:199)")
}

// waitForPath polls until a filesystem path exists or the timeout expires.
func waitForPath(path string, timeout time.Duration) error {
	deadline := time.Now().Add(timeout)
	for time.Now().Before(deadline) {
		if _, err := os.Stat(path); err == nil {
			return nil
		}
		time.Sleep(50 * time.Millisecond)
	}
	return fmt.Errorf("timeout waiting for %s", path)
}

// getUserShell returns the login shell for the given UID.
func getUserShell(uid string) string {
	data, err := os.ReadFile("/etc/passwd")
	if err != nil {
		return "/bin/sh"
	}
	for _, line := range strings.Split(string(data), "\n") {
		fields := strings.Split(line, ":")
		if len(fields) >= 7 && fields[2] == uid {
			return fields[6]
		}
	}
	return "/bin/sh"
}

// supplementaryGroups returns the supplementary group IDs for a user.
func supplementaryGroups(u *user.User) ([]uint32, error) {
	gids, err := u.GroupIds()
	if err != nil {
		return nil, err
	}
	var groups []uint32
	for _, g := range gids {
		id, err := strconv.ParseUint(g, 10, 32)
		if err != nil {
			continue
		}
		groups = append(groups, uint32(id))
	}
	return groups, nil
}

// sessionManager tracks active virtual sessions by username.
type sessionManager struct {
	mu       sync.Mutex
	sessions map[string]*VirtualSession
	log      *log.Entry
}

func newSessionManager(logger *log.Entry) *sessionManager {
	return &sessionManager{
		sessions: make(map[string]*VirtualSession),
		log:      logger,
	}
}

// GetOrCreate returns an existing virtual session or creates a new one.
// If a previous session for this user is stopped or its X server died, it is replaced.
func (sm *sessionManager) GetOrCreate(username string) (vncSession, error) {
	sm.mu.Lock()
	defer sm.mu.Unlock()

	if vs, ok := sm.sessions[username]; ok {
		if vs.isAlive() {
			return vs, nil
		}
		sm.log.Infof("replacing dead virtual session for %s", username)
		vs.Stop()
		delete(sm.sessions, username)
	}

	vs, err := StartVirtualSession(username, sm.log)
	if err != nil {
		return nil, err
	}
	vs.onIdle = func() {
		sm.mu.Lock()
		defer sm.mu.Unlock()
		if cur, ok := sm.sessions[username]; ok && cur == vs {
			delete(sm.sessions, username)
			sm.log.Infof("removed idle virtual session for %s", username)
		}
	}
	sm.sessions[username] = vs
	return vs, nil
}

// hasDummyDriver checks common paths for the Xorg dummy video driver.
func hasDummyDriver() bool {
	paths := []string{
		"/usr/lib/xorg/modules/drivers/dummy_drv.so",                // Debian/Ubuntu
		"/usr/lib64/xorg/modules/drivers/dummy_drv.so",              // RHEL/Fedora
		"/usr/local/lib/xorg/modules/drivers/dummy_drv.so",          // FreeBSD
		"/usr/lib/x86_64-linux-gnu/xorg/modules/drivers/dummy_drv.so", // Debian multiarch
	}
	for _, p := range paths {
		if _, err := os.Stat(p); err == nil {
			return true
		}
	}
	return false
}

// StopAll terminates all active virtual sessions.
func (sm *sessionManager) StopAll() {
	sm.mu.Lock()
	defer sm.mu.Unlock()

	for username, vs := range sm.sessions {
		vs.Stop()
		delete(sm.sessions, username)
		sm.log.Infof("stopped virtual session for %s", username)
	}
}