package runtime

import (
	"github.com/tidbcloud/AgentTeam/internal/knowledgeqa/protocol"

	"context"
	"encoding/json"
	"errors"
	"strings"
	"testing"
	"time"

	"github.com/stretchr/testify/require"
)

type memoryHandleStore struct {
	handles   map[string]protocol.RootHandle
	loadCalls int
	loadErrAt int
}

func (store *memoryHandleStore) PutRootHandles(_ context.Context, handles []protocol.RootHandle) error {
	if store.handles == nil {
		store.handles = make(map[string]protocol.RootHandle)
	}
	for _, handle := range handles {
		store.handles[handle.ID] = handle
	}
	return nil
}

func (store *memoryHandleStore) LoadRootHandles(_ context.Context, principal protocol.Principal, handleIDs []string) ([]protocol.RootHandle, error) {
	store.loadCalls++
	if store.loadErrAt == store.loadCalls {
		return nil, errors.New("load failed")
	}
	handles := make([]protocol.RootHandle, 0, len(handleIDs))
	for _, handleID := range normalizeUniqueStrings(handleIDs) {
		handle, ok := store.handles[handleID]
		if !ok || handle.OrganizationID != principal.OrganizationID ||
			handle.WorkspaceID != principal.WorkspaceID ||
			handle.RequesterID != principal.RequesterID || handle.TurnID != principal.TurnID {
			continue
		}
		handles = append(handles, handle)
	}
	return handles, nil
}

func (store *memoryHandleStore) DeleteRootHandles(_ context.Context, principal protocol.Principal, handleIDs []string) error {
	for _, id := range normalizeUniqueStrings(handleIDs) {
		handle, ok := store.handles[id]
		if ok && handle.OrganizationID == principal.OrganizationID &&
			handle.WorkspaceID == principal.WorkspaceID &&
			handle.RequesterID == principal.RequesterID && handle.TurnID == principal.TurnID {
			delete(store.handles, id)
		}
	}
	return nil
}

type fakeRepository struct {
	grepSnapshot RepositorySnapshot
	readSnapshot RepositorySnapshot
	readRequest  RepositoryReadRequest
	grepErr      error
	readErr      error
	grepCalls    int
	readCalls    int
	readByLine   bool
	refName      string
	refKind      string
	refSHA       string
	refErr       error
}

func (repository *fakeRepository) ResolveRef(_ context.Context, _ RepositorySnapshot, refName, refKind string) (string, error) {
	repository.refName = refName
	repository.refKind = refKind
	if repository.refErr != nil {
		return "", repository.refErr
	}
	if repository.refSHA != "" {
		return repository.refSHA, nil
	}
	return strings.Repeat("a", 40), nil
}

func (repository *fakeRepository) Grep(_ context.Context, snapshot RepositorySnapshot, _ RepositoryGrepRequest, emit func(RepositoryGrepMatch) error) error {
	repository.grepCalls++
	repository.grepSnapshot = snapshot
	if repository.grepErr != nil {
		return repository.grepErr
	}
	return emit(RepositoryGrepMatch{Path: "docs/design.md", Line: 12, Preview: "server hosted"})
}

func (repository *fakeRepository) Read(_ context.Context, snapshot RepositorySnapshot, request RepositoryReadRequest, emit func(RepositoryReadChunk) error) error {
	repository.readCalls++
	repository.readSnapshot = snapshot
	repository.readRequest = request
	if repository.readErr != nil {
		return repository.readErr
	}
	if repository.readByLine {
		for line := request.LineStart; line <= request.LineEnd; line++ {
			if err := emit(RepositoryReadChunk{
				Content: []byte("server hosted\n"), LineStart: line, LineEnd: line,
			}); err != nil {
				return err
			}
		}
		return nil
	}
	return emit(RepositoryReadChunk{Content: []byte("server hosted\n"), LineStart: request.LineStart, LineEnd: request.LineEnd})
}

