package conpty

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"net"
	"os"
	"strings"
	"testing"
	"time"

	"github.com/aoagents/agent-orchestrator/backend/internal/adapters/runtime/conpty/ptyregistry"
	"github.com/aoagents/agent-orchestrator/backend/internal/domain"
	"github.com/aoagents/agent-orchestrator/backend/internal/ports"
)

// livePID returns a PID that is guaranteed to be alive (the current process).
// Using this as the fake pty-host PID means ptyregistry.List() will not prune
// the entry during tests. Do NOT use this for the Destroy test: Destroy calls
// Kill on the pid, so use deadPID() there instead.
func livePID() int { return os.Getpid() }

// deadPID returns a PID that is guaranteed to be dead (no process). This is
// used in Destroy tests so the force-kill step is a safe no-op.
// ponytail: PID 2147483647 (MaxInt32) is never a real process; signal-0 returns ESRCH.
func deadPID() int { return 2147483647 }

// ---------------------------------------------------------------------------
// Test harness: in-process pty-host backed by a fakePTY.
// ---------------------------------------------------------------------------

// inProcHost starts a Serve engine with a fakePTY on a real 127.0.0.1:0
// listener and returns a fake spawner that returns that addr and a fake pid.
// The caller must call cleanup() to shut down the host.
type inProcHost struct {
	addr   string
	pid    int
	pty    *fakePTY
	ring   *Ring
	cancel context.CancelFunc
	done   chan error
	ln     net.Listener
}

func startInProcHost(t *testing.T, sessionID string, fakePID int) *inProcHost {
	t.Helper()
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("listen: %v", err)
	}
	pty := newFakePTY(fakePID)
	ring := NewRing()
	ctx, cancel := context.WithCancel(context.Background())
	done := make(chan error, 1)
	go func() {
		done <- Serve(ctx, ServeConfig{
			SessionID: sessionID,
			Listener:  ln,
			PTY:       pty,
			Ring:      ring,
		})
	}()
	return &inProcHost{
		addr:   ln.Addr().String(),
		pid:    fakePID,
		pty:    pty,
		ring:   ring,
		cancel: cancel,
		done:   done,
		ln:     ln,
	}
}

func (h *inProcHost) cleanup(t *testing.T) {
	t.Helper()
	h.cancel()
	select {
	case <-h.done:
	case <-time.After(2 * time.Second):
		t.Log("warning: inProcHost did not stop within 2s")
	}
}

// fakeSpawnerFor returns a hostSpawner that starts an in-process host for a
// single session ID and records which sessions have been spawned.
// The returned map maps sessionID -> *inProcHost for test inspection.
func fakeSpawnerFor(t *testing.T, hosts map[string]*inProcHost, fakePID int) hostSpawner {
	t.Helper()
	return func(ctx context.Context, sessionID, cwd string, argv []string, env map[string]string) (string, int, error) {
		h := startInProcHost(t, sessionID, fakePID)
		if hosts != nil {
			hosts[sessionID] = h
		}
		return h.addr, h.pid, nil
	}
}

// ---------------------------------------------------------------------------
// Redirect ptyregistry to a temp HOME so tests don't pollute ~/.ao
// ---------------------------------------------------------------------------

