package session

import (
	"fmt"
	"io"
	"os"
	"os/exec"
	"strconv"
	"strings"
	"syscall"

	"latchd/config"
)

type PostLoginEnv struct {
	Kind           string // "x11", "wayland", "tty"
	Title          string
	XinitrcPath    string
	ScriptPath     string
	DesktopNames   string // for XDG_CURRENT_DESKTOP (colon-separated)
	SessionDesktop string // for XDG_SESSION_DESKTOP
}

type SpawnedEnv struct {
	Kind   string
	Server *exec.Cmd
	Client *exec.Cmd
	logFd  *os.File
}

const logLimit = 67_108_864

type limitWriter struct {
	w   io.Writer
	n   int64
	max int64
}

func (l *limitWriter) Write(p []byte) (int, error) {
	if l.n >= l.max {
		return len(p), nil
	}
	rem := l.max - l.n
	if int64(len(p)) > rem {
		p = p[:rem]
	}
	n, err := l.w.Write(p)
	l.n += int64(n)
	return n, err
}

func Spawn(env *PostLoginEnv, uid, gid uint32, groups []uint32, cfg *config.Config) (*SpawnedEnv, error) {
	var loginFlag string
	switch cfg.ShellLoginFlag {
	case config.ShellLoginShort:
		loginFlag = "-l"
	case config.ShellLoginLong:
		loginFlag = "--login"
	}

	switch env.Kind {
	case "x11":
		return spawnX11(env, uid, gid, groups, cfg, loginFlag)
	case "wayland":
		return spawnWayland(env, uid, gid, groups, cfg, loginFlag)
	default:
		return spawnTTY(env, uid, gid, groups, cfg)
	}
}

func spawnX11(env *PostLoginEnv, uid, gid uint32, groups []uint32, cfg *config.Config, loginFlag string) (*SpawnedEnv, error) {
	server, err := setupX(uid, gid, groups, cfg)
	if err != nil {
		return nil, fmt.Errorf("X setup: %w", err)
	}

	args := shellCmd(cfg.SystemShell, loginFlag, cfg.X11.SetupPath+" "+env.XinitrcPath)
	client := privCmd(args, uid, gid, groups)

	logF := openLog(cfg.ClientLogPath)
	if logF != nil {
		lw := &limitWriter{w: logF, max: logLimit}
		client.Stdout = lw
		client.Stderr = lw
	}

	if err := client.Start(); err != nil {
		return nil, fmt.Errorf("X11 client: %w", err)
	}
	return &SpawnedEnv{Kind: "x11", Server: server, Client: client, logFd: logF}, nil
}

func spawnWayland(env *PostLoginEnv, uid, gid uint32, groups []uint32, cfg *config.Config, loginFlag string) (*SpawnedEnv, error) {
	args := shellCmd(cfg.SystemShell, loginFlag, env.ScriptPath)
	client := privCmd(args, uid, gid, groups)

	logF := openLog(cfg.ClientLogPath)
	if logF != nil {
		lw := &limitWriter{w: logF, max: logLimit}
		client.Stdout = lw
		client.Stderr = lw
	}

	if err := client.Start(); err != nil {
		return nil, fmt.Errorf("Wayland: %w", err)
	}
	return &SpawnedEnv{Kind: "wayland", Client: client, logFd: logF}, nil
}

func spawnTTY(env *PostLoginEnv, uid, gid uint32, groups []uint32, cfg *config.Config) (*SpawnedEnv, error) {
	shell := userShell(int(uid))
	_ = env
	_ = cfg

	client := privCmd([]string{shell}, uid, gid, groups)
	client.Stdin = os.Stdin
	client.Stdout = os.Stdout
	client.Stderr = os.Stderr

	if err := client.Start(); err != nil {
		return nil, fmt.Errorf("TTY: %w", err)
	}
	return &SpawnedEnv{Kind: "tty", Client: client}, nil
}

func (s *SpawnedEnv) Wait() error {
	defer func() {
		if s.logFd != nil {
			s.logFd.Close()
		}
	}()
	if s.Kind == "x11" && s.Server != nil {
		defer func() {
			s.Server.Process.Signal(syscall.SIGTERM)
			s.Server.Wait()
		}()
	}
	if s.Client != nil {
		return s.Client.Wait()
	}
	return nil
}

func (s *SpawnedEnv) PID() int {
	if s.Client != nil && s.Client.Process != nil {
		return s.Client.Process.Pid
	}
	return 0
}

func shellCmd(shell, loginFlag, command string) []string {
	a := []string{shell}
	if loginFlag != "" {
		a = append(a, loginFlag)
	}
	return append(a, "-c", command)
}

func privCmd(args []string, uid, gid uint32, groups []uint32) *exec.Cmd {
	cmd := exec.Command(args[0], args[1:]...)
	cmd.SysProcAttr = &syscall.SysProcAttr{
		Credential: &syscall.Credential{
			Uid:    uid,
			Gid:    gid,
			Groups: groups,
		},
		Setsid:  true,
		Setctty: false,
	}
	return cmd
}

func openLog(path string) *os.File {
	f, err := os.OpenFile(path, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0644)
	if err != nil {
		return nil
	}
	return f
}

func userShell(uid int) string {
	data, err := os.ReadFile("/etc/passwd")
	if err != nil {
		return "/bin/sh"
	}
	uidStr := strconv.Itoa(uid)
	for _, line := range strings.Split(string(data), "\n") {
		parts := strings.Split(line, ":")
		if len(parts) >= 7 && parts[2] == uidStr && parts[6] != "" {
			sh := parts[6]
			if sh != "/sbin/nologin" && sh != "/usr/sbin/nologin" && sh != "/bin/false" {
				return sh
			}
		}
	}
	return "/bin/sh"
}