func (repository *fakeRepository) List(_ context.Context, snapshot RepositorySnapshot, request RepositoryListRequest, emit func(RepositoryListEntry) error) error {
	repository.readSnapshot = snapshot
	entryPath := "docs"
	entryKind := "directory"
	if request.Path != "" {
		entryPath = request.Path + "/design.md"
		entryKind = "file"
	}
	return emit(RepositoryListEntry{Path: entryPath, Kind: entryKind})
}

func TestExecutionGrepUsesFrozenRevisionWithoutCitationMetadata(t *testing.T) {
	now := time.Date(2026, 7, 25, 10, 0, 0, 0, time.UTC)
	handles := &memoryHandleStore{handles: map[string]protocol.RootHandle{
		"kqr_1": testRootHandle(now),
	}}
	repository := &fakeRepository{}
	service := &Service{Handles: handles, Repository: repository, Clock: func() time.Time { return now }}
	execution := authorizeTestExecution(t, service, "kqr_1")
	var matches []GrepMatch
	completion, err := execution.Grep(context.Background(), GrepRequest{
		HandleIDs: []string{"kqr_1"}, Pattern: "server.*hosted", PathGlobs: []string{"docs/*.md"},
	}, func(match GrepMatch) error {
		matches = append(matches, match)
		return nil
	})
	require.NoError(t, err)
	require.True(t, completion.Complete)
	require.Equal(t, RepositorySnapshot{
		RepositoryRef: "repo-internal-1", Revision: "sha-fixed",
		WorkspaceID: "wsp_1", RequesterID: "usr_1", SourceID: "src_1",
	}, repository.grepSnapshot)
	require.Len(t, matches, 1)
	raw, err := json.Marshal(matches[0])
	require.NoError(t, err)
	require.NotContains(t, string(raw), "revision")
	require.NotContains(t, string(raw), "citation")
}

func TestExecutionOpenRefBindsAttemptLocalImmutableSnapshot(t *testing.T) {
	now := time.Date(2026, 8, 4, 13, 0, 0, 0, time.UTC)
	handle := testRootHandle(now)
	handle.RepoRole = RepoRoleMirror
	handle.DefaultRef = "pingcap-master"
	repository := &fakeRepository{refSHA: strings.Repeat("b", 40)}
	session := authorizeTestExecution(t, &Service{
		Handles:    &memoryHandleStore{handles: map[string]protocol.RootHandle{handle.ID: handle}},
		Repository: repository,
		Clock:      func() time.Time { return now },
	}, handle.ID)

	snapshot, err := session.OpenRef(context.Background(), handle.ID, "v7.1.9-0", RefKindTag)
	require.NoError(t, err)
	require.Equal(t, "v7.1.9-0", repository.refName)
	require.Equal(t, RefKindTag, repository.refKind)
	require.Equal(t, "src_1", snapshot.SourceID)
	require.Equal(t, RepoRoleMirror, snapshot.RepoRole)
	require.Equal(t, "pingcap-master", snapshot.CurrentRef)
	require.Equal(t, "v7.1.9-0", snapshot.TargetRef)
	require.Equal(t, RefKindTag, snapshot.RefKind)
	require.NotEqual(t, handle.ID, snapshot.HandleID)
	require.Equal(t, strings.Repeat("b", 40), session.handles[snapshot.HandleID].Revision)
	require.Equal(t, handle.RepositoryRef, session.handles[snapshot.HandleID].RepositoryRef)
	require.Equal(t, "v7.1.9-0", session.handles[snapshot.HandleID].ResolvedRef)
	require.Equal(t, RefKindTag, session.handles[snapshot.HandleID].ResolvedRefKind)
}

