package daemon

import (
	"context"
	"encoding/json"
	"errors"
	"sync"
	"testing"
	"time"

	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/AgentTeam/internal/runtimecontrol"
	"github.com/tidbcloud/AgentTeam/internal/store/agent"
)

type testDaemonTransportMetrics struct {
	mu             sync.Mutex
	requests       []testDaemonTransportRequestMetric
	routeDecisions []testDaemonTransportRouteMetric
	elasticEvents  []testDaemonTransportElasticMetric
	warmEvents     []testDaemonTransportWarmMetric
	pool           map[string]float64
}

type testDaemonTransportRequestMetric struct {
	endpoint string
	linkMode string
	outcome  string
	fallback string
}

type testDaemonTransportRouteMetric struct {
	endpoint   string
	linkPolicy string
	mode       string
	reason     string
}

type testDaemonTransportElasticMetric struct {
	event  string
	reason string
}

type testDaemonTransportWarmMetric struct {
	origin  string
	profile string
	event   string
	reason  string
}

func newTestDaemonTransportMetrics() *testDaemonTransportMetrics {
	return &testDaemonTransportMetrics{pool: make(map[string]float64)}
}

func (metrics *testDaemonTransportMetrics) RecordRemoteRequest(endpoint string, linkMode string, outcome string, fallback string, _ time.Duration) {
	metrics.mu.Lock()
	defer metrics.mu.Unlock()
	metrics.requests = append(metrics.requests, testDaemonTransportRequestMetric{endpoint: endpoint, linkMode: linkMode, outcome: outcome, fallback: fallback})
}

func (metrics *testDaemonTransportMetrics) RecordRouteDecision(endpoint string, linkPolicy string, mode string, reason string) {
	metrics.mu.Lock()
	defer metrics.mu.Unlock()
	metrics.routeDecisions = append(metrics.routeDecisions, testDaemonTransportRouteMetric{endpoint: endpoint, linkPolicy: linkPolicy, mode: mode, reason: reason})
}

func (metrics *testDaemonTransportMetrics) RecordElasticEvent(event string, reason string) {
	metrics.mu.Lock()
	defer metrics.mu.Unlock()
	metrics.elasticEvents = append(metrics.elasticEvents, testDaemonTransportElasticMetric{event: event, reason: reason})
}

func (metrics *testDaemonTransportMetrics) RecordWarmEvent(origin string, profile string, event string, reason string) {
	metrics.mu.Lock()
	defer metrics.mu.Unlock()
	metrics.warmEvents = append(metrics.warmEvents, testDaemonTransportWarmMetric{origin: origin, profile: profile, event: event, reason: reason})
}

func (metrics *testDaemonTransportMetrics) SetElasticPool(state string, value float64) {
	metrics.mu.Lock()
	defer metrics.mu.Unlock()
	metrics.pool[state] = value
}

func (metrics *testDaemonTransportMetrics) hasRequest(endpoint string, linkMode string, outcome string, fallback string) bool {
	metrics.mu.Lock()
	defer metrics.mu.Unlock()
	for _, request := range metrics.requests {
		if request.endpoint == endpoint && request.linkMode == linkMode && request.outcome == outcome && request.fallback == fallback {
			return true
		}
	}
	return false
}

func (metrics *testDaemonTransportMetrics) hasRoute(endpoint string, linkPolicy string, mode string, reason string) bool {
	metrics.mu.Lock()
	defer metrics.mu.Unlock()
	for _, route := range metrics.routeDecisions {
		if route.endpoint == endpoint && route.linkPolicy == linkPolicy && route.mode == mode && route.reason == reason {
			return true
		}
	}
	return false
}

func (metrics *testDaemonTransportMetrics) hasElasticEvent(event string, reason string) bool {
	metrics.mu.Lock()
	defer metrics.mu.Unlock()
	for _, elastic := range metrics.elasticEvents {
		if elastic.event == event && elastic.reason == reason {
			return true
		}
	}
	return false
}

