// Copyright (c) 2026 Storm Dragon
// SPDX-License-Identifier: GPL-3.0-or-later

package control

import (
	"fmt"
	"os"
	"path/filepath"
	"strconv"
	"syscall"
	"testing"
)

type screenFixture struct {
	root, executable string
	t                *testing.T
}

func new_screen_fixture(t *testing.T) screenFixture {
	t.Helper()
	f := screenFixture{root: t.TempDir(), t: t}
	f.executable = filepath.Join(f.root, "screen")
	if err := os.WriteFile(f.executable, nil, 0700); err != nil {
		t.Fatal(err)
	}
	f.process(100, 1, 0, "/usr/bin/SCREEN", f.executable)
	f.process(120, 100, 0, "notifier", "")
	f.process(300, 1, 0, "reader", "")
	return f
}

func (f screenFixture) process(pid, parent int, tty uint32, command, executable string) {
	f.t.Helper()
	path := filepath.Join(f.root, strconv.Itoa(pid))
	if err := os.MkdirAll(filepath.Join(path, "fd"), 0700); err != nil {
		f.t.Fatal(err)
	}
	// A parenthesized name checks the /proc/stat field boundary too.
	stat := fmt.Sprintf("%d (process (name)) S %d %d %d %d 0 0", pid, parent, pid, pid, int32(tty))
	for name, value := range map[string]string{"stat": stat, "cmdline": command + "\x00"} {
		if err := os.WriteFile(filepath.Join(path, name), []byte(value), 0600); err != nil {
			f.t.Fatal(err)
		}
	}
	if executable != "" {
		if err := os.Symlink(executable, filepath.Join(path, "exe")); err != nil {
			f.t.Fatal(err)
		}
	}
}

func (f screenFixture) attach(pid, reader int, device string, sameExecutable bool) {
	f.t.Helper()
	info, err := os.Stat(device)
	if err != nil {
		f.t.Fatal(err)
	}
	executable := f.executable
	if !sameExecutable {
		executable = filepath.Join(f.root, "other")
		if err := os.WriteFile(executable, nil, 0700); err != nil {
			f.t.Fatal(err)
		}
	}
	f.process(pid, reader, uint32(info.Sys().(*syscall.Stat_t).Rdev), "screen", executable)
	if err := os.Symlink(device, filepath.Join(f.root, "100", "fd", strconv.Itoa(pid))); err != nil {
		f.t.Fatal(err)
	}
}

func TestScreenAttachmentRouting(t *testing.T) {
	for _, scenario := range []string{"attached", "detached", "other executable", "ambiguous", "duplicate displays", "reattached", "no reader", "wrong uid"} {
		t.Run(scenario, func(t *testing.T) {
			f := new_screen_fixture(t)
			uid := os.Geteuid()
			readers := []instance{{PID: 300}, {PID: 400}}
			if scenario != "detached" {
				f.attach(200, 300, "/dev/null", scenario != "other executable")
			}
			if scenario == "ambiguous" {
				f.process(400, 1, 0, "reader", "")
				f.attach(201, 400, "/dev/zero", true)
			}
			if scenario == "duplicate displays" {
				f.attach(201, 300, "/dev/zero", true)
			}
			if scenario == "no reader" {
				readers = nil
			}
			if scenario == "wrong uid" {
				uid++
			}
			if scenario == "reattached" {
				if err := os.Remove(filepath.Join(f.root, "100", "fd", "200")); err != nil {
					t.Fatal(err)
				}
				f.process(400, 1, 0, "reader", "")
				f.attach(201, 400, "/dev/zero", true)
			}
			reader, handled, err := attached_screen_reader(f.root, []int{120, 100}, uid, readers)
			valid := scenario == "attached" || scenario == "duplicate displays" || scenario == "reattached"
			if handled != (scenario != "wrong uid") || (err == nil && handled) != valid {
				t.Fatal(reader, handled, err)
			}
			if valid {
				want := 300
				if scenario == "reattached" {
					want = 400
				}
				if reader.PID != want {
					t.Fatal("wrong reader", reader.PID, want)
				}
			}
		})
	}
}

func TestScreenRoutingIgnoresEnvironmentWithoutServerAncestor(t *testing.T) {
	f := new_screen_fixture(t)
	f.attach(200, 300, "/dev/null", true)
	t.Setenv("STY", "100.session")
	if _, handled, err := attached_screen_reader(f.root, []int{120}, os.Geteuid(), []instance{{PID: 300}}); handled || err != nil {
		t.Fatal("environment selected an unrelated Screen session", handled, err)
	}
	// A broken ancestry cycle must terminate without choosing an unrelated reader.
	f.process(301, 302, 0, "cycle", "")
	f.process(302, 301, 0, "cycle", "")
	f.attach(201, 301, "/dev/zero", true)
	reader, handled, err := attached_screen_reader(f.root, []int{120, 100}, os.Geteuid(), []instance{{PID: 300}})
	if !handled || err != nil || reader.PID != 300 {
		t.Fatal(reader, handled, err)
	}
}