package agent

import (
	"context"
	"crypto/ed25519"
	"encoding/base64"
	"encoding/json"
	"errors"
	"strings"
	"testing"
	"time"

	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/AgentTeam/internal/maas/capability"
	"github.com/tidbcloud/AgentTeam/internal/service/maas"
	storeagent "github.com/tidbcloud/AgentTeam/internal/store/agent"
	"github.com/tidbcloud/AgentTeam/internal/store/storeerr"
)

type fakeMaaSLLMProxyScopeStore struct {
	scope storeagent.MaaSLLMProxyScope
	err   error
}

func (fake fakeMaaSLLMProxyScopeStore) LoadMaaSLLMProxyScope(context.Context, storeagent.AuthenticatedAgent) (storeagent.MaaSLLMProxyScope, error) {
	return fake.scope, fake.err
}

type fakeRuntimeSessionLookup struct {
	sessionID           string
	ok                  bool
	running             bool
	publishAfterLookups int
	sessionLookups      int
}

func (fake *fakeRuntimeSessionLookup) AgentRuntimeSessionID(string, string) (string, bool) {
	fake.sessionLookups++
	if !fake.ok && fake.publishAfterLookups > 0 && fake.sessionLookups >= fake.publishAfterLookups {
		fake.ok = true
	}
	return fake.sessionID, fake.ok
}

func (fake *fakeRuntimeSessionLookup) AgentRuntimeRunning(string, string) bool {
	return fake.running
}

type fakeMaaSProjectionPusher struct {
	projection capability.AuthorityProjection
	err        error
	count      int
}

func (fake *fakeMaaSProjectionPusher) UpsertAuthorityProjection(_ context.Context, projection capability.AuthorityProjection) error {
	fake.count++
	fake.projection = projection
	return fake.err
}

type fakeHostedAgentValidator struct {
	last maas.AgentBindingValidation
	err  error
}

func (fake *fakeHostedAgentValidator) ValidateAgentBinding(_ context.Context, params maas.AgentBindingValidation) error {
	fake.last = params
	return fake.err
}

func TestMaaSLLMProxySessionServiceIssuesCapabilityAndProjection(t *testing.T) {
	publicKey, privateKey, err := ed25519.GenerateKey(nil)
	require.NoError(t, err)

	now := time.Date(2026, 6, 1, 10, 0, 0, 0, time.UTC)
	projections := &fakeMaaSProjectionPusher{}
	hosted := &fakeHostedAgentValidator{}
	service := NewMaaSLLMProxySessionService(MaaSLLMProxySessionConfig{
		Scopes: fakeMaaSLLMProxyScopeStore{scope: storeagent.MaaSLLMProxyScope{
			OrganizationID:      "org_test",
			WorkspaceID:         "wsp_test",
			AgentID:             "agt_0000000000test",
			MemberID:            "00000000test",
			MachineID:           "mac_test",
			Runtime:             "claude",
			ModelReference:      "agentteam/deepseek-v4-pro",
			LLMMode:             storeagent.AgentLLMModeAgentTeamHosted,
			AuthorityGeneration: 1,
			OrganizationStatus:  capability.StatusActive,
			WorkspaceStatus:     capability.StatusActive,
			AgentStatus:         capability.StatusActive,
			RuntimeStatus:       capability.StatusActive,
		}},
		HostedAgents: hosted,
		Signer: capability.Signer{
			KeyID:      "dev",
			PrivateKey: privateKey,
		},
		RuntimeSessions: &fakeRuntimeSessionLookup{sessionID: "ars_test", ok: true},
		ProjectionPush:  projections,
		RuntimeBaseURL:  "http://127.0.0.1:8080/maas/v1",
		Clock:           func() time.Time { return now },
	})

	session, err := service.IssueMaaSLLMProxySession(context.Background(), storeagent.AuthenticatedAgent{
		AgentID: "agt_0000000000test",
	}, "lps_prev")
	require.NoError(t, err)

	require.NotEmpty(t, session.SessionID)
	require.NotEmpty(t, session.ProxyToken)
	require.Equal(t, now.Add(15*time.Minute), session.ExpiresAt)
	require.Equal(t, now.Add(10*time.Minute), *session.RefreshAt)
	require.Equal(t, "lps_prev", session.SupersedesSessionID)
	require.Equal(t, "http://127.0.0.1:8080/maas/v1", session.MaasBaseURL)
	require.Equal(t, LLMProxyCredentialKindMaaSCapability, session.CredentialKind)

	require.Equal(t, 1, projections.count)
	require.Equal(t, "org_test", projections.projection.OrganizationID)
	require.Equal(t, "mac_test", projections.projection.MachineID)
	require.Equal(t, "ars_test", projections.projection.RuntimeSessionID)
	require.Equal(t, capability.ProjectionID(projections.projection), projections.projection.ID)
	require.Equal(t, "agentteam/deepseek-v4-pro", hosted.last.ModelReference)
	require.Equal(t, storeagent.AgentLLMModeAgentTeamHosted, hosted.last.LLMMode)

	parts := strings.Split(session.ProxyToken, ".")
	require.Len(t, parts, 3)
	payloadRaw, err := base64.RawURLEncoding.DecodeString(parts[1])
	require.NoError(t, err)
	var payload map[string]any
	require.NoError(t, json.Unmarshal(payloadRaw, &payload))
	require.Equal(t, capability.DefaultIssuer, payload["iss"])
	require.Equal(t, capability.DefaultAudience, payload["aud"])
	require.Equal(t, session.SessionID, payload["sub"])
	require.Equal(t, session.SessionID, payload["capability_session_id"])
	require.Equal(t, "org_test", payload["organization_id"])
	require.Equal(t, "wsp_test", payload["workspace_id"])
	require.Equal(t, "agt_0000000000test", payload["agent_id"])
	require.Equal(t, "00000000test", payload["member_id"])
	require.Equal(t, "ars_test", payload["runtime_session_id"])
	require.True(t, ed25519.Verify(publicKey, []byte(parts[0]+"."+parts[1]), mustDecodeSignature(t, parts[2])))
}