func (metrics *testDaemonTransportMetrics) hasWarmEvent(origin string, profile string, event string, reason string) bool {
	metrics.mu.Lock()
	defer metrics.mu.Unlock()
	for _, warm := range metrics.warmEvents {
		if warm.origin == origin && warm.profile == profile && warm.event == event && warm.reason == reason {
			return true
		}
	}
	return false
}

func (metrics *testDaemonTransportMetrics) poolValue(state string) float64 {
	metrics.mu.Lock()
	defer metrics.mu.Unlock()
	return metrics.pool[state]
}

func TestRemoteEndpointCapabilityAndLinkPolicyMatrixCoversMigratedEndpoints(t *testing.T) {
	now := time.Date(2026, 7, 1, 3, 0, 0, 0, time.UTC)
	for _, descriptor := range defaultRemoteEndpointRegistry().Descriptors() {
		t.Run(descriptor.EndpointID, func(t *testing.T) {
			manager := newTestElasticRequestLinkManager(t, &now, ElasticRequestLinkManagerConfig{})
			client := newElasticReadyClient()

			lightDecision := manager.DecideRoute(client, descriptor, RemoteRequestRouteOptions{})
			if descriptor.LinkPolicy == RemoteEndpointLinkPolicySessionPreferred {
				require.Equal(t, RemoteRequestRouteSessionLink, lightDecision.Mode)
			} else {
				require.Equal(t, RemoteRequestRouteOpenElasticRequestLink, lightDecision.Mode)
			}

			manager = newTestElasticRequestLinkManager(t, &now, ElasticRequestLinkManagerConfig{})
			forcedDecision := manager.DecideRoute(newElasticReadyClient(), descriptor, RemoteRequestRouteOptions{ForceElastic: true})
			require.Equal(t, RemoteRequestRouteOpenElasticRequestLink, forcedDecision.Mode)

			missingCapability := newElasticReadyClient()
			missingCapability.capabilities = map[string]struct{}{}
			fallbackDecision := manager.DecideRoute(missingCapability, descriptor, RemoteRequestRouteOptions{ForceElastic: true})
			expectedFallbackMode := RemoteRequestRouteLegacySessionLinkFallback
			if descriptor.FallbackCommand == commandFallbackNone {
				expectedFallbackMode = RemoteRequestRouteStableError
			}
			require.Equal(t, expectedFallbackMode, fallbackDecision.Mode)
			require.Equal(t, "missing_transport_capability", fallbackDecision.Reason)

			registryMismatch := newElasticReadyClient()
			registryMismatch.endpointRegistryVersion = "daemon-endpoints:v0"
			registryDecision := manager.DecideRoute(registryMismatch, descriptor, RemoteRequestRouteOptions{ForceElastic: true})
			require.Equal(t, expectedFallbackMode, registryDecision.Mode)
			require.Equal(t, "endpoint_registry_mismatch", registryDecision.Reason)

			protocolMismatch := newElasticReadyClient()
			protocolMismatch.transportProtocol = defaultElasticRequestLinkProtocolVersion + 1
			protocolDecision := manager.DecideRoute(protocolMismatch, descriptor, RemoteRequestRouteOptions{ForceElastic: true})
			require.Equal(t, expectedFallbackMode, protocolDecision.Mode)
			require.Equal(t, "missing_transport_capability", protocolDecision.Reason)
		})
	}
}

