package kb

import (
	"context"
	"errors"
	"strings"
	"sync/atomic"
	"testing"
	"time"

	"github.com/stretchr/testify/require"
	"github.com/tidbcloud/AgentTeam/internal/service/content"
	storekb "github.com/tidbcloud/AgentTeam/internal/store/kb"
	"gorm.io/driver/sqlite"
	"gorm.io/gorm"
)

type kbContentWorkspaceAccess struct {
	allowed map[string]bool
	admins  map[string]bool
	err     error
}

func (access *kbContentWorkspaceAccess) HasWorkspaceAccess(_ context.Context, workspaceID, subject string) (bool, error) {
	return access.allowed[workspaceID+"\x00"+subject], access.err
}

func (access *kbContentWorkspaceAccess) HasAgentWorkspaceAccess(_ context.Context, workspaceID, subject string) (bool, error) {
	return access.allowed[workspaceID+"\x00"+subject], access.err
}

func (access *kbContentWorkspaceAccess) CanAdministerWorkspace(_ context.Context, workspaceID, subject string) (bool, error) {
	return access.admins[workspaceID+"\x00"+subject], access.err
}

func TestKBSourceContentAccessPolicyRechecksLiveWorkspaceAndKBGrant(t *testing.T) {
	db, err := gorm.Open(sqlite.Open("file:"+t.Name()+"?mode=memory&cache=shared"), &gorm.Config{})
	require.NoError(t, err)
	require.NoError(t, db.AutoMigrate(&storekb.KnowledgeBase{}, &storekb.Source{}, &storekb.KnowledgeBaseGrant{}, &storekb.KnowledgeBaseSourceRepositoryAccessOverride{}, &storekb.SourceContentMapping{}))
	now := time.Now().UTC()
	require.NoError(t, db.Create(&storekb.KnowledgeBase{
		ID: "kb_1", WorkspaceID: "wsp_1", Name: "Runbooks", Description: "ops", IsActive: true,
		PublicationState: "official", DefaultWorkspaceMemberRole: "view", DefaultRepositoryAccessProfile: storekb.KnowledgeBaseRepositoryAccessProfileRead,
		CreatedBy: "usr_owner", UpdatedBy: "usr_owner", CreatedAt: now, UpdatedAt: now,
	}).Error)
	require.NoError(t, db.Create(&storekb.Source{
		ID: "src_1", WorkspaceID: "wsp_1", KnowledgeBaseID: "kb_1", SourceType: "codebase",
		Name: "AgentTeam", CreatedBy: "usr_owner", UpdatedBy: "usr_owner", CreatedAt: now, UpdatedAt: now,
	}).Error)
	require.NoError(t, db.Create(&storekb.KnowledgeBaseGrant{
		ID: "kbg_edit", WorkspaceID: "wsp_1", KnowledgeBaseID: "kb_1", SubjectType: "user",
		SubjectID: "usr_editor", Role: "edit", RepositoryAccessProfile: storekb.KnowledgeBaseRepositoryAccessProfileWrite, CreatedBy: "usr_owner", UpdatedBy: "usr_owner", CreatedAt: now, UpdatedAt: now,
	}).Error)
	require.NoError(t, db.Create(&storekb.KnowledgeBaseGrant{
		ID: "kbg_agent", WorkspaceID: "wsp_1", KnowledgeBaseID: "kb_1", SubjectType: "agent",
		SubjectID: "agt_reader", Role: "view", RepositoryAccessProfile: storekb.KnowledgeBaseRepositoryAccessProfileRead, CreatedBy: "usr_owner", UpdatedBy: "usr_owner", CreatedAt: now, UpdatedAt: now,
	}).Error)
	require.NoError(t, db.Create(&storekb.SourceContentMapping{
		ID: "kbm_1", WorkspaceID: "wsp_1", KnowledgeBaseID: "kb_1", SourceID: "src_1", SourceRevisionID: "kbr_1",
		RepositoryRole: "canonical", RepositoryID: "repo_1", CanonicalBranch: "main", CommitSHA: "abc123",
		ContentID: "ctn_1", ContentVersionID: "ctv_1", RootDigest: strings.Repeat("a", 64), FileCount: 1, SizeBytes: 1,
		SourceCommittedAt: now, VerifiedAt: &now, CreatedAt: now, UpdatedAt: now,
	}).Error)
	access := &kbContentWorkspaceAccess{
		allowed: map[string]bool{
			"wsp_1\x00usr_owner": true, "wsp_1\x00usr_editor": true,
			"wsp_1\x00usr_viewer": true, "wsp_1\x00usr_admin": true, "wsp_1\x00agt_reader": true,
		},
		admins: map[string]bool{"wsp_1\x00usr_admin": true},
	}
	policy, err := NewKBSourceContentAccessPolicy(access, access, storekb.New(db))
	require.NoError(t, err)
	request := func(subject string, action content.AccessAction) content.AccessRequest {
		principal := content.AccessPrincipal{
			Subject: subject, SubjectType: kbContentSubjectUser, WorkspaceID: "wsp_1", Domain: kbSourceContentDomain,
			BindingKey: "src_1", ContentID: "ctn_1", Generation: 2, Actions: []content.AccessAction{action},
		}
		row := content.AccessRequest{
			Principal: principal, WorkspaceID: "wsp_1", Domain: kbSourceContentDomain,
			BindingKey: "src_1", ContentID: "ctn_1", Generation: 2, Action: action,
		}
		switch action {
		case content.AccessReadPayload, content.AccessReadArchive:
			row.VersionIDs = []string{"ctv_1"}
		case content.AccessCompare:
			row.VersionIDs = []string{"ctv_1", "ctv_1"}
		}
		return row
	}

	require.NoError(t, policy.AuthorizeContentAccess(t.Context(), request("usr_owner", content.AccessPublish)))
	dashboardPublish := request("usr_owner", content.AccessPublish)
	dashboardPublish.Domain = kbUADashboardContentDomain
	dashboardPublish.Principal.Domain = kbUADashboardContentDomain
	require.NoError(t, policy.AuthorizeContentAccess(t.Context(), dashboardPublish))
	require.NoError(t, policy.AuthorizeContentAccess(t.Context(), request("usr_editor", content.AccessPublish)))
	require.NoError(t, policy.AuthorizeContentAccess(t.Context(), request("usr_viewer", content.AccessReadPayload)))
	require.ErrorIs(t, policy.AuthorizeContentAccess(t.Context(), request("usr_viewer", content.AccessPublish)), content.ErrAccessDenied)
	require.NoError(t, policy.AuthorizeContentAccess(t.Context(), request("usr_admin", content.AccessManagePin)))
	agentRead := request("agt_reader", content.AccessReadPayload)
	agentRead.Principal.SubjectType = kbContentSubjectAgent
	require.NoError(t, policy.AuthorizeContentAccess(t.Context(), agentRead))
	agentPublish := request("agt_reader", content.AccessPublish)
	agentPublish.Principal.SubjectType = kbContentSubjectAgent
	require.ErrorIs(t, policy.AuthorizeContentAccess(t.Context(), agentPublish), content.ErrAccessDenied)

	// Live product intent, not the AGS-era materialization projection, is the
	// native authorization source. Canonical and mirror versions may share one
	// ContentID, so exact VersionIDs must remain independently authorized.
	require.NoError(t, db.Create(&storekb.KnowledgeBaseSourceRepositoryAccessOverride{
		ID:          storekb.KnowledgeBaseSourceRepositoryAccessOverrideID("wsp_1", "kb_1", "kbg_edit", "src_1", "canonical"),
		WorkspaceID: "wsp_1", KnowledgeBaseID: "kb_1", GrantID: "kbg_edit", SourceID: "src_1", RepoRole: "canonical",
		SubjectType: "user", SubjectID: "usr_editor", RepositoryAccessProfile: storekb.KnowledgeBaseRepositoryAccessProfileNone,
		CreatedBy: "usr_owner", UpdatedBy: "usr_owner", CreatedAt: now, UpdatedAt: now,
	}).Error)
	require.ErrorIs(t, policy.AuthorizeContentAccess(t.Context(), request("usr_editor", content.AccessReadPayload)), content.ErrAccessDenied)
	dashboardRead := request("usr_editor", content.AccessReadPayload)
	dashboardRead.Domain = kbUADashboardContentDomain
	dashboardRead.Principal.Domain = kbUADashboardContentDomain
	require.NoError(t, policy.AuthorizeContentAccess(t.Context(), dashboardRead))
	require.NoError(t, db.Create(&storekb.SourceContentMapping{
		ID: "kbm_2", WorkspaceID: "wsp_1", KnowledgeBaseID: "kb_1", SourceID: "src_1", SourceRevisionID: "kbr_2",
		RepositoryRole: "mirror", RepositoryID: "repo_2", CanonicalBranch: "main", CommitSHA: "def456",
		ContentID: "ctn_1", ContentVersionID: "ctv_2", RootDigest: strings.Repeat("b", 64), FileCount: 1, SizeBytes: 1,
		SourceCommittedAt: now, VerifiedAt: &now, CreatedAt: now, UpdatedAt: now,
	}).Error)
	mirrorRead := request("usr_editor", content.AccessReadPayload)
	mirrorRead.VersionIDs = []string{"ctv_2"}
	require.NoError(t, policy.AuthorizeContentAccess(t.Context(), mirrorRead), "canonical deny must not revoke the independently allowed mirror version")
	canonicalMirrorCompare := request("usr_editor", content.AccessCompare)
	canonicalMirrorCompare.VersionIDs = []string{"ctv_1", "ctv_2"}
	require.ErrorIs(t, policy.AuthorizeContentAccess(t.Context(), canonicalMirrorCompare), content.ErrAccessDenied, "compare must require both exact version roles")
	require.ErrorIs(t, policy.AuthorizeContentAccess(t.Context(), request("usr_editor", content.AccessReadHistory)), content.ErrAccessDenied, "lineage-wide reads must require every represented repository role")

	require.NoError(t, db.Model(&storekb.KnowledgeBaseSourceRepositoryAccessOverride{}).
		Where("grant_id = ? AND source_id = ? AND repo_role = ?", "kbg_edit", "src_1", "canonical").
		Update("repository_access_profile", storekb.KnowledgeBaseRepositoryAccessProfileWrite).Error)
	require.NoError(t, db.Create(&storekb.KnowledgeBaseSourceRepositoryAccessOverride{
		ID:          storekb.KnowledgeBaseSourceRepositoryAccessOverrideID("wsp_1", "kb_1", "kbg_edit", "src_1", "mirror"),
		WorkspaceID: "wsp_1", KnowledgeBaseID: "kb_1", GrantID: "kbg_edit", SourceID: "src_1", RepoRole: "mirror",
		SubjectType: "user", SubjectID: "usr_editor", RepositoryAccessProfile: storekb.KnowledgeBaseRepositoryAccessProfileNone,
		CreatedBy: "usr_owner", UpdatedBy: "usr_owner", CreatedAt: now, UpdatedAt: now,
	}).Error)
	require.NoError(t, policy.AuthorizeContentAccess(t.Context(), request("usr_editor", content.AccessReadPayload)), "mirror deny must not revoke the independently allowed canonical version")
	require.ErrorIs(t, policy.AuthorizeContentAccess(t.Context(), mirrorRead), content.ErrAccessDenied)
	require.ErrorIs(t, policy.AuthorizeContentAccess(t.Context(), canonicalMirrorCompare), content.ErrAccessDenied)

	require.NoError(t, db.Model(&storekb.KnowledgeBaseSourceRepositoryAccessOverride{}).
		Where("grant_id = ? AND source_id = ? AND repo_role = ?", "kbg_edit", "src_1", "mirror").
		Update("repository_access_profile", storekb.KnowledgeBaseRepositoryAccessProfileWrite).Error)
	var queryCount atomic.Int64
	require.NoError(t, db.Callback().Query().Before("gorm:query").Register("test:count-native-kb-acl-queries", func(*gorm.DB) {
		queryCount.Add(1)
	}))
	queryCount.Store(0)
	require.NoError(t, policy.AuthorizeContentAccess(t.Context(), request("usr_editor", content.AccessReadPayload)))
	singleVersionQueries := queryCount.Load()
	queryCount.Store(0)
	require.NoError(t, policy.AuthorizeContentAccess(t.Context(), canonicalMirrorCompare))
	require.Equal(t, singleVersionQueries, queryCount.Load(), "compare must batch both exact version mappings instead of querying once per version")
	require.EqualValues(t, 5, singleVersionQueries, "native ACL must remain a fixed set of exact product reads")
	require.NoError(t, db.Model(&storekb.SourceContentMapping{}).Where("id = ?", "kbm_2").Update("repository_role", "unknown").Error)
	require.ErrorIs(t, policy.AuthorizeContentAccess(t.Context(), mirrorRead), content.ErrAccessDenied)
	require.ErrorIs(t, policy.AuthorizeContentAccess(t.Context(), request("usr_editor", content.AccessReadHistory)), content.ErrAccessDenied, "lineage-wide access must fail closed on an unknown persisted role")
	require.NoError(t, db.Model(&storekb.SourceContentMapping{}).Where("id = ?", "kbm_2").Update("repository_role", "mirror").Error)

	for name, mutate := range map[string]func(*content.AccessRequest){
		"wrong subject type": func(row *content.AccessRequest) { row.Principal.SubjectType = "agent" },
		"forged workspace":   func(row *content.AccessRequest) { row.Principal.WorkspaceID = "wsp_other" },
		"wrong domain":       func(row *content.AccessRequest) { row.Domain = "skill" },
		"wrong binding":      func(row *content.AccessRequest) { row.Principal.BindingKey = "src_other" },
		"wrong content":      func(row *content.AccessRequest) { row.Principal.ContentID = "ctn_other" },
		"unknown action": func(row *content.AccessRequest) {
			row.Action = content.AccessAction("unknown")
		},
	} {
		t.Run(name, func(t *testing.T) {
			row := request("usr_owner", content.AccessReadPayload)
			mutate(&row)
			require.ErrorIs(t, policy.AuthorizeContentAccess(t.Context(), row), content.ErrAccessDenied)
		})
	}

	require.NoError(t, db.Model(&storekb.KnowledgeBaseGrant{}).
		Where("workspace_id = ? AND knowledge_base_id = ? AND subject_id = ?", "wsp_1", "kb_1", "usr_editor").
		Update("role", "none").Error)
	require.ErrorIs(t, policy.AuthorizeContentAccess(t.Context(), request("usr_editor", content.AccessReadPayload)), content.ErrAccessDenied)
	require.NoError(t, db.Model(&storekb.KnowledgeBaseGrant{}).
		Where("workspace_id = ? AND knowledge_base_id = ? AND subject_id = ?", "wsp_1", "kb_1", "agt_reader").
		Update("role", "none").Error)
	require.ErrorIs(t, policy.AuthorizeContentAccess(t.Context(), agentRead), content.ErrAccessDenied)
	access.allowed["wsp_1\x00usr_owner"] = false
	require.ErrorIs(t, policy.AuthorizeContentAccess(t.Context(), request("usr_owner", content.AccessReadPayload)), content.ErrAccessDenied)
	access.allowed["wsp_1\x00usr_owner"] = true
	require.NoError(t, db.Model(&storekb.Source{}).Where("workspace_id = ? AND id = ?", "wsp_1", "src_1").Update("deleted_at", now).Error)
	require.ErrorIs(t, policy.AuthorizeContentAccess(t.Context(), request("usr_owner", content.AccessReadPayload)), content.ErrAccessDenied)
}

func TestKBSourceContentAccessPolicyFailsClosedWhenDependenciesFail(t *testing.T) {
	access := &kbContentWorkspaceAccess{allowed: map[string]bool{"wsp_1\x00usr_1": true}, err: errors.New("workspace unavailable")}
	policy, err := NewKBSourceContentAccessPolicy(access, access, storekb.New(nil))
	require.NoError(t, err)
	request := content.AccessRequest{
		Principal: content.AccessPrincipal{
			Subject: "usr_1", SubjectType: kbContentSubjectUser, WorkspaceID: "wsp_1", Domain: kbSourceContentDomain,
			BindingKey: "src_1", Actions: []content.AccessAction{content.AccessReadPayload},
		},
		WorkspaceID: "wsp_1", Domain: kbSourceContentDomain, BindingKey: "src_1", Action: content.AccessReadPayload,
	}
	require.ErrorContains(t, policy.AuthorizeContentAccess(t.Context(), request), "workspace unavailable")
}
