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

package control

import (
	"bufio"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"testing"
	"time"
)

type processCommand struct {
	Path, Text string
	Stop       bool
}

type processReply struct {
	PID      int
	Response string
	Error    string
}

type requestReceipt struct {
	InstancePID, CallerPID, CallerUID int
	Text                              string
}

// Each fixture reader runs in a separate process so PID endpoints, ancestry,
// registry ownership and descriptor credentials are exercised through the OS.
func TestControlProcess(t *testing.T) {
	mode := os.Getenv("FENRIR_CONTROL_PROCESS")
	if mode == "" {
		return
	}
	if mode == "client" {
		reply := processReply{PID: os.Getpid()}
		connection, err := net.DialTimeout("unix", os.Getenv("FENRIR_CONTROL_PATH"), time.Second)
		if err == nil {
			defer connection.Close()
			_ = connection.SetDeadline(time.Now().Add(3 * time.Second))
			_, err = io.WriteString(connection, os.Getenv("FENRIR_CONTROL_TEXT")+"\n")
			if err == nil {
				data, readErr := io.ReadAll(connection)
				reply.Response, err = string(data), readErr
			}
		}
		if err != nil {
			reply.Error = err.Error()
		}
		_ = json.NewEncoder(os.Stdout).Encode(reply)
		os.Exit(0)
	}
	run_reader_process(mode)
	os.Exit(0)
}

func child_command(mode string) *exec.Cmd {
	command := exec.Command(os.Args[0], "-test.run=^TestControlProcess$")
	command.Env = append(os.Environ(), "FENRIR_CONTROL_PROCESS="+mode)
	return command
}

func run_reader_process(mode string) {
	encoder := json.NewEncoder(os.Stdout)
	directory := os.Getenv("FENRIR_CONTROL_DIRECTORY")
	server, err := Open(context.Background(), Options{Directory: directory, RegistryDirectory: filepath.Join(directory, "registry"), TerminalMode: mode == "terminal", WindowID: os.Getenv("WINDOWID")})
	if err != nil {
		_ = encoder.Encode(processReply{Error: err.Error()})
		return
	}
	defer server.Close()
	go func() {
		for {
			select {
			case request := <-server.Requests():
				data, _ := json.Marshal(requestReceipt{os.Getpid(), request.PID, request.UID, request.Text})
				request.Reply <- string(data)
			case <-server.ctx.Done():
				return
			}
		}
	}()
	_ = encoder.Encode(processReply{PID: os.Getpid()})
	scanner := bufio.NewScanner(os.Stdin)
	for scanner.Scan() {
		var request processCommand
		if json.Unmarshal(scanner.Bytes(), &request) != nil || request.Stop {
			return
		}
		command := child_command("client")
		command.Env = append(command.Env, "FENRIR_CONTROL_PATH="+request.Path, "FENRIR_CONTROL_TEXT="+request.Text)
		data, err := command.Output()
		var reply processReply
		if err != nil {
			reply.Error = err.Error()
		} else if err := json.Unmarshal(data, &reply); err != nil {
			reply.Error = err.Error()
		}
		_ = encoder.Encode(reply)
	}
}

type readerProcess struct {
	command *exec.Cmd
	input   *json.Encoder
	output  *json.Decoder
	pid     int
	stopped bool
}

func start_reader_process(t *testing.T, directory, mode string) *readerProcess {
	t.Helper()
	command := child_command(mode)
	command.Env = append(command.Env, "FENRIR_CONTROL_DIRECTORY="+directory)
	command.Stderr = os.Stderr
	input, err := command.StdinPipe()
	if err != nil {
		t.Fatal(err)
	}
	output, err := command.StdoutPipe()
	if err != nil {
		t.Fatal(err)
	}
	if err := command.Start(); err != nil {
		t.Fatal(err)
	}
	process := &readerProcess{command: command, input: json.NewEncoder(input), output: json.NewDecoder(output)}
	t.Cleanup(func() { process.stop(t) })
	var ready processReply
	if err := process.output.Decode(&ready); err != nil || ready.Error != "" {
		t.Fatal("reader startup", ready, err)
	}
	process.pid = ready.PID
	return process
}

