package dispatcher

import (
	"context"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"runtime"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/AgentTeam/desktop/internal/agent"
	"github.com/tidbcloud/AgentTeam/desktop/internal/agent/drivercontract"
	"github.com/tidbcloud/AgentTeam/desktop/internal/agent/runtimeconfig"
	"github.com/tidbcloud/AgentTeam/desktop/internal/types"
	"github.com/tidbcloud/AgentTeam/desktop/internal/wire/manager"
)

func newPiOwnershipRetryTestSession(t *testing.T, d *Dispatcher) *interactiveSession {
	t.Helper()
	workDir := t.TempDir()
	require.NoError(t, saveSessionResumeState(workDir, sessionResumeState{
		RuntimeDeliveryOwnershipVersion: runtimeDeliveryOwnershipVersion,
		RuntimeDeliveryOwnership: []agent.DeliveryOwnershipRecord{{
			Message:    types.InboxMessage{MessageID: "pi-startup-msg", Target: "#ops", Content: "accepted"},
			State:      agent.DeliveryOwnershipQueued,
			AcceptedAt: time.Now().UTC(),
		}},
	}))
	sess := &interactiveSession{
		sessionKey:    "pi-startup-session",
		agentID:       "pi-startup-agent",
		serverAgentID: "pi-startup-agent",
		workDir:       workDir,
		runtime:       runtimeconfig.Runtime{Type: runtimeconfig.TypePi},
		binding:       conversationBinding{channelID: "ops", serverAgentID: "pi-startup-agent"},
	}
	d.sessionManager.Set(sess.sessionKey, sess)
	return sess
}

func writePiCatalogFailureBinary(t *testing.T) string {
	t.Helper()
	dir := t.TempDir()
	binary := filepath.Join(dir, "pi")
	if runtime.GOOS == "windows" {
		binary += ".cmd"
		require.NoError(t, os.WriteFile(binary, []byte("@echo off\r\necho Error: Unknown provider \"inline\" 1>&2\r\nexit /b 1\r\n"), 0o755))
		return binary
	}
	require.NoError(t, os.WriteFile(binary, []byte("#!/bin/sh\nprintf '%s\\n' 'Error: Unknown provider \"inline\"' >&2\nexit 1\n"), 0o755))
	return binary
}

func spawnPiCatalogFailure(t *testing.T) error {
	t.Helper()
	driver, err := agent.NewDriver(runtimeconfig.TypePi)
	require.NoError(t, err)
	workDir := t.TempDir()
	_, err = driver.Spawn(context.Background(), testLaunchConfig(drivercontract.LaunchConfig{AgentID: "pi-catalog-failure", RuntimeType: runtimeconfig.TypePi, Binary: writePiCatalogFailureBinary(t), WorkDir: workDir, RuntimeStateDir: t.TempDir(), PermissionMode: "full_access"}, testLaunchSurfaces{FreshPrompt: "accepted"}))
	return err
}

func TestPiSpawnCatalogFailureCircuitBreaksAndPreservesRootCause(t *testing.T) {
	d := mustNewDispatcher("http://server.test")
	defer d.Close()
	ch := &recordingChannel{id: "ops"}
	mgr := manager.New(nil, nil)
	mgr.Register(ch)
	d.serverConnections = mgr
	sess := newPiOwnershipRetryTestSession(t, d)

	spawnErr := spawnPiCatalogFailure(t)
	require.Error(t, spawnErr)
	require.Contains(t, spawnErr.Error(), "initial RPC prompt")
	require.Contains(t, spawnErr.Error(), `Unknown provider "inline"`)
	require.True(t, agent.IsDeterministicRuntimeStartupError(spawnErr))

	var attempts atomic.Int32
	var delayCalls atomic.Int32
	d.runtimeOwnershipContinuationStarter = func(*interactiveSession) (bool, error) {
		attempts.Add(1)
		return false, fmt.Errorf("must not start")
	}
	d.runtimeOwnershipContinuationDelay = func(int) time.Duration {
		delayCalls.Add(1)
		return time.Hour
	}
	d.scheduleRuntimeDeliveryOwnershipContinuationRetry(sess, spawnErr)

	require.Eventually(t, func() bool {
		state, err := loadSessionResumeState(sess.workDir)
		if err != nil || len(state.RuntimeDeliveryOwnership) != 1 {
			return false
		}
		return state.RuntimeDeliveryOwnership[0].IsFailed()
	}, time.Second, time.Millisecond)
	require.Zero(t, attempts.Load())
	require.Zero(t, delayCalls.Load())
	state, err := loadSessionResumeState(sess.workDir)
	require.NoError(t, err)
	require.Contains(t, state.RuntimeDeliveryOwnership[0].LastError, `Unknown provider "inline"`)
	require.NotContains(t, state.RuntimeDeliveryOwnership[0].LastError, "runtime_input_reconciliation_exhausted")
	require.Eventually(t, func() bool {
		return len(ch.events) > 0
	}, time.Second, time.Millisecond)
	require.Contains(t, ch.events[0].(types.RuntimeEvent).Detail, `Unknown provider "inline"`)
}

