package memory

import (
	"archive/zip"
	"bytes"
	"context"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"sync/atomic"
	"testing"
	"time"

	"github.com/stretchr/testify/require"
	"github.com/tidbcloud/AgentTeam/desktop/internal/agent"
	"github.com/tidbcloud/AgentTeam/desktop/internal/dispatcher"
	"github.com/tidbcloud/AgentTeam/desktop/internal/nativememory"
	"github.com/tidbcloud/AgentTeam/desktop/internal/profile"
	"github.com/tidbcloud/AgentTeam/desktop/internal/provisioning/asset"
	"github.com/tidbcloud/AgentTeam/desktop/internal/types"
)

// accessBatchRequest is the test-local minimal view of the access-batch request
// body the provider posts; tests decode into it to assert the agent id.
type accessBatchRequest struct {
	AgentIDs []string `json:"agentIds"`
}

// fakeResolver is a test ServerConnectionResolver that returns a fixed connection info.
type fakeResolver struct {
	conn asset.ConnectionInfo
	ok   bool
}

func (f *fakeResolver) ResolveServerForAgent(ctx context.Context, agentID string) (asset.ConnectionInfo, bool) {
	return f.conn, f.ok
}

func newProvider(serverURL, userToken string) *Provider {
	return New(&fakeResolver{
		conn: asset.ConnectionInfo{ServerURL: serverURL, UserToken: userToken, WorkspaceID: "ws_1"},
		ok:   true,
	})
}

func TestLookupResolvesGrant(t *testing.T) {
	var capturedBody string
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		require.Equal(t, "/api/v1/workspaces/ws_1/memory/agents/access:batch", r.URL.Path)
		require.Equal(t, "Bearer tok", r.Header.Get("Authorization"))
		require.Equal(t, "ws_1", r.Header.Get("X-Workspace-Id"))
		require.Equal(t, r.Header.Get("X-Workspace-Id"), r.Header.Get("X-AgentTeam-Route-Workspace-Id"))

		var req accessBatchRequest
		require.NoError(t, json.NewDecoder(r.Body).Decode(&req))
		require.Equal(t, []string{"agt_001"}, req.AgentIDs)
		reqBytes, err := json.Marshal(req)
		require.NoError(t, err)
		capturedBody = string(reqBytes)

		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"agents":{"agt_001":{"id":"mem_abc","accessLevel":"write","syncProtocol":"legacy_git_v1","branch":"main","importExistingLocalContent":true,"repository":{"repoOwner":"ws_xxx","repoName":"mem_abc","repoFullName":"ws_xxx/mem_abc"}}},"syncToken":"ags_rat_sync","syncTokenExpires":"2099-01-02T03:04:05Z"}`))
	}))
	t.Cleanup(srv.Close)

	p := New(&fakeResolver{conn: asset.ConnectionInfo{ServerURL: srv.URL, UserToken: "tok", WorkspaceID: "ws_1"}, ok: true})
	access, state, err := p.Lookup(context.Background(), "agt_001")
	require.NoError(t, err)
	require.Equal(t, dispatcher.MemoryAccessBound, state)
	require.JSONEq(t, `{"agentIds":["agt_001"]}`, capturedBody)
	require.Equal(t, "mem_abc", access.MemoryID)
	require.Equal(t, dispatcher.MemorySyncProtocolLegacyGitV1, access.Protocol)
	require.Equal(t, srv.URL+"/api/v1/workspaces/ws_1/ags/ws_xxx/mem_abc.git", access.OriginURL)
	require.Equal(t, "main", access.Branch)
	require.True(t, access.ImportExistingLocalContent)
	require.Equal(t, "ags_rat_sync", access.Token)
	require.NotNil(t, access.Auth)
	req, err := http.NewRequest(http.MethodGet, access.OriginURL+"/info/refs", nil)
	require.NoError(t, err)
	access.Auth.SetAuth(req)
	require.Equal(t, "Bearer ags_rat_sync", req.Header.Get("Authorization"))
}

func TestLookupReturnsUnknownWhenBoundResponseHasNoSyncLease(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"agents":{"agt_001":{"id":"mem_abc","accessLevel":"write","syncProtocol":"legacy_git_v1","branch":"main","repository":{"repoOwner":"ws_xxx","repoName":"mem_abc","repoFullName":"ws_xxx/mem_abc"}}}}`))
	}))
	t.Cleanup(srv.Close)

	p := newProvider(srv.URL, "tok")
	_, state, err := p.Lookup(context.Background(), "agt_001")
	require.Error(t, err)
	require.Equal(t, dispatcher.MemoryAccessUnknown, state)
}