func TestDaemonTransportMetricsRecordRemoteEndpointRequestPaths(t *testing.T) {
	now := time.Date(2026, 7, 1, 3, 10, 0, 0, time.UTC)
	metrics := newTestDaemonTransportMetrics()
	manager := newTestElasticRequestLinkManager(t, &now, ElasticRequestLinkManagerConfig{Metrics: metrics})
	hub := &Hub{remoteEndpoints: defaultRemoteEndpointRegistry(), elasticRequestLinks: manager}
	client := newElasticReadyClient()
	descriptor := defaultWorkspaceListEndpoint(t)
	open, _, err := manager.CreateLinkGrant(client, descriptor)
	require.NoError(t, err)
	sender := newCaptureElasticSender()
	_, err = manager.AcceptLink(client.workspaceID, open.LinkID, open.Token, sender)
	require.NoError(t, err)

	go func() {
		message := <-sender.messages
		assignment, ok := message.(RemoteRequestAssignment[json.RawMessage])
		require.True(t, ok)
		require.Equal(t, remoteEndpointWorkspaceList, assignment.EndpointID)
		manager.CompleteRemoteRequest(RemoteRequestResult[json.RawMessage]{
			Type:      sessionLinkRemoteRequestResultType,
			RequestID: assignment.RequestID,
			Success:   true,
			Outcome:   "ok",
			Payload:   workspaceListRemotePayload(t, "metrics.txt"),
		})
	}()

	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	ctx = withRemoteRequestRouteOptions(ctx, RemoteRequestRouteOptions{SessionLinkUnderPressure: true})
	page, err := hub.listWorkspace(ctx, client, "agt_1", "", "", agent.WorkspaceListOptions{Limit: 10})
	require.NoError(t, err)
	require.Equal(t, "metrics.txt", page.Files[0].Name)

	require.True(t, metrics.hasRoute("workspace.list", string(RemoteEndpointLinkPolicySessionPreferred), string(RemoteRequestRouteElasticRequestLink), "idle_elastic_request_link"))
	require.True(t, metrics.hasRequest("workspace.list", "elastic_request_link", "success", "none"))
	require.True(t, metrics.hasElasticEvent("grant_created", "ok"))
	require.True(t, metrics.hasElasticEvent("link_accepted", "ok"))
	require.True(t, metrics.hasElasticEvent("assignment_completed", "completed"))
	require.Equal(t, float64(0), metrics.poolValue("assignments"))
	require.Equal(t, float64(1), metrics.poolValue("links"))
}

func TestDaemonTransportMetricsRecordFallbackAndStableErrors(t *testing.T) {
	now := time.Date(2026, 7, 1, 3, 20, 0, 0, time.UTC)
	metrics := newTestDaemonTransportMetrics()
	manager := newTestElasticRequestLinkManager(t, &now, ElasticRequestLinkManagerConfig{Metrics: metrics})
	descriptor := defaultWorkspaceListEndpoint(t)
	oldDesktop := newElasticReadyClient()
	oldDesktop.capabilities = map[string]struct{}{}

	_, err := invokeRemoteEndpoint(RemoteEndpointInvocation[string]{
		Context:      context.Background(),
		Client:       oldDesktop,
		Manager:      manager,
		Descriptor:   descriptor,
		RouteOptions: RemoteRequestRouteOptions{ForceElastic: true},
		Fallback: func() (string, error) {
			return "fallback", nil
		},
		BuildSpec: func() (RemoteRequestAssignmentSpec, error) {
			return RemoteRequestAssignmentSpec{}, errors.New("spec builder must not run for route fallback")
		},
		MapResult: func(RemoteRequestResult[json.RawMessage]) (string, error) {
			return "", errors.New("result mapper must not run for route fallback")
		},
	})
	require.NoError(t, err)
	require.True(t, metrics.hasRoute("workspace.list", string(RemoteEndpointLinkPolicySessionPreferred), string(RemoteRequestRouteLegacySessionLinkFallback), "missing_transport_capability"))
	require.True(t, metrics.hasRequest("workspace.list", "legacy_session_link", "success", "route_fallback"))

	noFallback := descriptor
	noFallback.MetricsName = "workspace.list"
	noFallback.FallbackCommand = commandFallbackNone
	_, err = invokeRemoteEndpoint(RemoteEndpointInvocation[string]{
		Context:      context.Background(),
		Client:       oldDesktop,
		Manager:      manager,
		Descriptor:   noFallback,
		RouteOptions: RemoteRequestRouteOptions{ForceElastic: true},
		Fallback: func() (string, error) {
			return "", errors.New("fallback must not run for stable error")
		},
		BuildSpec: func() (RemoteRequestAssignmentSpec, error) {
			return RemoteRequestAssignmentSpec{}, errors.New("spec builder must not run for stable error")
		},
		MapResult: func(RemoteRequestResult[json.RawMessage]) (string, error) {
			return "", errors.New("result mapper must not run for stable error")
		},
	})
	require.ErrorContains(t, err, "desktop_unavailable")
	require.True(t, metrics.hasRoute("workspace.list", string(RemoteEndpointLinkPolicySessionPreferred), string(RemoteRequestRouteStableError), "missing_transport_capability"))
	require.True(t, metrics.hasRequest("workspace.list", "none", "unavailable", "none"))
}