func TestPiTransientStartupRetrySucceedsWithInjectedBackoff(t *testing.T) {
	d := mustNewDispatcher("http://server.test")
	defer d.Close()
	sess := newPiOwnershipRetryTestSession(t, d)

	var attempts atomic.Int32
	var done = make(chan struct{})
	var once sync.Once
	var delaysMu sync.Mutex
	var delays []int
	d.runtimeOwnershipContinuationStarter = func(*interactiveSession) (bool, error) {
		if attempts.Add(1) == 2 {
			once.Do(func() { close(done) })
			return true, nil
		}
		return false, errors.New("temporary startup transport failure")
	}
	d.runtimeOwnershipContinuationDelay = func(attempt int) time.Duration {
		delaysMu.Lock()
		delays = append(delays, attempt)
		delaysMu.Unlock()
		return 0
	}
	d.scheduleRuntimeDeliveryOwnershipContinuationRetry(sess, errors.New("temporary startup transport failure"))
	select {
	case <-done:
	case <-time.After(time.Second):
		t.Fatal("transient startup retry did not converge")
	}
	require.Equal(t, int32(2), attempts.Load())
	delaysMu.Lock()
	require.Equal(t, []int{1, 2}, delays)
	delaysMu.Unlock()
}

func TestPiTransientStartupRetryUsesExactBudgetAndStops(t *testing.T) {
	d := mustNewDispatcher("http://server.test")
	defer d.Close()
	sess := newPiOwnershipRetryTestSession(t, d)

	var attempts atomic.Int32
	var delaysMu sync.Mutex
	var delays []int
	d.runtimeOwnershipContinuationStarter = func(*interactiveSession) (bool, error) {
		attempts.Add(1)
		return false, errors.New("temporary startup transport failure")
	}
	d.runtimeOwnershipContinuationDelay = func(attempt int) time.Duration {
		delaysMu.Lock()
		delays = append(delays, attempt)
		delaysMu.Unlock()
		return 0
	}
	d.scheduleRuntimeDeliveryOwnershipContinuationRetry(sess, errors.New("temporary startup transport failure"))
	require.Eventually(t, func() bool {
		state, err := loadSessionResumeState(sess.workDir)
		return err == nil && len(state.RuntimeDeliveryOwnership) == 1 && state.RuntimeDeliveryOwnership[0].IsFailed()
	}, time.Second, time.Millisecond)
	require.Equal(t, int32(runtimeOwnershipContinuationRetryAttempts), attempts.Load())
	delaysMu.Lock()
	require.Equal(t, []int{1, 2, 3, 4, 5, 6, 7, 8}, delays)
	delaysMu.Unlock()
}

func TestPiStartupRetryCancellationStopsWaitingAndAttempts(t *testing.T) {
	d := mustNewDispatcher("http://server.test")
	defer d.Close()
	sess := newPiOwnershipRetryTestSession(t, d)
	var attempts atomic.Int32
	d.runtimeOwnershipContinuationStarter = func(*interactiveSession) (bool, error) {
		attempts.Add(1)
		return false, errors.New("temporary startup transport failure")
	}
	d.runtimeOwnershipContinuationDelay = func(int) time.Duration { return time.Hour }
	d.scheduleRuntimeDeliveryOwnershipContinuationRetry(sess, errors.New("temporary startup transport failure"))
	require.Eventually(t, func() bool {
		sess.mu.Lock()
		defer sess.mu.Unlock()
		return sess.ownershipContinuationRetryCancel != nil
	}, time.Second, time.Millisecond)
	require.NoError(t, d.blockAndWaitRuntimeOwnershipRetries(context.Background(), sess))
	require.Zero(t, attempts.Load())
}