func TestLookupReturnsFalseWhenNoGrant(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		resp := batchResponse{
			Agents: map[string]*accessibleMemory{
				"agt_001": nil,
			},
		}
		w.Header().Set("Content-Type", "application/json")
		json.NewEncoder(w).Encode(resp)
	}))
	t.Cleanup(srv.Close)

	p := newProvider(srv.URL, "tok")
	_, state, err := p.Lookup(context.Background(), "agt_001")
	require.NoError(t, err)
	require.Equal(t, dispatcher.MemoryAccessUnbound, state)
}

func TestLookupReturnsFalseWhenAgentNotInResponse(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		resp := batchResponse{
			Agents: map[string]*accessibleMemory{
				"agt_other": {ID: "mem_x", SyncProtocol: string(dispatcher.MemorySyncProtocolLegacyGitV1), Branch: "main", Repository: &repositoryRef{RepoOwner: "ws", RepoName: "mem_x"}},
			},
		}
		w.Header().Set("Content-Type", "application/json")
		json.NewEncoder(w).Encode(resp)
	}))
	t.Cleanup(srv.Close)

	p := newProvider(srv.URL, "tok")
	_, state, err := p.Lookup(context.Background(), "agt_001")
	require.Error(t, err)
	require.Equal(t, dispatcher.MemoryAccessUnknown, state)
}

func TestLookupReturnsFalseWhenGrantHasNoRepository(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		resp := batchResponse{
			Agents: map[string]*accessibleMemory{
				"agt_001": {ID: "mem_abc", AccessLevel: "read", SyncProtocol: string(dispatcher.MemorySyncProtocolLegacyGitV1), Branch: "main"},
			},
		}
		w.Header().Set("Content-Type", "application/json")
		json.NewEncoder(w).Encode(resp)
	}))
	t.Cleanup(srv.Close)

	p := newProvider(srv.URL, "tok")
	_, state, err := p.Lookup(context.Background(), "agt_001")
	require.Error(t, err)
	require.Equal(t, dispatcher.MemoryAccessUnknown, state)
}

func TestLookupSelectsNativeProtocolWithoutParsingAGSCredential(t *testing.T) {
	var archive bytes.Buffer
	archiveWriter := zip.NewWriter(&archive)
	archiveFile, err := archiveWriter.Create("MEMORY.md")
	require.NoError(t, err)
	_, err = archiveFile.Write([]byte("# cloud\n"))
	require.NoError(t, err)
	require.NoError(t, archiveWriter.Close())
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		switch r.URL.Path {
		case "/api/v1/workspaces/ws_1/memory/agents/access:batch":
			_, _ = w.Write([]byte(`{"agents":{"agt_001":{"id":"mem_abc","accessLevel":"write","syncProtocol":"native_content_v1","contentId":"ctn_abc","contentVersionId":"ctv_head","contentGeneration":7,"acknowledgedSequence":23,"activeReplicaId":"replica-desktop"}},"syncToken":"malformed-token-that-must-be-ignored","syncTokenExpires":"invalid"}`))
		case "/api/v1/workspaces/ws_1/memory/mem_abc/checkpoints":
			require.Equal(t, "Bearer tok", r.Header.Get("Authorization"))
			var request nativeCheckpointRequest
			require.NoError(t, json.NewDecoder(r.Body).Decode(&request))
			require.Equal(t, "checkpoint-1", request.OperationID)
			require.Equal(t, base64.StdEncoding.EncodeToString([]byte("# local\n")), request.Files[0].ContentBase64)
			w.WriteHeader(http.StatusCreated)
			_, _ = w.Write([]byte(`{"versionId":"ctv_next","operationId":"checkpoint-1","replicaId":"replica-desktop","generation":7,"coveredFromSequence":24,"coveredThroughSequence":24}`))
		case "/api/v1/workspaces/ws_1/memory/mem_abc/versions/ctv_head/archive":
			_, _ = w.Write(archive.Bytes())
		default:
			t.Fatalf("unexpected path %s", r.URL.Path)
		}
	}))
	t.Cleanup(srv.Close)

	p := newProvider(srv.URL, "tok")
	access, state, err := p.Lookup(context.Background(), "agt_001")
	require.NoError(t, err)
	require.Equal(t, dispatcher.MemoryAccessBound, state)
	require.Equal(t, dispatcher.MemorySyncProtocolNativeContentV1, access.Protocol)
	require.Equal(t, "mem_abc", access.MemoryID)
	require.Empty(t, access.OriginURL)
	require.Empty(t, access.Branch)
	require.Nil(t, access.Auth)
	require.Empty(t, access.Token)
	require.Equal(t, "ctn_abc", access.ContentID)
	require.Equal(t, "ctv_head", access.ContentVersionID)
	require.EqualValues(t, 7, access.ContentGeneration)
	require.EqualValues(t, 23, access.AcknowledgedSequence)
	require.Equal(t, "replica-desktop", access.ActiveReplicaID)
	require.NotNil(t, access.NativeRemote)
	result, err := access.NativeRemote.PublishCheckpoint(context.Background(), access.MemoryID, nativememory.Checkpoint{
		OperationID: "checkpoint-1", ReplicaID: "replica-desktop", ExpectedVersionID: "ctv_head", Generation: 7,
		CoveredFromSequence: 24, CoveredThroughSequence: 24,
		Files: []nativememory.File{{Path: "MEMORY.md", Content: []byte("# local\n")}},
	})
	require.NoError(t, err)
	require.Equal(t, "ctv_next", result.VersionID)
	reader, err := access.NativeRemote.OpenVersionArchive(context.Background(), access.MemoryID, access.ContentVersionID)
	require.NoError(t, err)
	archiveBytes, err := io.ReadAll(reader)
	require.NoError(t, err)
	require.NoError(t, reader.Close())
	require.Equal(t, archive.Bytes(), archiveBytes)
}