func TestMaaSLLMProxySessionServiceRefreshRequiresSelector(t *testing.T) {
	projections := &fakeMaaSProjectionPusher{}
	service := NewMaaSLLMProxySessionService(MaaSLLMProxySessionConfig{
		Scopes:         fakeMaaSLLMProxyScopeStore{},
		ProjectionPush: projections,
		RuntimeBaseURL: "http://127.0.0.1:8080/maas/v1",
	})

	_, err := service.RefreshMaaSLLMProxySession(context.Background(), storeagent.AuthenticatedAgent{}, "")
	require.ErrorIs(t, err, storeerr.ErrConflict)
	require.Zero(t, projections.count)
}

func TestMaaSLLMProxySessionServiceRejectsNonHostedScope(t *testing.T) {
	projections := &fakeMaaSProjectionPusher{}
	service := NewMaaSLLMProxySessionService(MaaSLLMProxySessionConfig{
		Scopes: fakeMaaSLLMProxyScopeStore{scope: storeagent.MaaSLLMProxyScope{
			OrganizationID:      "org_test",
			WorkspaceID:         "wsp_test",
			AgentID:             "agt_0000000000test",
			MemberID:            "00000000test",
			MachineID:           "mac_test",
			Runtime:             "claude",
			ModelReference:      "claude-3",
			LLMMode:             storeagent.AgentLLMModeAgentTeamProxy,
			AuthorityGeneration: 1,
			OrganizationStatus:  capability.StatusActive,
			WorkspaceStatus:     capability.StatusActive,
			AgentStatus:         capability.StatusActive,
			RuntimeStatus:       capability.StatusActive,
		}},
		HostedAgents:    &fakeHostedAgentValidator{},
		RuntimeSessions: &fakeRuntimeSessionLookup{sessionID: "ars_test", ok: true},
		ProjectionPush:  projections,
		RuntimeBaseURL:  "http://127.0.0.1:8080/maas/v1",
	})

	_, err := service.IssueMaaSLLMProxySession(context.Background(), storeagent.AuthenticatedAgent{}, "")
	require.ErrorIs(t, err, storeerr.ErrConflict)
	require.Contains(t, err.Error(), maas.ReasonRuntimeUnavailable)
	require.Zero(t, projections.count)
}

func TestMaaSLLMProxySessionServiceWrapsProjectionPushFailure(t *testing.T) {
	_, privateKey, err := ed25519.GenerateKey(nil)
	require.NoError(t, err)

	service := NewMaaSLLMProxySessionService(MaaSLLMProxySessionConfig{
		Scopes: fakeMaaSLLMProxyScopeStore{scope: storeagent.MaaSLLMProxyScope{
			OrganizationID:      "org_test",
			WorkspaceID:         "wsp_test",
			AgentID:             "agt_0000000000test",
			MemberID:            "00000000test",
			MachineID:           "mac_test",
			Runtime:             "claude",
			ModelReference:      "agentteam/deepseek-v4-pro",
			LLMMode:             storeagent.AgentLLMModeAgentTeamHosted,
			AuthorityGeneration: 1,
			OrganizationStatus:  capability.StatusActive,
			WorkspaceStatus:     capability.StatusActive,
			AgentStatus:         capability.StatusActive,
			RuntimeStatus:       capability.StatusActive,
		}},
		HostedAgents: &fakeHostedAgentValidator{},
		Signer: capability.Signer{
			KeyID:      "dev",
			PrivateKey: privateKey,
		},
		RuntimeSessions: &fakeRuntimeSessionLookup{sessionID: "ars_test", ok: true},
		ProjectionPush:  &fakeMaaSProjectionPusher{err: errors.New("agentteam-maas unavailable")},
		RuntimeBaseURL:  "http://127.0.0.1:8080/maas/v1",
	})

	_, err = service.IssueMaaSLLMProxySession(context.Background(), storeagent.AuthenticatedAgent{}, "")
	require.Error(t, err)
	require.Contains(t, err.Error(), "projection_push")
	require.Contains(t, err.Error(), "agentteam-maas unavailable")
}