func TestReadUsesSameFrozenRevisionAsCitation(t *testing.T) {
	now := time.Date(2026, 7, 25, 10, 0, 0, 0, time.UTC)
	repository := &fakeRepository{}
	service := &Service{
		Handles:    &memoryHandleStore{handles: map[string]protocol.RootHandle{"kqr_1": testRootHandle(now)}},
		Repository: repository,
		Clock:      func() time.Time { return now },
	}
	execution := authorizeTestExecution(t, service, "kqr_1")
	var chunks []ReadChunk
	err := execution.Read(context.Background(), ReadRequest{
		HandleID: "kqr_1", Path: "docs/design.md", LineStart: 10, LineEnd: 20,
	}, func(chunk ReadChunk) error {
		chunks = append(chunks, chunk)
		return nil
	})
	require.NoError(t, err)
	require.Equal(t, "sha-fixed", repository.readSnapshot.Revision)
	require.Len(t, chunks, 1)
	require.Equal(t, repository.readSnapshot.Revision, chunks[0].Citation.Revision)
	require.Equal(t, 10, chunks[0].Citation.LineStart)
	require.Equal(t, 20, chunks[0].Citation.LineEnd)
}

func TestExecutionRepositoryAuthorizesRootsOnceAcrossToolCalls(t *testing.T) {
	now := time.Date(2026, 7, 25, 10, 0, 0, 0, time.UTC)
	first := testRootHandle(now)
	second := first
	second.ID = "kqr_2"
	second.SourceID = "src_2"
	second.RepositoryRef = "repo-internal-2"
	handles := &memoryHandleStore{handles: map[string]protocol.RootHandle{first.ID: first, second.ID: second}}
	service := &Service{Handles: handles, Repository: &fakeRepository{}, Clock: func() time.Time { return now }}

	repository, roots, err := service.AuthorizeExecution(context.Background(), testPrincipal(), []string{first.ID, second.ID})
	require.NoError(t, err)
	require.Len(t, roots, 2)
	require.Equal(t, 1, handles.loadCalls)

	_, err = repository.Grep(context.Background(), GrepRequest{
		HandleIDs: []string{first.ID, second.ID}, Pattern: "server",
	}, func(GrepMatch) error { return nil })
	require.NoError(t, err)
	require.NoError(t, repository.List(context.Background(), ListRequest{HandleID: first.ID}, func(ListEntry) error { return nil }))
	require.NoError(t, repository.Read(context.Background(), ReadRequest{
		HandleID: second.ID, Path: "docs/design.md", LineStart: 1, LineEnd: 20,
	}, func(ReadChunk) error { return nil }))
	require.Equal(t, 1, handles.loadCalls, "tool calls must not reload authorized handles")
}

func TestHandleReplayAcrossOrganizationRequesterOrTurnIsRejectedUniformly(t *testing.T) {
	now := time.Date(2026, 7, 25, 10, 0, 0, 0, time.UTC)
	service := &Service{
		Handles:    &memoryHandleStore{handles: map[string]protocol.RootHandle{"kqr_1": testRootHandle(now)}},
		Repository: &fakeRepository{},
		Clock:      func() time.Time { return now },
	}
	for _, principal := range []protocol.Principal{
		{OrganizationID: "org_other", WorkspaceID: "wsp_1", RequesterID: "usr_1", TurnID: "kqt_1"},
		{OrganizationID: "org_1", WorkspaceID: "wsp_other", RequesterID: "usr_1", TurnID: "kqt_1"},
		{OrganizationID: "org_1", WorkspaceID: "wsp_1", RequesterID: "usr_other", TurnID: "kqt_1"},
		{OrganizationID: "org_1", WorkspaceID: "wsp_1", RequesterID: "usr_1", TurnID: "kqt_other"},
	} {
		_, _, err := service.AuthorizeExecution(context.Background(), principal, []string{"kqr_1"})
		var toolErr *Error
		require.ErrorAs(t, err, &toolErr)
		require.Equal(t, CodeHandleRejected, toolErr.Code)
		require.Equal(t, "Knowledge root handle was rejected", toolErr.Message)
	}
}