func TestLookupFailsClosedForUnknownOrIncompleteProtocol(t *testing.T) {
	tests := []string{
		`{"agents":{"agt_001":{"id":"mem_abc","syncProtocol":"future_v2"}}}`,
		`{"agents":{"agt_001":{"id":"mem_abc","syncProtocol":"native_content_v1","contentId":"ctn_abc","contentGeneration":7}}}`,
		`{"agents":{"agt_001":{"id":"mem_abc","syncProtocol":"native_content_v1","contentId":"ctn_abc","contentVersionId":"ctv_head","contentGeneration":7,"branch":"main"}}}`,
		`{"agents":{"agt_001":{"id":"mem_abc","branch":"main","repository":{"repoFullName":"ws/mem"}}},"syncToken":"token","syncTokenExpires":"2099-01-02T03:04:05Z"}`,
	}
	for _, body := range tests {
		t.Run(body, func(t *testing.T) {
			srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				w.Header().Set("Content-Type", "application/json")
				_, _ = w.Write([]byte(body))
			}))
			t.Cleanup(srv.Close)

			_, state, err := newProvider(srv.URL, "tok").Lookup(context.Background(), "agt_001")
			require.Error(t, err)
			require.Equal(t, dispatcher.MemoryAccessUnknown, state)
		})
	}
}

func TestLookupErrorOnNon200(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusUnauthorized)
	}))
	t.Cleanup(srv.Close)

	p := newProvider(srv.URL, "tok")
	_, state, err := p.Lookup(context.Background(), "agt_001")
	require.Error(t, err)
	require.Equal(t, dispatcher.MemoryAccessUnknown, state)
}

func TestLookupReturnsFalseWhenResolverHasNoChannel(t *testing.T) {
	p := New(&fakeResolver{ok: false})
	_, state, err := p.Lookup(context.Background(), "agt_001")
	require.Error(t, err)
	require.Equal(t, dispatcher.MemoryAccessUnknown, state)
}

func TestLookupErrorOnUnreachableServer(t *testing.T) {
	p := newProvider("http://127.0.0.1:1", "tok")
	_, state, err := p.Lookup(context.Background(), "agt_001")
	require.Error(t, err)
	require.Equal(t, dispatcher.MemoryAccessUnknown, state)
}

func TestComposeOriginURL(t *testing.T) {
	url := composeOriginURL("https://agentteam.example.com", "wsp_1", &repositoryRef{RepoOwner: "ws_x", RepoName: "mem_y", RepoFullName: "ws_x/mem_y"})
	require.Equal(t, "https://agentteam.example.com/api/v1/workspaces/wsp_1/ags/ws_x/mem_y.git", url)
}

func TestComposeOriginURLWithServerPathPrefix(t *testing.T) {
	url := composeOriginURL("https://app.example.com/agentteam/", "wsp_1", &repositoryRef{RepoOwner: "ws_x", RepoName: "mem_y", RepoFullName: "ws_x/mem_y"})
	require.Equal(t, "https://app.example.com/agentteam/api/v1/workspaces/wsp_1/ags/ws_x/mem_y.git", url)
}

func TestComposeOriginURLEmptyComponents(t *testing.T) {
	require.Equal(t, "", composeOriginURL("https://example.com", "wsp_1", nil))
	require.Equal(t, "", composeOriginURL("https://example.com", "wsp_1", &repositoryRef{}))
	require.Equal(t, "", composeOriginURL("https://example.com", "wsp_1", &repositoryRef{RepoOwner: "ws_x", RepoName: "mem_y"}))
}