func isolateRegistry(t *testing.T) {
	t.Helper()
	dir := t.TempDir()
	t.Setenv("HOME", dir)
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

// TestCreate_RegistersSession verifies Create returns {ID: sessionID}, writes
// to the in-memory map, and registers in the ptyregistry.
func TestCreate_RegistersSession(t *testing.T) {
	isolateRegistry(t)
	hosts := map[string]*inProcHost{}
	rt := New(Options{Spawner: fakeSpawnerFor(t, hosts, livePID())})

	ctx := context.Background()
	handle, err := rt.Create(ctx, ports.RuntimeConfig{
		SessionID:     domain.SessionID("sess-abc"),
		WorkspacePath: "/tmp/workspace",
		Argv:          []string{"claude-code"},
	})
	if err != nil {
		t.Fatalf("Create: %v", err)
	}

	if handle.ID != "sess-abc" {
		t.Fatalf("handle.ID = %q, want %q", handle.ID, "sess-abc")
	}

	// In-memory map must have the entry.
	rt.mu.Lock()
	sess := rt.sessions["sess-abc"]
	rt.mu.Unlock()
	if sess == nil {
		t.Fatal("session not in in-memory map after Create")
	}

	// Registry must have the entry.
	entries, err := ptyregistry.List()
	if err != nil {
		t.Fatalf("List: %v", err)
	}
	var found bool
	for _, e := range entries {
		if e.SessionID == "sess-abc" {
			found = true
		}
	}
	if !found {
		t.Fatal("session not in registry after Create")
	}

	hosts["sess-abc"].cleanup(t)
}

// TestCreate_DuplicateErrors verifies a second Create for the same session id fails.
func TestCreate_DuplicateErrors(t *testing.T) {
	isolateRegistry(t)
	hosts := map[string]*inProcHost{}
	rt := New(Options{Spawner: fakeSpawnerFor(t, hosts, livePID())})
	ctx := context.Background()

	if _, err := rt.Create(ctx, ports.RuntimeConfig{
		SessionID:     "sess-dup",
		WorkspacePath: "/tmp/w",
		Argv:          []string{"sh"},
	}); err != nil {
		t.Fatalf("first Create: %v", err)
	}

	_, err := rt.Create(ctx, ports.RuntimeConfig{
		SessionID:     "sess-dup",
		WorkspacePath: "/tmp/w",
		Argv:          []string{"sh"},
	})
	if err == nil {
		t.Fatal("expected error on duplicate Create, got nil")
	}
	if !strings.Contains(err.Error(), "already exists") {
		t.Fatalf("error %q should contain 'already exists'", err.Error())
	}

	hosts["sess-dup"].cleanup(t)
}

// TestCreate_InvalidIDErrors verifies Create rejects invalid session ids.
func TestCreate_InvalidIDErrors(t *testing.T) {
	isolateRegistry(t)
	rt := New(Options{Spawner: fakeSpawnerFor(t, nil, livePID())})
	ctx := context.Background()

	for _, bad := range []string{"", "has space", "has/slash", "has.dot"} {
		_, err := rt.Create(ctx, ports.RuntimeConfig{
			SessionID:     domain.SessionID(bad),
			WorkspacePath: "/tmp/w",
			Argv:          []string{"sh"},
		})
		if err == nil {
			t.Fatalf("Create(%q): expected error for invalid id, got nil", bad)
		}
	}
}

// TestSendMessage_DeliversChunkedTextAndEnter verifies clientSendMessage sends
// the text + "\r" to the fakePTY input.
func TestSendMessage_DeliversChunkedTextAndEnter(t *testing.T) {
	isolateRegistry(t)
	hosts := map[string]*inProcHost{}
	rt := New(Options{Spawner: fakeSpawnerFor(t, hosts, livePID())})
	ctx := context.Background()

	handle, err := rt.Create(ctx, ports.RuntimeConfig{
		SessionID:     "sess-sm",
		WorkspacePath: "/tmp/w",
		Argv:          []string{"sh"},
	})
	if err != nil {
		t.Fatalf("Create: %v", err)
	}
	h := hosts["sess-sm"]
	defer h.cleanup(t)

	msg := "hello world"
	// Collect PTY input in background.
	inputC := make(chan []byte, 4)
	go func() {
		buf := make([]byte, 1024)
		for {
			n, err := h.pty.inR.Read(buf)
			if n > 0 {
				cp := make([]byte, n)
				copy(cp, buf[:n])
				inputC <- cp
			}
			if err != nil {
				return
			}
		}
	}()

	if err := rt.SendMessage(ctx, handle, msg); err != nil {
		t.Fatalf("SendMessage: %v", err)
	}

	// Collect all received bytes within 2s.
	var received []byte
	deadline := time.After(2 * time.Second)
	// Expect at least msg + "\r".
	for !bytes.Contains(received, []byte("\r")) {
		select {
		case chunk := <-inputC:
			received = append(received, chunk...)
		case <-deadline:
			t.Fatalf("timeout waiting for PTY input; got %q so far", received)
		}
	}

	if !bytes.HasPrefix(received, []byte(msg)) {
		t.Fatalf("PTY input = %q, want prefix %q then \\r", received, msg)
	}
	if !bytes.Contains(received, []byte("\r")) {
		t.Fatalf("PTY input = %q, missing trailing \\r", received)
	}
}

// TestSendMessage_LargeMessageChunked verifies a message > 512 runes is
// delivered correctly (host receives full text + "\r").
func TestSendMessage_LargeMessageChunked(t *testing.T) {
	isolateRegistry(t)
	hosts := map[string]*inProcHost{}
	rt := New(Options{Spawner: fakeSpawnerFor(t, hosts, livePID())})
	ctx := context.Background()

	handle, _ := rt.Create(ctx, ports.RuntimeConfig{
		SessionID:     "sess-lg",
		WorkspacePath: "/tmp/w",
		Argv:          []string{"sh"},
	})
	h := hosts["sess-lg"]
	defer h.cleanup(t)

	// Build a message longer than 512 runes (use multi-byte runes to test
	// rune-boundary splitting).
	var sb strings.Builder
	for i := 0; i < 600; i++ {
		sb.WriteRune('A' + rune(i%26))
	}
	msg := sb.String()

	inputDone := make(chan []byte, 1)
	go func() {
		// Read until we see "\r".
		var acc []byte
		buf := make([]byte, 4096)
		for {
			n, err := h.pty.inR.Read(buf)
			if n > 0 {
				acc = append(acc, buf[:n]...)
			}
			if bytes.Contains(acc, []byte("\r")) {
				inputDone <- acc
				return
			}
			if err != nil {
				inputDone <- acc
				return
			}
		}
	}()

	if err := rt.SendMessage(ctx, handle, msg); err != nil {
		t.Fatalf("SendMessage: %v", err)
	}

	select {
	case got := <-inputDone:
		// Strip trailing \r for comparison.
		trimmed := strings.TrimSuffix(string(got), "\r")
		if trimmed != msg {
			t.Fatalf("PTY received %d chars, want %d\ngot:  %q\nwant: %q", len(trimmed), len(msg), trimmed[:min(50, len(trimmed))], msg[:min(50, len(msg))])
		}
	case <-time.After(5 * time.Second):
		t.Fatal("timeout waiting for large message delivery")
	}
}

// TestGetOutput_ReturnsRingTail verifies GetOutput returns the ring's tail.
func TestGetOutput_ReturnsRingTail(t *testing.T) {
	isolateRegistry(t)
	hosts := map[string]*inProcHost{}
	rt := New(Options{Spawner: fakeSpawnerFor(t, hosts, livePID())})
	ctx := context.Background()

	handle, _ := rt.Create(ctx, ports.RuntimeConfig{
		SessionID:     "sess-go",
		WorkspacePath: "/tmp/w",
		Argv:          []string{"sh"},
	})
	h := hosts["sess-go"]
	defer h.cleanup(t)

	// Seed the ring.
	h.ring.Append([]byte("line1\nline2\nline3\n"))

	text, err := rt.GetOutput(ctx, handle, 2)
	if err != nil {
		t.Fatalf("GetOutput: %v", err)
	}
	want := h.ring.Tail(2)
	if text != want {
		t.Fatalf("GetOutput = %q, want %q", text, want)
	}
}

// TestIsAlive_TrueWhileServing_FalseAfterClose verifies IsAlive returns true
// while the host listens and false after its listener is closed.
func TestIsAlive_TrueWhileServing_FalseAfterClose(t *testing.T) {
	isolateRegistry(t)
	hosts := map[string]*inProcHost{}
	rt := New(Options{Spawner: fakeSpawnerFor(t, hosts, livePID())})
	ctx := context.Background()

	handle, _ := rt.Create(ctx, ports.RuntimeConfig{
		SessionID:     "sess-ia",
		WorkspacePath: "/tmp/w",
		Argv:          []string{"sh"},
	})
	h := hosts["sess-ia"]

	alive, err := rt.IsAlive(ctx, handle)
	if err != nil {
		t.Fatalf("IsAlive: %v", err)
	}
	if !alive {
		t.Fatal("expected IsAlive=true while serving")
	}

	// Shut down the host.
	h.cancel()
	<-h.done

	// Give the listener a moment to close.
	time.Sleep(100 * time.Millisecond)

	alive2, err2 := rt.IsAlive(ctx, handle)
	if err2 != nil {
		t.Fatalf("IsAlive after close: %v", err2)
	}
	if alive2 {
		t.Fatal("expected IsAlive=false after host closed")
	}
}

// TestIsAlive_FalseForUnknownSession verifies IsAlive returns (false, nil) for
// a session not in the map or registry.
func TestIsAlive_FalseForUnknownSession(t *testing.T) {
	isolateRegistry(t)
	rt := New(Options{Spawner: fakeSpawnerFor(t, nil, livePID())})
	ctx := context.Background()

	alive, err := rt.IsAlive(ctx, ports.RuntimeHandle{ID: "ghost-session"})
	if err != nil {
		t.Fatalf("IsAlive: unexpected error: %v", err)
	}
	if alive {
		t.Fatal("expected IsAlive=false for unknown session")
	}
}

// TestDestroy_KillsHostAndCleansUp verifies Destroy triggers clientKill,
// removes the map + registry entry, and is idempotent on second call.
// Uses deadPID() so the force-kill step is a safe no-op (the fake pty-host
// has no real OS process; clientKill already shut it down via the loopback).
func TestDestroy_KillsHostAndCleansUp(t *testing.T) {
	isolateRegistry(t)
	hosts := map[string]*inProcHost{}
	rt := New(Options{Spawner: fakeSpawnerFor(t, hosts, deadPID())})
	ctx := context.Background()

	handle, err := rt.Create(ctx, ports.RuntimeConfig{
		SessionID:     "sess-destroy",
		WorkspacePath: "/tmp/w",
		Argv:          []string{"sh"},
	})
	if err != nil {
		t.Fatalf("Create: %v", err)
	}
	h := hosts["sess-destroy"]

	// Destroy should succeed.
	if err := rt.Destroy(ctx, handle); err != nil {
		t.Fatalf("Destroy: %v", err)
	}

	// Wait for Serve to stop (clientKill triggers shutdown).
	select {
	case <-h.done:
	case <-time.After(3 * time.Second):
		t.Fatal("host did not stop after Destroy")
	}

	// fakePTY.Close must have been called.
	h.pty.closeMu.Lock()
	closed := h.pty.closed
	h.pty.closeMu.Unlock()
	if !closed {
		t.Fatal("expected fakePTY.Close() after Destroy")
	}

	// Map entry must be gone.
	rt.mu.Lock()
	_, exists := rt.sessions["sess-destroy"]
	rt.mu.Unlock()
	if exists {
		t.Fatal("expected map entry removed after Destroy")
	}

	// Registry entry must be gone.
	entries, _ := ptyregistry.List()
	for _, e := range entries {
		if e.SessionID == "sess-destroy" {
			t.Fatal("expected registry entry removed after Destroy")
		}
	}

	// Second Destroy must be idempotent (returns nil).
	if err := rt.Destroy(ctx, handle); err != nil {
		t.Fatalf("second Destroy: expected nil, got %v", err)
	}
}

// TestResolveViaRegistry verifies that with an empty in-memory map but a
// registry entry pointing at a live in-process host, IsAlive and SendMessage
// still work (simulates a daemon restart).
func TestResolveViaRegistry(t *testing.T) {
	isolateRegistry(t)

	// Start a host directly (not through Create) to simulate a pre-existing
	// pty-host from a previous daemon run. Use the current process PID so
	// ptyregistry.List() does not prune the entry as dead.
	h := startInProcHost(t, "sess-reg", livePID())
	defer h.cleanup(t)

	// Manually register the host in the registry.
	err := ptyregistry.Register(ptyregistry.Entry{
		SessionID:    "sess-reg",
		PtyHostPID:   h.pid,
		PipePath:     h.addr, // addr stored in PipePath field
		RegisteredAt: fmt.Sprintf("%d", time.Now().Unix()),
	})
	if err != nil {
		t.Fatalf("Register: %v", err)
	}

	// Create a Runtime with an empty in-memory map (simulates daemon restart).
	rt := New(Options{Spawner: fakeSpawnerFor(t, nil, livePID())})
	ctx := context.Background()

	// IsAlive must work via registry resolution.
	alive, err := rt.IsAlive(ctx, ports.RuntimeHandle{ID: "sess-reg"})
	if err != nil {
		t.Fatalf("IsAlive via registry: %v", err)
	}
	if !alive {
		t.Fatal("expected IsAlive=true via registry resolution")
	}

	// SendMessage must work via registry resolution.
	inputC := make(chan []byte, 4)
	go func() {
		buf := make([]byte, 512)
		for {
			n, err := h.pty.inR.Read(buf)
			if n > 0 {
				cp := make([]byte, n)
				copy(cp, buf[:n])
				inputC <- cp
			}
			if err != nil {
				return
			}
		}
	}()

	if err := rt.SendMessage(ctx, ports.RuntimeHandle{ID: "sess-reg"}, "ping"); err != nil {
		t.Fatalf("SendMessage via registry: %v", err)
	}

	// Collect PTY input.
	var received []byte
	deadline := time.After(3 * time.Second)
	for !bytes.Contains(received, []byte("\r")) {
		select {
		case chunk := <-inputC:
			received = append(received, chunk...)
		case <-deadline:
			t.Fatalf("timeout waiting for PTY input via registry; got %q", received)
		}
	}
	if !bytes.Contains(received, []byte("ping")) {
		t.Fatalf("PTY did not receive 'ping'; got %q", received)
	}
}

// ---------------------------------------------------------------------------
// Unit tests for client helpers (dial a fresh in-proc host directly).
// ---------------------------------------------------------------------------

// TestClientGetOutput_TimesOutReturnsEmpty verifies clientGetOutput returns ""
// (no error) if no response arrives within the timeout. We test the happy path
// instead (timeout path would require a non-responding server).
func TestClientGetOutput_HappyPath(t *testing.T) {
	f := startServe(t, 3001)
	defer f.cancel()

	f.ring.Append([]byte("alpha\nbeta\ngamma\n"))

	text, err := clientGetOutput(f.addr, 2)
	if err != nil {
		t.Fatalf("clientGetOutput: %v", err)
	}
	want := f.ring.Tail(2)
	if text != want {
		t.Fatalf("clientGetOutput = %q, want %q", text, want)
	}
}

// TestClientIsAlive_TrueAndFalse verifies clientIsAlive returns (true, nil) for
// a live host and (false, nil) for a refused address (definitively gone).
func TestClientIsAlive_TrueAndFalse(t *testing.T) {
	f := startServe(t, 3002)
	defer f.cancel()

	if alive, err := clientIsAlive(f.addr); err != nil || !alive {
		t.Fatalf("clientIsAlive(live) = (%v, %v), want (true, nil)", alive, err)
	}

	f.cancel()
	// Wait for listener to close.
	select {
	case <-f.done:
	case <-time.After(2 * time.Second):
	}
	time.Sleep(50 * time.Millisecond)

	// After close the OS refuses the connection on the freed port -> gone.
	if alive, err := clientIsAlive(f.addr); alive || err != nil {
		t.Fatalf("clientIsAlive(closed) = (%v, %v), want (false, nil)", alive, err)
	}
}

// TestIsAlive_RefusedIsGone_TimeoutIsTransient is the reaper-safety regression
// test. It asserts the dead-vs-transient split that keeps a single transient
// loopback hiccup from spuriously reaping a live idle session:
//
//	(a) a resolved-but-REFUSED host -> IsAlive == (false, nil)  [ProbeDead]
//	(b) a resolved host whose probe TIMES OUT -> (false, non-nil) [ProbeFailed]
func TestIsAlive_RefusedIsGone_TimeoutIsTransient(t *testing.T) {
	isolateRegistry(t)

	// (a) Refused: bind+close a listener to obtain a port nothing listens on.
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("listen: %v", err)
	}
	refusedAddr := ln.Addr().String()
	_ = ln.Close()

	rtRefused := New(Options{Spawner: fakeSpawnerFor(t, nil, livePID())})
	rtRefused.mu.Lock()
	rtRefused.sessions["gone"] = &hostSession{addr: refusedAddr, pid: livePID()}
	rtRefused.mu.Unlock()

	alive, err := rtRefused.IsAlive(context.Background(), ports.RuntimeHandle{ID: "gone"})
	if alive || err != nil {
		t.Fatalf("IsAlive(refused) = (%v, %v), want (false, nil) definitively gone", alive, err)
	}

	// (b) Transient timeout: a listener that Accepts but never replies. The
	// short isAliveTimeout read deadline fires before any STATUS_RES arrives,
	// which must surface as a non-nil (transient) error, not a death.
	silent, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("listen silent: %v", err)
	}
	defer silent.Close()
	go func() {
		for {
			c, err := silent.Accept()
			if err != nil {
				return
			}
			// Hold the connection open without ever sending a STATUS_RES.
			go func(c net.Conn) {
				time.Sleep(isAliveTimeout + time.Second)
				_ = c.Close()
			}(c)
		}
	}()

	rtSilent := New(Options{Spawner: fakeSpawnerFor(t, nil, livePID())})
	rtSilent.mu.Lock()
	rtSilent.sessions["stuck"] = &hostSession{addr: silent.Addr().String(), pid: livePID()}
	rtSilent.mu.Unlock()

	alive, err = rtSilent.IsAlive(context.Background(), ports.RuntimeHandle{ID: "stuck"})
	if alive {
		t.Fatalf("IsAlive(silent) alive=true, want false")
	}
	if err == nil {
		t.Fatal("IsAlive(silent) err=nil, want non-nil transient error so the reaper records ProbeFailed")
	}
}

// TestClientKill_Idempotent verifies clientKill on a dead address returns nil.
func TestClientKill_Idempotent(t *testing.T) {
	if err := clientKill("127.0.0.1:1"); err != nil {
		t.Fatalf("clientKill on unreachable addr: %v", err)
	}
}

// Ensure the packages compile (import check).
var _ = io.Discard