func TestReadRejectsTraversalBeforeCallingRepository(t *testing.T) {
	now := time.Date(2026, 7, 25, 10, 0, 0, 0, time.UTC)
	repository := &fakeRepository{}
	service := &Service{
		Handles:    &memoryHandleStore{handles: map[string]protocol.RootHandle{"kqr_1": testRootHandle(now)}},
		Repository: repository,
		Clock:      func() time.Time { return now },
	}
	execution := authorizeTestExecution(t, service, "kqr_1")
	err := execution.Read(context.Background(), ReadRequest{
		HandleID: "kqr_1", Path: "docs/../secrets.txt", LineStart: 1, LineEnd: 1,
	}, func(ReadChunk) error { return nil })
	var toolErr *Error
	require.ErrorAs(t, err, &toolErr)
	require.Equal(t, CodeInvalidRequest, toolErr.Code)
	require.Zero(t, repository.readCalls)
}

func TestGrepReportsResourceExhaustionWithoutClaimingCompleteSearch(t *testing.T) {
	now := time.Date(2026, 7, 25, 10, 0, 0, 0, time.UTC)
	service := &Service{
		Handles:    &memoryHandleStore{handles: map[string]protocol.RootHandle{"kqr_1": testRootHandle(now)}},
		Repository: &fakeRepository{grepErr: ErrRepositoryExhausted},
		Clock:      func() time.Time { return now },
	}
	execution := authorizeTestExecution(t, service, "kqr_1")
	completion, err := execution.Grep(context.Background(), GrepRequest{
		HandleIDs: []string{"kqr_1"}, Pattern: "x",
	}, func(GrepMatch) error { return nil })
	var toolErr *Error
	require.ErrorAs(t, err, &toolErr)
	require.Equal(t, CodeResourceExhausted, toolErr.Code)
	require.False(t, completion.Complete)
	require.Equal(t, CodeResourceExhausted, completion.Reason)
}

func TestGrepRejectsRegexOutsideGitPOSIXDialect(t *testing.T) {
	now := time.Date(2026, 7, 25, 10, 0, 0, 0, time.UTC)
	repository := &fakeRepository{}
	service := &Service{
		Handles:    &memoryHandleStore{handles: map[string]protocol.RootHandle{"kqr_1": testRootHandle(now)}},
		Repository: repository,
		Clock:      func() time.Time { return now },
	}
	execution := authorizeTestExecution(t, service, "kqr_1")

	_, err := execution.Grep(context.Background(), GrepRequest{
		HandleIDs: []string{"kqr_1"}, Pattern: "(?i)server",
	}, func(GrepMatch) error { return nil })
	var toolErr *Error
	require.ErrorAs(t, err, &toolErr)
	require.Equal(t, CodeInvalidRequest, toolErr.Code)
	require.Zero(t, repository.grepCalls)
}

func authorizeTestExecution(t *testing.T, service *Service, handleIDs ...string) *ExecutionRepository {
	t.Helper()
	repository, _, err := service.AuthorizeExecution(context.Background(), testPrincipal(), handleIDs)
	require.NoError(t, err)
	return repository
}

func testPrincipal() protocol.Principal {
	return protocol.Principal{OrganizationID: "org_1", WorkspaceID: "wsp_1", RequesterID: "usr_1", TurnID: "kqt_1"}
}

func testRootHandle(now time.Time) protocol.RootHandle {
	return protocol.RootHandle{
		ID: "kqr_1", OrganizationID: "org_1", WorkspaceID: "wsp_1",
		RequesterID: "usr_1", TurnID: "kqt_1", KnowledgeBaseID: "kb_1",
		SourceID: "src_1", RootKind: RootKindPrimary, RepoRole: RepoRolePrimary,
		Revision: "sha-fixed", RepositoryRef: "repo-internal-1",
		ExpiresAt: now.Add(time.Minute), CreatedAt: now,
	}
}

func TestRootHandleViewJSONNeverIncludesRepositoryRef(t *testing.T) {
	raw, err := json.Marshal(RootHandleView{
		HandleID: "kqr_1", KnowledgeBaseID: "kb_1", SourceID: "src_1",
		RootKind: RootKindPrimary, RepoRole: RepoRolePrimary, Revision: "sha-fixed",
	})
	require.NoError(t, err)
	require.NotContains(t, string(raw), "repository")
	require.NotContains(t, string(raw), "repo-internal")
}