func TestLookupSatisfiesMemoryAccessProviderInterface(t *testing.T) {
	var _ dispatcher.MemoryAccessProvider = (*Provider)(nil)
	var _ dispatcher.MemoryAccessBatchProvider = (*Provider)(nil)
}

func TestLookupBatchUsesOneRequestPerTwoHundredAgents(t *testing.T) {
	for _, tc := range []struct {
		count int
		calls int32
	}{{1, 1}, {199, 1}, {201, 2}} {
		t.Run(fmt.Sprintf("agents_%d", tc.count), func(t *testing.T) {
			var calls atomic.Int32
			srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				calls.Add(1)
				var request accessBatchRequest
				require.NoError(t, json.NewDecoder(r.Body).Decode(&request))
				require.LessOrEqual(t, len(request.AgentIDs), 200)
				response := batchResponse{Agents: make(map[string]*accessibleMemory, len(request.AgentIDs))}
				for _, id := range request.AgentIDs {
					response.Agents[id] = nil
				}
				require.NoError(t, json.NewEncoder(w).Encode(response))
			}))
			t.Cleanup(srv.Close)
			ids := make([]string, tc.count)
			for index := range ids {
				ids[index] = fmt.Sprintf("agt_%03d", index)
			}
			result := collectBatchLookups(newProvider(srv.URL, "tok").LookupBatch(t.Context(), ids))
			require.Len(t, result, tc.count)
			for _, id := range ids {
				require.Equal(t, dispatcher.MemoryAccessUnbound, result[id].State)
				require.NoError(t, result[id].Err)
			}
			require.Equal(t, tc.calls, calls.Load())
		})
	}
}

type connectionByAgentResolver map[string]asset.ConnectionInfo

func (r connectionByAgentResolver) ResolveServerForAgent(_ context.Context, agentID string) (asset.ConnectionInfo, bool) {
	conn, ok := r[agentID]
	return conn, ok
}

func TestLookupBatchIsolatesFailedConnectionGroup(t *testing.T) {
	good := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		_ = json.NewEncoder(w).Encode(batchResponse{Agents: map[string]*accessibleMemory{"agt_good": nil}})
	}))
	bad := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(http.StatusServiceUnavailable) }))
	t.Cleanup(good.Close)
	t.Cleanup(bad.Close)
	provider := New(connectionByAgentResolver{
		"agt_good": {ServerURL: good.URL, UserToken: "tok", WorkspaceID: "ws_1"},
		"agt_bad":  {ServerURL: bad.URL, UserToken: "tok", WorkspaceID: "ws_2"},
	})
	result := collectBatchLookups(provider.LookupBatch(t.Context(), []string{"agt_bad", "agt_good"}))
	require.Equal(t, dispatcher.MemoryAccessUnbound, result["agt_good"].State)
	require.NoError(t, result["agt_good"].Err)
	require.Equal(t, dispatcher.MemoryAccessUnknown, result["agt_bad"].State)
	require.Error(t, result["agt_bad"].Err)
}

func TestLookupBatchYieldsHealthyGroupBeforeSlowGroup(t *testing.T) {
	releaseSlow := make(chan struct{}, 1)
	slowStarted := make(chan struct{})
	slow := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		close(slowStarted)
		<-releaseSlow
		_ = json.NewEncoder(w).Encode(batchResponse{Agents: map[string]*accessibleMemory{"agt_slow": nil}})
	}))
	fast := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		_ = json.NewEncoder(w).Encode(batchResponse{Agents: map[string]*accessibleMemory{"agt_fast": nil}})
	}))
	t.Cleanup(slow.Close)
	t.Cleanup(fast.Close)
	t.Cleanup(func() {
		select {
		case releaseSlow <- struct{}{}:
		default:
		}
	})
	provider := New(connectionByAgentResolver{
		"agt_slow": {ServerURL: slow.URL, UserToken: "tok", WorkspaceID: "ws_1"},
		"agt_fast": {ServerURL: fast.URL, UserToken: "tok", WorkspaceID: "ws_2"},
	})
	results := provider.LookupBatch(t.Context(), []string{"agt_slow", "agt_fast"})
	<-slowStarted
	first := <-results
	require.Contains(t, first, "agt_fast", "healthy connection result must not wait for the slow group")
	require.Equal(t, dispatcher.MemoryAccessUnbound, first["agt_fast"].State)
	releaseSlow <- struct{}{}
	rest := collectBatchLookups(results)
	require.Equal(t, dispatcher.MemoryAccessUnbound, rest["agt_slow"].State)
}