func TestMaaSLLMProxySessionServiceRejectsMissingRuntimeSession(t *testing.T) {
	projections := &fakeMaaSProjectionPusher{}
	service := NewMaaSLLMProxySessionService(MaaSLLMProxySessionConfig{
		Scopes: fakeMaaSLLMProxyScopeStore{scope: storeagent.MaaSLLMProxyScope{
			OrganizationID:      "org_test",
			WorkspaceID:         "wsp_test",
			AgentID:             "agt_0000000000test",
			MemberID:            "00000000test",
			MachineID:           "mac_test",
			Runtime:             "claude",
			ModelReference:      "agentteam/deepseek-v4-pro",
			LLMMode:             storeagent.AgentLLMModeAgentTeamHosted,
			AuthorityGeneration: 1,
			OrganizationStatus:  capability.StatusActive,
			WorkspaceStatus:     capability.StatusActive,
			AgentStatus:         capability.StatusActive,
			RuntimeStatus:       capability.StatusActive,
		}},
		RuntimeSessions: &fakeRuntimeSessionLookup{running: true},
		ProjectionPush:  projections,
		RuntimeBaseURL:  "http://127.0.0.1:8080/maas/v1",
		ReadyWait:       10 * time.Millisecond,
		ReadyPoll:       time.Millisecond,
	})

	_, err := service.IssueMaaSLLMProxySession(context.Background(), storeagent.AuthenticatedAgent{}, "")
	require.ErrorIs(t, err, storeerr.ErrConflict)
	require.Contains(t, err.Error(), "agent runtime session is not ready")
	require.Zero(t, projections.count)
}

func TestMaaSLLMProxySessionServiceWaitsForRunningRuntimeSession(t *testing.T) {
	_, privateKey, err := ed25519.GenerateKey(nil)
	require.NoError(t, err)

	runtimeSessions := &fakeRuntimeSessionLookup{
		sessionID:           "ars_late",
		running:             true,
		publishAfterLookups: 2,
	}

	projections := &fakeMaaSProjectionPusher{}
	service := NewMaaSLLMProxySessionService(MaaSLLMProxySessionConfig{
		Scopes: fakeMaaSLLMProxyScopeStore{scope: storeagent.MaaSLLMProxyScope{
			OrganizationID:      "org_test",
			WorkspaceID:         "wsp_test",
			AgentID:             "agt_0000000000test",
			MemberID:            "00000000test",
			MachineID:           "mac_test",
			Runtime:             "claude",
			ModelReference:      "agentteam/deepseek-v4-pro",
			LLMMode:             storeagent.AgentLLMModeAgentTeamHosted,
			AuthorityGeneration: 1,
			OrganizationStatus:  capability.StatusActive,
			WorkspaceStatus:     capability.StatusActive,
			AgentStatus:         capability.StatusActive,
			RuntimeStatus:       capability.StatusActive,
		}},
		HostedAgents:    &fakeHostedAgentValidator{},
		Signer:          capability.Signer{KeyID: "dev", PrivateKey: privateKey},
		RuntimeSessions: runtimeSessions,
		ProjectionPush:  projections,
		RuntimeBaseURL:  "http://127.0.0.1:8080/maas/v1",
		ReadyWait:       100 * time.Millisecond,
		ReadyPoll:       time.Millisecond,
	})

	session, err := service.IssueMaaSLLMProxySession(context.Background(), storeagent.AuthenticatedAgent{}, "")
	require.NoError(t, err)
	require.NotEmpty(t, session.ProxyToken)
	require.Equal(t, "ars_late", projections.projection.RuntimeSessionID)
}

func mustDecodeSignature(t *testing.T, raw string) []byte {
	t.Helper()
	signature, err := base64.RawURLEncoding.DecodeString(raw)
	require.NoError(t, err)
	return signature
}