func TestElasticRequestLinkDrainCleanupAndMetrics(t *testing.T) {
	now := time.Date(2026, 7, 1, 3, 30, 0, 0, time.UTC)
	metrics := newTestDaemonTransportMetrics()
	manager := newTestElasticRequestLinkManager(t, &now, ElasticRequestLinkManagerConfig{Metrics: metrics})
	client := newElasticReadyClient()
	descriptor := defaultWorkspaceListEndpoint(t)
	open, _, err := manager.CreateLinkGrant(client, descriptor)
	require.NoError(t, err)
	sender := &fakeElasticRequestLinkSender{}
	_, err = manager.AcceptLink(client.workspaceID, open.LinkID, open.Token, sender)
	require.NoError(t, err)

	require.True(t, manager.DrainLink(open.LinkID, "drain"))
	require.True(t, sender.closed)
	require.Equal(t, 0, manager.activeCount(client.machineID))
	require.True(t, metrics.hasElasticEvent("link_drained", "drain"))
	require.Equal(t, float64(0), metrics.poolValue("links"))
	require.Equal(t, float64(0), metrics.poolValue("grants"))
	require.Equal(t, float64(0), metrics.poolValue("assignments"))
}

func TestWarmPoolMetricsRecordBoundedLabels(t *testing.T) {
	now := time.Date(2026, 7, 1, 3, 40, 0, 0, time.UTC)
	metrics := newTestDaemonTransportMetrics()
	var opens []ElasticRequestLinkOpenCommand
	manager := newTestElasticRequestLinkManager(t, &now, ElasticRequestLinkManagerConfig{
		Metrics: metrics,
		WarmOpenSender: func(_ *clientState, open ElasticRequestLinkOpenCommand) error {
			opens = append(opens, open)
			return nil
		},
	})
	client := newElasticReadyClient()

	manager.OnRequestActivity(context.Background(), client, defaultWorkspaceListEndpoint(t), userInteractiveRemoteRequestContext())
	require.Len(t, opens, 1)
	require.True(t, metrics.hasWarmEvent("user_interactive", defaultWarmEndpointProfileName, "lease_refreshed", "web_api"))
	require.True(t, metrics.hasWarmEvent("user_interactive", defaultWarmEndpointProfileName, "grant_created", "ok"))

	_, err := manager.AcceptLink(client.workspaceID, opens[0].LinkID, opens[0].Token, &fakeElasticRequestLinkSender{})
	require.NoError(t, err)
	require.True(t, metrics.hasWarmEvent("user_interactive", defaultWarmEndpointProfileName, "connected", "ok"))

	manager.OnRequestActivity(context.Background(), client, defaultWorkspaceListEndpoint(t), runtimecontrol.RemoteRequestContext{})
	require.True(t, metrics.hasWarmEvent("background", defaultWarmEndpointProfileName, "skipped_background", "background"))
}