func (p *readerProcess) stop(t *testing.T) {
	t.Helper()
	if p.stopped {
		return
	}
	p.stopped = true
	_ = p.input.Encode(processCommand{Stop: true})
	done := make(chan error, 1)
	go func() { done <- p.command.Wait() }()
	select {
	case err := <-done:
		if err != nil {
			t.Error("reader exit", err)
		}
	case <-time.After(5 * time.Second):
		_ = p.command.Process.Kill()
		<-done
		t.Error("reader shutdown blocked")
	}
}

func (p *readerProcess) send(t *testing.T, path, text string) processReply {
	t.Helper()
	if err := p.input.Encode(processCommand{Path: path, Text: text}); err != nil {
		t.Fatal(err)
	}
	var reply processReply
	if err := p.output.Decode(&reply); err != nil || reply.Error != "" {
		t.Fatal("remote request", reply, err)
	}
	return reply
}

func check_receipt(t *testing.T, reply processReply, instancePID int, text string) {
	t.Helper()
	var receipt requestReceipt
	if err := json.Unmarshal([]byte(reply.Response), &receipt); err != nil {
		t.Fatal("invalid remote reply", reply, err)
	}
	if receipt.InstancePID != instancePID || receipt.CallerPID != reply.PID || receipt.CallerUID != os.Geteuid() || receipt.Text != text {
		t.Fatalf("receipt=%+v; want instance=%d caller=%d uid=%d text=%q", receipt, instancePID, reply.PID, os.Geteuid(), text)
	}
}

func wait_main_owner(t *testing.T, observer *Server, expectedCount int, excludedPID int) int {
	t.Helper()
	deadline := time.Now().Add(5 * time.Second)
	for time.Now().Before(deadline) {
		instances := observer.instances()
		owners, owner := 0, 0
		for _, data := range instances {
			if data.MainSocket {
				owners++
				owner = data.PID
			}
		}
		if len(instances) == expectedCount && owners == 1 && owner != excludedPID {
			return owner
		}
		time.Sleep(20 * time.Millisecond)
	}
	t.Fatal("shared endpoint ownership did not settle", observer.instances())
	return 0
}

func TestReadersCoexistInBothStartupOrdersAndSurviveOwnerExit(t *testing.T) {
	for _, order := range []string{"tty-first", "terminal-first"} {
		t.Run(order, func(t *testing.T) {
			options := private_options(t)
			observer := &Server{ctx: context.Background(), options: options}
			mainPath := observer.main_path()
			var tty, first, second, owner *readerProcess
			if order == "tty-first" {
				tty = start_reader_process(t, options.Directory, "tty")
				owner = tty
				first = start_reader_process(t, options.Directory, "terminal")
			} else {
				first = start_reader_process(t, options.Directory, "terminal")
				owner = first
				tty = start_reader_process(t, options.Directory, "tty")
			}
			second = start_reader_process(t, options.Directory, "terminal")
			if got := wait_main_owner(t, observer, 3, 0); got != owner.pid {
				t.Fatal("startup replaced the shared owner", got, owner.pid)
			}
			listing := first.send(t, mainPath, "COMMAND LS").Response
			for _, reader := range []*readerProcess{tty, first, second} {
				if !strings.Contains(listing, fmt.Sprintf("pid=%d ", reader.pid)) {
					t.Fatal("missing reader", reader.pid, listing)
				}
			}
			for _, reader := range []*readerProcess{tty, first, second} {
				check_receipt(t, reader.send(t, mainPath, "COMMAND SAY shared"), reader.pid, "COMMAND SAY shared")
				privatePath := filepath.Join(options.Directory, fmt.Sprintf("fenrirscreenreader-%d.sock", reader.pid))
				check_receipt(t, second.send(t, privatePath, "COMMAND SAY private"), reader.pid, "COMMAND SAY private")
			}
			// An unrelated caller must not fall through to the console reader.
			client := child_command("client")
			client.Env = append(client.Env, "FENRIR_CONTROL_PATH="+mainPath, "FENRIR_CONTROL_TEXT=COMMAND SAY legacy")
			data, err := client.Output()
			if err != nil {
				t.Fatal(err)
			}
			var reply processReply
			if err := json.Unmarshal(data, &reply); err != nil {
				t.Fatal(err)
			}
			if !strings.Contains(reply.Response, "unique reader") {
				t.Fatal("unidentified caller reached a reader", reply)
			}
			owner.stop(t)
			wait_main_owner(t, observer, 2, owner.pid)
			for _, reader := range []*readerProcess{tty, first, second} {
				if reader != owner {
					check_receipt(t, reader.send(t, mainPath, "COMMAND SAY survivor"), reader.pid, "COMMAND SAY survivor")
				}
			}
		})
	}
}

