//go:build linux

package microvm

import (
	"bytes"
	"context"
	"crypto/ed25519"
	"crypto/rand"
	"encoding/pem"
	"fmt"
	"io"
	"log/slog"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"testing"
	"time"

	gliderssh "github.com/gliderlabs/ssh"
	gossh "golang.org/x/crypto/ssh"
	agentv1 "tangled.org/core/spindle/agentproto/gen"
	"tangled.org/core/spindle/models"
)

func TestDebugShellStaysOpenWhileIdleE2E(t *testing.T) {
	imageSpecPath := os.Getenv("SPINDLE_E2E_IMAGE_SPEC")
	if imageSpecPath == "" {
		t.Skip("SPINDLE_E2E_IMAGE_SPEC is not set")
	}

	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
	defer cancel()

	imageSpec, err := LoadImageSpec(imageSpecPath)
	if err != nil {
		t.Fatal(err)
	}

	agentListener, agentPort, err := listenRandomVsockPort(ctx)
	if err != nil {
		t.Fatal(err)
	}
	defer agentListener.Close()
	imageSpec.BootArgs = fmt.Sprintf("%s shuttle.vsock_port=%d", imageSpec.BootArgs, agentPort)

	workDir, err := os.MkdirTemp("/tmp", "spindle-debug-e2e-")
	if err != nil {
		t.Fatal(err)
	}
	defer os.RemoveAll(workDir)

	logger := slog.Default()
	vm, err := StartVM(ctx, VMConfig{
		Image:       imageSpec,
		BootTimeout: 30 * time.Second,
		EnableKVM:   true,
		WorkDir:     workDir,
		Dev:         true,
	}, logger)
	if err != nil {
		t.Fatal(err)
	}
	defer vm.Close()

	conn, err := acceptWithTimeout(ctx, &cidFilteredVsockListener{
		Listener: agentListener,
		cid:      vm.CID(),
		logger:   logger,
	}, 60*time.Second)
	if err != nil {
		t.Fatal(err)
	}
	defer conn.Close()

	agent := NewAgentSession(conn, vm.CID(), logger)
	if err := agent.Init(ctx, &agentv1.Init{JobId: "debug-idle-e2e"}); err != nil {
		t.Fatal(err)
	}

	failedStep := &agentv1.ExecStart{
		Argv: []string{
			"/run/current-system/sw/bin/bash",
			"-lc",
			"exit 37",
		},
		Env: []string{
			"HOME=/workspace",
			"LOGNAME=spindle-workflow",
			"PATH=/run/current-system/sw/bin:/usr/bin:/bin",
			"USER=spindle-workflow",
			"CI=true",
			"WF_ONLY=workflow",
			"STEP_ONLY=step",
			"OVERRIDE=step",
		},
		Cwd:  "/workspace/repo",
		User: "spindle-workflow",
	}
	if exitCode, err := agent.Exec(ctx, AgentExec{
		ID:        "failed-step",
		ExecStart: failedStep,
		Stdout:    io.Discard,
		Stderr:    io.Discard,
	}); err != nil || exitCode != 37 {
		t.Fatalf("failed step: exit code %d, error %v", exitCode, err)
	}
	if failedStep.TimeoutSeconds == 0 {
		t.Fatal("AgentSession.Exec did not retain the effective step timeout")
	}

	wid := models.WorkflowId{
		PipelineId: models.PipelineId{Knot: "knot.test", Rkey: "pipeline"},
		Name:       "debug-idle-e2e.yml",
	}
	engine := &Engine{l: logger, debug: make(map[string]debugTarget)}
	engine.registerDebugTarget(wid, debugTarget{
		cid:       vm.CID(),
		agent:     agent,
		connected: make(chan struct{}),
		released:  make(chan struct{}),
	})
	handle := newDebugHandle(wid)

	sshListener, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	sshServer := &gliderssh.Server{
		Handler: engine.debugHandle,
		PublicKeyHandler: func(gliderssh.Context, gliderssh.PublicKey) bool {
			return true
		},
	}
	hostKeyPath := filepath.Join(workDir, "host-key")
	if err := ensureDebugHostKey(hostKeyPath); err != nil {
		t.Fatal(err)
	}
	if err := sshServer.SetOption(gliderssh.HostKeyFile(hostKeyPath)); err != nil {
		t.Fatal(err)
	}
	serverDone := make(chan error, 1)
	go func() { serverDone <- sshServer.Serve(sshListener) }()
	defer func() {
		_ = sshServer.Close()
		if err := <-serverDone; err != nil && err != gliderssh.ErrServerClosed {
			t.Errorf("debug SSH server: %v", err)
		}
	}()

	_, privateKey, err := ed25519.GenerateKey(rand.Reader)
	if err != nil {
		t.Fatal(err)
	}
	privateKeyBlock, err := gossh.MarshalPrivateKey(privateKey, "")
	if err != nil {
		t.Fatal(err)
	}
	clientKeyPath := filepath.Join(workDir, "client-key")
	if err := os.WriteFile(clientKeyPath, pem.EncodeToMemory(privateKeyBlock), 0o600); err != nil {
		t.Fatal(err)
	}

	host, port, err := net.SplitHostPort(sshListener.Addr().String())
	if err != nil {
		t.Fatal(err)
	}
	target := handle + "@" + host
	controlPath := filepath.Join(workDir, "control")
	commonArgs := []string{
		"-o", "BatchMode=yes",
		"-o", "IdentitiesOnly=yes",
		"-o", "StrictHostKeyChecking=no",
		"-o", "UserKnownHostsFile=/dev/null",
		"-i", clientKeyPath,
		"-p", port,
	}
	missingContextArgs := append([]string{"-S", "none", "-tt"}, commonArgs...)
	missingContextArgs = append(missingContextArgs, target)
	missingContextOutput, err := exec.CommandContext(ctx, "ssh", missingContextArgs...).CombinedOutput()
	if err == nil {
		t.Fatalf("debug shell without failed step context exited successfully\noutput:\n%s", missingContextOutput)
	}
	wantMissingContext := "error: debug shell failed: debug shell request is missing failed step context"
	if !strings.Contains(string(missingContextOutput), wantMissingContext) {
		t.Fatalf(
			"missing-context error did not reach SSH client\nwant: %s\noutput:\n%s",
			wantMissingContext,
			missingContextOutput,
		)
	}
	if _, ok := engine.lookupDebugTarget(handle); ok {
		t.Fatal("debug target remained registered after missing-context failure")
	}
	engine.registerDebugTarget(wid, debugTarget{
		cid:        vm.CID(),
		agent:      agent,
		failedStep: failedStep,
		connected:  make(chan struct{}),
		released:   make(chan struct{}),
	})
	masterArgs := append([]string{"-M", "-N", "-S", controlPath}, commonArgs...)
	masterArgs = append(masterArgs, target)
	master := exec.CommandContext(ctx, "ssh", masterArgs...)
	var masterOutput bytes.Buffer
	master.Stdout = &masterOutput
	master.Stderr = &masterOutput
	if err := master.Start(); err != nil {
		t.Fatal(err)
	}
	masterDone := make(chan error, 1)
	go func() { masterDone <- master.Wait() }()
	masterStopped := false
	defer func() {
		if !masterStopped {
			_ = master.Process.Kill()
			<-masterDone
		}
	}()

	controlDeadline := time.Now().Add(10 * time.Second)
	for {
		if _, err := os.Stat(controlPath); err == nil {
			break
		}
		select {
		case err := <-masterDone:
			masterStopped = true
			t.Fatalf("SSH control master exited before creating its socket: %v\noutput:\n%s", err, masterOutput.String())
		case <-time.After(10 * time.Millisecond):
		}
		if time.Now().After(controlDeadline) {
			t.Fatalf("SSH control master did not create %s\noutput:\n%s", controlPath, masterOutput.String())
		}
	}

	sessionArgs := append([]string{"-S", controlPath, "-tt"}, commonArgs...)
	sessionArgs = append(sessionArgs, target)
	session := exec.CommandContext(ctx, "ssh", sessionArgs...)
	stdin, err := session.StdinPipe()
	if err != nil {
		t.Fatal(err)
	}
	var output bytes.Buffer
	session.Stdout = &output
	session.Stderr = &output
	session.Env = append(os.Environ(), "TERM=xterm-idle-e2e")
	if err := session.Start(); err != nil {
		t.Fatal(err)
	}
	waitDone := make(chan error, 1)
	go func() { waitDone <- session.Wait() }()

	const idleTime = 5 * time.Second
	select {
	case err := <-waitDone:
		t.Fatalf("debug shell exited while idle before %s: %v\noutput:\n%s", idleTime, err, output.String())
	case <-time.After(idleTime):
	}

	probe := `printf '__DEBUG_IDLE_E2E__|CI=%s|WF=%s|STEP=%s|OVERRIDE=%s|XDG=%s|PWD=%s|USER=%s|LOGNAME=%s|HOME=%s|TERM=%s|\n' "$CI" "$WF_ONLY" "$STEP_ONLY" "$OVERRIDE" "$XDG_CACHE_HOME" "$PWD" "$USER" "$LOGNAME" "$HOME" "$TERM"; exit` + "\n"
	if _, err := io.WriteString(stdin, probe); err != nil {
		t.Fatal(err)
	}

	select {
	case err := <-waitDone:
		if err != nil {
			t.Fatalf("debug shell exit: %v\noutput:\n%s", err, output.String())
		}
	case <-time.After(15 * time.Second):
		t.Fatal("debug shell did not exit after probe")
	}

	if err := master.Process.Kill(); err != nil {
		t.Fatal(err)
	}
	<-masterDone
	masterStopped = true

	want := "__DEBUG_IDLE_E2E__|CI=true|WF=workflow|STEP=step|OVERRIDE=step|XDG=/workspace/.cache|PWD=/workspace/repo|USER=spindle-workflow|LOGNAME=spindle-workflow|HOME=/workspace|TERM=xterm-idle-e2e|"
	if !strings.Contains(output.String(), want) {
		t.Fatalf("debug shell did not preserve the failed step context\nwant: %s\noutput:\n%s", want, output.String())
	}
	if _, ok := engine.lookupDebugTarget(handle); ok {
		t.Fatal("debug target remained registered after shell exit")
	}
}