package dispatcher

import (
	"context"
	"io"
	"os"
	"path/filepath"
	"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"
)

func initTokenForTest(t *testing.T) string {
	t.Helper()
	tokenDir := t.TempDir()
	t.Setenv("AGENTTEAM_TOKEN_DIR", tokenDir)
	return tokenDir
}

func boolForTest(value bool) *bool {
	return &value
}

func stubRuntimeBinaries(t *testing.T, names ...string) {
	t.Helper()
	runtimeDir := t.TempDir()
	for _, name := range names {
		path := filepath.Join(runtimeDir, name)
		require.NoError(t, os.WriteFile(path, []byte("#!/bin/sh\nexit 0\n"), 0o755))
		if name == "opencode" {
			require.NoError(t, os.WriteFile(filepath.Join(runtimeDir, "models.json"), []byte(`{"local":{"id":"local","name":"Local","env":[],"models":{}}}`+"\n"), 0o644))
			t.Setenv("AGENTTEAM_MANAGED_OPENCODE_PATH", path)
		}
		if name == "codex" {
			oldProbe := runtimeconfig.CodexAppServerAvailable
			runtimeconfig.CodexAppServerAvailable = func(string) bool { return true }
			t.Cleanup(func() { runtimeconfig.CodexAppServerAvailable = oldProbe })
		}
		if name == "kimi" {
			oldProbe := runtimeconfig.KimiACPAvailable
			runtimeconfig.KimiACPAvailable = func(string) bool { return true }
			t.Cleanup(func() { runtimeconfig.KimiACPAvailable = oldProbe })
		}
	}
	pathValue := runtimeDir
	if current := os.Getenv("PATH"); current != "" {
		pathValue += string(os.PathListSeparator) + current
	}
	t.Setenv("PATH", pathValue)
}

type stubDriver struct{}

func (*stubDriver) Name() string                    { return "codex" }
func (*stubDriver) SupportsStdinNotification() bool { return true }
func (*stubDriver) MCPToolPrefix() string           { return "mcp_agentteam_" }
func (*stubDriver) Spawn(context.Context, drivercontract.LaunchConfig) (drivercontract.SpawnResult, error) {
	return drivercontract.SpawnResult{}, nil
}
func (*stubDriver) ParseLine(string) []drivercontract.ParsedEvent { return nil }
func (*stubDriver) EncodeStdinMessage(text, _ string) []byte      { return []byte(text) }

type blockingDispatcherManagedDriver struct {
	ready      chan struct{}
	terminated chan struct{}
	done       chan struct{}
	r          *io.PipeReader
	w          *io.PipeWriter
}

func newBlockingDispatcherManagedDriver() *blockingDispatcherManagedDriver {
	r, w := io.Pipe()
	return &blockingDispatcherManagedDriver{
		ready:      make(chan struct{}),
		terminated: make(chan struct{}),
		done:       make(chan struct{}),
		r:          r,
		w:          w,
	}
}

func (d *blockingDispatcherManagedDriver) Name() string { return "builtin" }
func (d *blockingDispatcherManagedDriver) SupportsStdinNotification() bool {
	return true
}
func (d *blockingDispatcherManagedDriver) MCPToolPrefix() string { return "mcp_agentteam_" }
func (d *blockingDispatcherManagedDriver) Spawn(context.Context, drivercontract.LaunchConfig) (drivercontract.SpawnResult, error) {
	return drivercontract.SpawnResult{
		Stdin:  noopWriteCloser{},
		Stdout: d.r,
		Wait: func() error {
			close(d.ready)
			<-d.done
			return nil
		},
		Terminate: func() {
			select {
			case <-d.terminated:
			default:
				close(d.terminated)
			}
			_ = d.w.Close()
		},
	}, nil
}
func (d *blockingDispatcherManagedDriver) ParseLine(string) []drivercontract.ParsedEvent {
	return nil
}
func (d *blockingDispatcherManagedDriver) EncodeStdinMessage(string, string) []byte {
	return nil
}
func (d *blockingDispatcherManagedDriver) waitStartedWithin(timeout time.Duration) bool {
	select {
	case <-d.ready:
		return true
	case <-time.After(timeout):
		return false
	}
}
func (d *blockingDispatcherManagedDriver) terminatedWithin(timeout time.Duration) bool {
	select {
	case <-d.terminated:
		return true
	case <-time.After(timeout):
		return false
	}
}
func (d *blockingDispatcherManagedDriver) allowExit() {
	select {
	case <-d.done:
	default:
		close(d.done)
	}
}

type noopWriteCloser struct{}

func (noopWriteCloser) Write(p []byte) (int, error) { return len(p), nil }
func (noopWriteCloser) Close() error                { return nil }

func newShutdownDeadlineProcess(t *testing.T, agentID string) (*agent.Process, *blockingDispatcherManagedDriver) {
	t.Helper()
	driver := newBlockingDispatcherManagedDriver()
	process, err := agent.NewProcess(
		context.Background(),
		driver,
		drivercontract.LaunchConfig{
			AgentID: agentID,
			WorkDir: filepath.Join(t.TempDir(), agentID),
		},
		agent.ProcessOptions{},
	)
	require.NoError(t, err)
	require.True(t, driver.waitStartedWithin(time.Second))
	t.Cleanup(func() {
		driver.allowExit()
		process.Stop()
		require.True(t, process.WaitForExit(time.Second))
	})
	return process, driver
}

func newShutdownCallbackProcess(t *testing.T, agentID string, onExit func(*agent.Process)) (*agent.Process, *blockingDispatcherManagedDriver) {
	t.Helper()
	driver := newBlockingDispatcherManagedDriver()
	process, err := agent.NewProcess(
		context.Background(),
		driver,
		drivercontract.LaunchConfig{
			AgentID: agentID,
			WorkDir: filepath.Join(t.TempDir(), agentID),
		},
		agent.ProcessOptions{OnExit: onExit},
	)
	require.NoError(t, err)
	require.True(t, driver.waitStartedWithin(time.Second))
	return process, driver
}