package daemoncontrol

import (
	"errors"
	"testing"
	"time"

	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/AgentTeam/x/proc"
)

type memoryReceiptRewriter struct {
	receipt Receipt
	err     error
	writes  int
}

func (s *memoryReceiptRewriter) RewriteMatching(expected, replacement Receipt) error {
	if s.err != nil {
		return s.err
	}
	if !sameReceipt(s.receipt, expected) {
		return ErrReceiptMismatch
	}
	s.receipt = replacement
	s.writes++
	return nil
}

func TestControllerIdentityRequiresHeaderCapability(t *testing.T) {
	receipt := validReceipt()
	store := &memoryReceiptRewriter{receipt: receipt}
	controller := newControllerForTest(t, receipt, store, map[int]proc.ProcessInfo{
		receipt.Owner.PID: ownerProcess(receipt.Owner),
	})

	_, err := controller.Identity("")
	require.ErrorIs(t, err, ErrControlUnauthorized)
	_, err = controller.Identity("wrong")
	require.ErrorIs(t, err, ErrControlUnauthorized)

	identity, err := controller.Identity(receipt.ControlToken)
	require.NoError(t, err)
	require.Equal(t, receipt.InstanceID, identity.InstanceID)
	require.Equal(t, receipt.DesktopAPIVersion, identity.DesktopAPIVersion)
	require.Equal(t, receipt.PID, identity.PID)
	require.Equal(t, receipt.BaseURL, identity.BaseURL)
	require.Equal(t, receipt.Owner.Kind, identity.OwnerKind)
	require.False(t, identity.Draining)
}

func TestControllerMutationRejectsWrongCapabilityInstanceAndOwner(t *testing.T) {
	receipt := validReceipt()
	store := &memoryReceiptRewriter{receipt: receipt}
	controller := newControllerForTest(t, receipt, store, map[int]proc.ProcessInfo{
		receipt.Owner.PID: ownerProcess(receipt.Owner),
	})
	valid := credentialsFor(receipt)

	tests := []struct {
		name   string
		mutate func(*Credentials)
	}{
		{"token", func(c *Credentials) { c.ControlToken = "wrong" }},
		{"instance", func(c *Credentials) { c.InstanceID = "wrong" }},
		{"owner id", func(c *Credentials) { c.OwnerID = "wrong" }},
		{"owner pid", func(c *Credentials) { c.Owner.PID++ }},
		{"owner start", func(c *Credentials) { c.Owner.StartedAt = c.Owner.StartedAt.Add(time.Second) }},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			credentials := valid
			tt.mutate(&credentials)
			_, err := controller.BeginShutdown(credentials)
			require.ErrorIs(t, err, ErrControlUnauthorized)
		})
	}
}

func TestControllerAdoptRefusesActiveOrNonElectronOwner(t *testing.T) {
	receipt := validReceipt()
	store := &memoryReceiptRewriter{receipt: receipt}
	controller := newControllerForTest(t, receipt, store, map[int]proc.ProcessInfo{
		receipt.Owner.PID: ownerProcess(receipt.Owner),
	})

	_, err := controller.Adopt(credentialsFor(receipt), nextElectronOwner(receipt.Owner.PID+1))
	require.ErrorIs(t, err, ErrOwnerActive)
	require.Zero(t, store.writes)

	for _, kind := range []OwnerKind{OwnerKindCLI, OwnerKindService} {
		t.Run(string(kind), func(t *testing.T) {
			externalReceipt := receipt
			externalReceipt.Owner.Kind = kind
			externalStore := &memoryReceiptRewriter{receipt: externalReceipt}
			externalController := newControllerForTest(t, externalReceipt, externalStore, nil)
			_, err := externalController.Adopt(
				credentialsFor(externalReceipt),
				nextElectronOwner(externalReceipt.Owner.PID+1),
			)
			require.ErrorIs(t, err, ErrOwnerUnsupported)
		})
	}
}

func TestControllerAdoptDeadElectronRotatesOwnerAndTokenAtomically(t *testing.T) {
	receipt := validReceipt()
	next := nextElectronOwner(receipt.Owner.PID + 1)
	store := &memoryReceiptRewriter{receipt: receipt}
	controller := newControllerForTest(t, receipt, store, map[int]proc.ProcessInfo{
		next.PID: ownerProcess(next),
	})
	controller.newToken = func() (string, error) { return "rotated-control-token", nil }

	adoption, err := controller.Adopt(credentialsFor(receipt), next)
	require.NoError(t, err)
	require.Equal(t, next, adoption.Owner)
	require.Equal(t, "rotated-control-token", adoption.ControlToken)
	require.Equal(t, 1, store.writes)
	require.Equal(t, next, store.receipt.Owner)
	require.Equal(t, "rotated-control-token", store.receipt.ControlToken)

	_, err = controller.Identity(receipt.ControlToken)
	require.ErrorIs(t, err, ErrControlUnauthorized)
	_, err = controller.Identity(adoption.ControlToken)
	require.NoError(t, err)
}