func collectBatchLookups(ch <-chan map[string]dispatcher.MemoryAccessLookup) map[string]dispatcher.MemoryAccessLookup {
	out := make(map[string]dispatcher.MemoryAccessLookup)
	for snapshot := range ch {
		for id, lookup := range snapshot {
			out[id] = lookup
		}
	}
	return out
}

// fakeChannel is a minimal types.ServerConnection + ConnectionInfoProvider for resolver tests.
type fakeChannel struct {
	serverURL        string
	authToken        string
	tokenStarted     chan<- struct{}
	releaseUserToken <-chan struct{}
}

func (f *fakeChannel) ID() string                      { return "chn_test" }
func (f *fakeChannel) Enabled() bool                   { return true }
func (f *fakeChannel) Start(ctx context.Context) error { return nil }
func (f *fakeChannel) Stop(ctx context.Context) error  { return nil }
func (f *fakeChannel) Send(ctx context.Context, meta types.Metadata, msg types.OutboundMessage) error {
	return nil
}
func (f *fakeChannel) ServerURL() string   { return f.serverURL }
func (f *fakeChannel) WorkspaceID() string { return "ws_test" }
func (f *fakeChannel) UserToken(ctx context.Context) (string, error) {
	if f.tokenStarted != nil {
		select {
		case f.tokenStarted <- struct{}{}:
		default:
		}
	}
	if f.releaseUserToken != nil {
		select {
		case <-f.releaseUserToken:
		case <-ctx.Done():
			return "", ctx.Err()
		}
	}
	return f.authToken, nil
}
func (f *fakeChannel) RefreshUserToken(ctx context.Context, stale string) (string, error) {
	return f.authToken, nil
}

type fakeServerConnectionAccessor map[string]types.ServerConnection

func (f fakeServerConnectionAccessor) ServerConnection(id string) (types.ServerConnection, bool) {
	conn, ok := f[id]
	return conn, ok
}

func TestLookupBatchYieldsHealthyServerBeforeSlowUserTokenResolver(t *testing.T) {
	t.Setenv(profile.ConfigDirEnvKey, t.TempDir())
	slowRelease := make(chan struct{})
	slowStarted := make(chan struct{}, 1)
	slowServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		_ = json.NewEncoder(w).Encode(batchResponse{Agents: map[string]*accessibleMemory{"agt_slow_resolver": nil}})
	}))
	fastServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		_ = json.NewEncoder(w).Encode(batchResponse{Agents: map[string]*accessibleMemory{"agt_fast_resolver": nil}})
	}))
	t.Cleanup(slowServer.Close)
	t.Cleanup(fastServer.Close)
	manager := profile.Default()
	for _, cfg := range []agent.Config{
		{ID: "agt_slow_resolver", Name: "Slow resolver", Runtime: "codex", Source: agent.AgentSource{Type: "agentteam", ID: "chn_slow"}},
		{ID: "agt_fast_resolver", Name: "Fast resolver", Runtime: "codex", Source: agent.AgentSource{Type: "agentteam", ID: "chn_fast"}},
	} {
		require.NoError(t, manager.Agents().Add(cfg))
		id := cfg.ID
		t.Cleanup(func() { _ = manager.Agents().Delete(id) })
	}
	resolver := NewAgentServerConnectionResolver(fakeServerConnectionAccessor{
		"chn_slow": &fakeChannel{serverURL: slowServer.URL, authToken: "tok", tokenStarted: slowStarted, releaseUserToken: slowRelease},
		"chn_fast": &fakeChannel{serverURL: fastServer.URL, authToken: "tok"},
	})
	results := New(resolver).LookupBatch(t.Context(), []string{"agt_slow_resolver", "agt_fast_resolver"})
	select {
	case <-slowStarted:
	case <-time.After(time.Second):
		t.Fatal("slow server token lookup did not start")
	}
	select {
	case first := <-results:
		require.Contains(t, first, "agt_fast_resolver", "healthy server must resolve and fetch while another UserToken call is blocked")
		require.Equal(t, dispatcher.MemoryAccessUnbound, first["agt_fast_resolver"].State)
	case <-time.After(time.Second):
		t.Fatal("healthy server was starved by slow UserToken resolution")
	}
	close(slowRelease)
	rest := collectBatchLookups(results)
	require.Equal(t, dispatcher.MemoryAccessUnbound, rest["agt_slow_resolver"].State)
}

func TestAgentServerRuntimeResolverTypeAssertion(t *testing.T) {
	var _ ConnectionInfoProvider = (*fakeChannel)(nil)
}