func external_request(t *testing.T, path, text string) processReply {
	t.Helper()
	client := child_command("client")
	client.Env = append(client.Env, "FENRIR_CONTROL_PATH="+path, "FENRIR_CONTROL_TEXT="+text)
	data, err := client.Output()
	if err != nil {
		t.Fatal(err)
	}
	var reply processReply
	if err := json.Unmarshal(data, &reply); err != nil || reply.Error != "" {
		t.Fatal(reply, err)
	}
	return reply
}

func TestTerminalReadersRejectAmbiguityAndRecoverAfterOwnerFailure(t *testing.T) {
	options := private_options(t)
	observer := &Server{ctx: context.Background(), options: options}
	mainPath := observer.main_path()
	first := start_reader_process(t, options.Directory, "terminal")
	second := start_reader_process(t, options.Directory, "terminal")
	third := start_reader_process(t, options.Directory, "terminal")
	wait_main_owner(t, observer, 3, 0)
	for _, reader := range []*readerProcess{first, second, third} {
		check_receipt(t, reader.send(t, mainPath, "COMMAND SAY terminal"), reader.pid, "COMMAND SAY terminal")
	}
	if reply := external_request(t, mainPath, "COMMAND SAY ambiguous"); !strings.Contains(reply.Response, "unique reader") {
		t.Fatal("arbitrary terminal selected", reply)
	}
	if reply := external_request(t, mainPath, "LIST"); strings.Count(reply.Response, "pid=") != 6 {
		// Each record contains both pid= and ppid=.
		t.Fatal("ambiguous caller could not discover instances", reply)
	}
	// Registration failure after binding a new private socket must preserve
	// the active shared endpoint and every existing registry record.
	before, err := os.Stat(mainPath)
	if err != nil {
		t.Fatal(err)
	}
	broken := options
	broken.RegistryDirectory = filepath.Join(options.Directory, "registry-file")
	if err := os.WriteFile(broken.RegistryDirectory, []byte("preserve"), 0600); err != nil {
		t.Fatal(err)
	}
	if server, err := Open(context.Background(), broken); err == nil {
		server.Close()
		t.Fatal("accepted a registry file")
	}
	privatePath := filepath.Join(options.Directory, fmt.Sprintf("fenrirscreenreader-%d.sock", os.Getpid()))
	if _, err := os.Lstat(privatePath); !os.IsNotExist(err) {
		t.Fatal("failed startup leaked a private socket", err)
	}
	after, err := os.Stat(mainPath)
	if err != nil || !os.SameFile(before, after) || len(observer.instances()) != 3 {
		t.Fatal("failed startup disturbed active endpoints", err)
	}
	first.stopped = true
	if err := first.command.Process.Kill(); err != nil {
		t.Fatal(err)
	}
	_ = first.command.Wait()
	wait_main_owner(t, observer, 2, first.pid)
	for _, reader := range []*readerProcess{second, third} {
		check_receipt(t, reader.send(t, mainPath, "COMMAND SAY recovered"), reader.pid, "COMMAND SAY recovered")
	}
}