func TestControllerAdoptAcceptsProvenPIDReuse(t *testing.T) {
	receipt := validReceipt()
	next := nextElectronOwner(receipt.Owner.PID + 1)
	processes := map[int]proc.ProcessInfo{
		receipt.Owner.PID: {
			PID:       receipt.Owner.PID,
			StartedAt: receipt.Owner.StartedAt.Add(time.Hour),
		},
		next.PID: ownerProcess(next),
	}
	store := &memoryReceiptRewriter{receipt: receipt}
	controller := newControllerForTest(t, receipt, store, processes)

	_, err := controller.Adopt(credentialsFor(receipt), next)
	require.NoError(t, err)
	require.Equal(t, 1, store.writes)
}

func TestControllerAdoptRejectsAmbiguousPreviousOwner(t *testing.T) {
	receipt := validReceipt()
	next := nextElectronOwner(receipt.Owner.PID + 1)
	store := &memoryReceiptRewriter{receipt: receipt}
	controller := newControllerForTest(t, receipt, store, map[int]proc.ProcessInfo{
		next.PID: ownerProcess(next),
	})
	controller.pidAlive = func(pid int) bool { return pid == receipt.Owner.PID }

	_, err := controller.Adopt(credentialsFor(receipt), next)
	require.ErrorIs(t, err, ErrOwnerIdentity)
	require.Zero(t, store.writes)
}

func TestControllerAdoptReceiptFailureLeavesOldOwnershipValid(t *testing.T) {
	receipt := validReceipt()
	next := nextElectronOwner(receipt.Owner.PID + 1)
	store := &memoryReceiptRewriter{
		receipt: receipt,
		err:     errors.New("disk full"),
	}
	controller := newControllerForTest(t, receipt, store, map[int]proc.ProcessInfo{
		next.PID: ownerProcess(next),
	})
	controller.newToken = func() (string, error) { return "rotated-control-token", nil }

	_, err := controller.Adopt(credentialsFor(receipt), next)
	require.ErrorContains(t, err, "disk full")
	require.Equal(t, receipt, controller.CurrentReceipt())
	_, err = controller.Identity(receipt.ControlToken)
	require.NoError(t, err)
	_, err = controller.Identity("rotated-control-token")
	require.ErrorIs(t, err, ErrControlUnauthorized)
}

func TestControllerShutdownOwnerOnlyAndDraining(t *testing.T) {
	receipt := validReceipt()
	store := &memoryReceiptRewriter{receipt: receipt}
	controller := newControllerForTest(t, receipt, store, map[int]proc.ProcessInfo{
		receipt.Owner.PID: ownerProcess(receipt.Owner),
	})

	already, err := controller.BeginShutdown(credentialsFor(receipt))
	require.NoError(t, err)
	require.False(t, already)
	require.True(t, controller.PublicIdentity().Draining)

	already, err = controller.BeginShutdown(credentialsFor(receipt))
	require.NoError(t, err)
	require.True(t, already)

	serviceReceipt := receipt
	serviceReceipt.Owner.Kind = OwnerKindService
	store = &memoryReceiptRewriter{receipt: serviceReceipt}
	controller = newControllerForTest(t, serviceReceipt, store, map[int]proc.ProcessInfo{
		serviceReceipt.Owner.PID: ownerProcess(serviceReceipt.Owner),
	})
	_, err = controller.BeginShutdown(credentialsFor(serviceReceipt))
	require.ErrorIs(t, err, ErrOwnerUnsupported)
}

func newControllerForTest(
	t *testing.T,
	receipt Receipt,
	store ReceiptRewriter,
	processes map[int]proc.ProcessInfo,
) *Controller {
	t.Helper()
	controller, err := NewController(ControllerOptions{
		Receipt: receipt,
		Store:   store,
		LookupProcess: func(pid int) (proc.ProcessInfo, bool) {
			process, ok := processes[pid]
			return process, ok
		},
		PIDAlive: func(int) bool { return false },
		NewToken: func() (string, error) {
			return "new-control-token", nil
		},
	})
	require.NoError(t, err)
	return controller
}

func credentialsFor(receipt Receipt) Credentials {
	return Credentials{
		ControlToken: receipt.ControlToken,
		InstanceID:   receipt.InstanceID,
		OwnerID:      receipt.Owner.ID,
		Owner:        receipt.Owner.ProcessIdentity,
	}
}

func ownerProcess(owner Owner) proc.ProcessInfo {
	return proc.ProcessInfo{
		PID:       owner.PID,
		StartedAt: owner.StartedAt,
		StartID:   owner.StartID,
	}
}

func nextElectronOwner(pid int) Owner {
	return Owner{
		Kind: OwnerKindElectron,
		ID:   "next-electron-owner",
		ProcessIdentity: ProcessIdentity{
			PID:       pid,
			StartedAt: time.Unix(1710001000, 0).UTC(),
			StartID:   "linux:boot-a:1000",
		},
	}
}