// Screen reattachment and Codex's app-server can leave the caller outside the
// reader's descendant tree. The connecting process retains its X11 identity.
func TestWindowRoutingWithoutReaderAncestor(t *testing.T) {
	for _, consoleFirst := range []bool{true, false} {
		name := "terminal-first"
		if consoleFirst {
			name = "console-first"
		}
		t.Run(name, func(t *testing.T) {
			options := private_options(t)
			observer := &Server{ctx: context.Background(), options: options}
			mainPath := observer.main_path()
			t.Setenv("DISPLAY", ":fixture")
			t.Setenv("WINDOWID", "42")
			var terminal *readerProcess
			if consoleFirst {
				start_reader_process(t, options.Directory, "tty")
			}
			terminal = start_reader_process(t, options.Directory, "terminal")
			if !consoleFirst {
				start_reader_process(t, options.Directory, "tty")
			}
			// This client is a sibling, not a child, of the terminal reader.
			text := "command application codex"
			check_receipt(t, external_request(t, mainPath, text), terminal.pid, text)
			t.Setenv("WINDOWID", "43")
			second := start_reader_process(t, options.Directory, "terminal")
			check_receipt(t, external_request(t, mainPath, text), second.pid, text)
			t.Setenv("WINDOWID", "42")
			check_receipt(t, external_request(t, mainPath, text), terminal.pid, text)
			// Even a sole terminal on a different display is not a match.
			t.Setenv("DISPLAY", ":other")
			if reply := external_request(t, mainPath, text); !strings.Contains(reply.Response, "unique reader") {
				t.Fatal("cross-display routing", reply)
			}
			t.Setenv("DISPLAY", ":fixture")
			t.Setenv("WINDOWID", "unknown")
			if reply := external_request(t, mainPath, text); !strings.Contains(reply.Response, "unique reader") {
				t.Fatal("unknown window routed", reply)
			}
			t.Setenv("WINDOWID", "")
			if reply := external_request(t, mainPath, text); !strings.Contains(reply.Response, "unique reader") {
				t.Fatal("unidentified background caller routed", reply)
			}
			t.Setenv("WINDOWID", "42")
			start_reader_process(t, options.Directory, "terminal")
			if reply := external_request(t, mainPath, text); !strings.Contains(reply.Response, "unique reader") {
				t.Fatal("duplicate window routed arbitrarily", reply)
			}
			// Discovery remains usable when routing is ambiguous.
			if reply := external_request(t, mainPath, "ls"); !strings.Contains(reply.Response, "pid=") {
				t.Fatal("discovery failed", reply)
			}
		})
	}
}

func TestConsoleIdentityAndUnknownCallers(t *testing.T) {
	options := private_options(t)
	console := instance{PID: os.Getpid(), uid: os.Geteuid(), ScreenDriver: "vcsaDriver"}
	server := &Server{ctx: context.Background(), options: options}
	for _, trial := range []struct {
		name     string
		identity processIdentity
		consoles []instance
		valid    bool
	}{
		{"console", processIdentity{console: true}, []instance{console}, true},
		{"console ignores stale window", processIdentity{console: true, window: "42", display: ":fixture"}, []instance{console}, true},
		{"multiple consoles", processIdentity{console: true}, []instance{console, console}, false},
		{"no console", processIdentity{console: true}, nil, false},
		{"background", processIdentity{}, []instance{console}, false},
		{"unknown window", processIdentity{window: "42", display: ":fixture"}, []instance{console}, false},
	} {
		t.Run(trial.name, func(t *testing.T) {
			_, err := server.route_identity(trial.identity, os.Geteuid(), nil, trial.consoles)
			if (err == nil) != trial.valid {
				t.Fatal("unexpected routing", err)
			}
		})
	}
	if got := process_identity(os.Getpid(), os.Geteuid()+1); got != (processIdentity{}) {
		t.Fatal("read identity of a different UID", got)
	}
}