package repository

import (
	"context"
	"testing"

	"github.com/Tencent/WeKnora/internal/accesspolicy"
	"github.com/Tencent/WeKnora/internal/types"
	"github.com/Tencent/WeKnora/internal/types/interfaces"
	"github.com/stretchr/testify/require"
	"gorm.io/driver/sqlite"
	"gorm.io/gorm"
)

func fileAccessFixture(t *testing.T) (*resourceRepository, *gorm.DB, string) {
	t.Helper()
	db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
	require.NoError(t, err)
	sqlDB, err := db.DB()
	require.NoError(t, err)
	sqlDB.SetMaxOpenConns(1)
	t.Cleanup(func() { _ = sqlDB.Close() })
	require.NoError(t, db.AutoMigrate(&types.StoredResource{}, &types.ResourceBinding{}))
	for _, schema := range []string{
		"CREATE TABLE knowledges (id TEXT PRIMARY KEY, tenant_id INTEGER, knowledge_base_id TEXT, file_path TEXT, deleted_at DATETIME)",
		"CREATE TABLE knowledge_bases (id TEXT PRIMARY KEY, tenant_id INTEGER, creator_id TEXT, is_temporary BOOLEAN DEFAULT FALSE, version_state TEXT DEFAULT '', deleted_at DATETIME)",
		"CREATE TABLE sessions (id TEXT PRIMARY KEY, tenant_id INTEGER, user_id TEXT, description TEXT, deleted_at DATETIME)",
		"CREATE TABLE messages (id TEXT PRIMARY KEY, session_id TEXT, knowledge_references TEXT, deleted_at DATETIME)",
		"CREATE TABLE temporary_documents (id TEXT PRIMARY KEY, tenant_id INTEGER, session_id TEXT, expires_at DATETIME, deleted_at DATETIME)",
		"CREATE TABLE custom_agents (id TEXT PRIMARY KEY, tenant_id INTEGER, enabled BOOLEAN, config TEXT, deleted_at DATETIME)",
		"CREATE TABLE embed_channels (id TEXT PRIMARY KEY, tenant_id INTEGER, agent_id TEXT, enabled BOOLEAN, deleted_at DATETIME)",
	} {
		require.NoError(t, db.Exec(schema).Error)
	}
	resource := &types.StoredResource{ID: "file-1", Handle: "AbCdEfGhIjKlMnOpQrStUv", TenantID: 7, PhysicalPath: "local://7/exports/private.png", Provider: "local", LocationHash: "unused"}
	require.NoError(t, db.Create(resource).Error)
	return &resourceRepository{db: db}, db, types.BuildResourcePath(resource.Handle)
}

func fileUserContext(user string) context.Context {
	ctx := context.WithValue(context.Background(), types.TenantIDContextKey, uint64(7))
	ctx = context.WithValue(ctx, types.UserIDContextKey, user)
	return accesspolicy.WithSnapshot(ctx, &accesspolicy.Snapshot{Publications: map[string]bool{"knowledge:kb-a": false}})
}

func TestResourceFileAccessKnowledgeBinding(t *testing.T) {
	r, db, ref := fileAccessFixture(t)
	require.NoError(t, db.Exec("INSERT INTO knowledge_bases (id,tenant_id,creator_id) VALUES ('kb-a',7,'alice'),('kb-b',7,'bob')").Error)
	require.NoError(t, db.Exec("INSERT INTO knowledges (id,tenant_id,knowledge_base_id,file_path) VALUES ('doc-a',7,'kb-a','local://7/legacy.doc')").Error)
	require.NoError(t, r.CreateBinding(context.Background(), &types.ResourceBinding{ResourceID: "file-1", TenantID: 7, OwnerType: types.ResourceOwnerKnowledge, OwnerID: "doc-a"}))
	scope := interfaces.ResourceFileScope{TenantID: 7}
	require.NoError(t, r.AuthorizeFileRead(fileUserContext("alice"), ref, scope))
	require.Error(t, r.AuthorizeFileRead(fileUserContext("bob"), ref, scope))
	require.Error(t, r.AuthorizeFileRead(fileUserContext("bob"), "local://7/legacy.doc", scope))
	require.NoError(t, r.AuthorizeFileRead(fileUserContext("alice"), "local://7/legacy.doc", scope))
	require.Error(t, r.AuthorizeFileRead(fileUserContext("alice"), "local://7/unregistered.png", scope))
	require.Error(t, r.AuthorizeFileRead(fileUserContext("alice"), ref, interfaces.ResourceFileScope{TenantID: 8}))
	require.Error(t, r.AuthorizeFileRead(fileUserContext("bob"), ref, interfaces.ResourceFileScope{TenantID: 7, KnowledgeBaseID: "kb-b"}))
	require.NoError(t, r.AuthorizeFileRead(fileUserContext("bob"), ref, interfaces.ResourceFileScope{TenantID: 7, KnowledgeBaseID: "kb-a"}))
	shared := accesspolicy.WithSnapshot(fileUserContext("bob"), &accesspolicy.Snapshot{Publications: map[string]bool{"knowledge:kb-a": false}, Grants: map[string]bool{"knowledge:kb-a": true}})
	require.NoError(t, r.AuthorizeFileRead(shared, ref, scope))
	denied := accesspolicy.WithSnapshot(fileUserContext("alice"), &accesspolicy.Snapshot{Restricted: true, Permissions: map[string]bool{}})
	require.Error(t, r.AuthorizeFileRead(denied, ref, scope))
	require.NoError(t, db.Exec("UPDATE knowledges SET deleted_at=CURRENT_TIMESTAMP WHERE id='doc-a'").Error)
	require.Error(t, r.AuthorizeFileRead(fileUserContext("alice"), ref, scope))
}

func TestResourceFileAccessSharedOriginalCannotBeDownloaded(t *testing.T) {
	r, db, ref := fileAccessFixture(t)
	require.NoError(t, db.Exec("INSERT INTO knowledge_bases (id,tenant_id,creator_id) VALUES ('kb-a',7,'alice')").Error)
	require.NoError(t, db.Exec("INSERT INTO knowledges (id,tenant_id,knowledge_base_id,file_path) VALUES ('doc-a',7,'kb-a','local://7/legacy.doc')").Error)
	require.NoError(t, r.CreateBinding(context.Background(), &types.ResourceBinding{ResourceID: "file-1", TenantID: 7, OwnerType: types.ResourceOwnerKnowledge, OwnerID: "doc-a", Relation: types.ResourceRelationSourceFile}))
	require.NoError(t, r.CreateBinding(context.Background(), &types.ResourceBinding{ResourceID: "file-1", TenantID: 7, OwnerType: types.ResourceOwnerKnowledge, OwnerID: "doc-a", Relation: types.ResourceRelationAttachment}))
	require.NoError(t, db.Exec("INSERT INTO sessions (id,tenant_id,user_id) VALUES ('s-b',7,'bob')").Error)
	require.NoError(t, db.Exec("INSERT INTO messages (id,session_id,knowledge_references) VALUES ('msg-b','s-b',?)", []byte(`[{"knowledge_id":"doc-a"}]`)).Error)
	shared := accesspolicy.WithSnapshot(fileUserContext("bob"), &accesspolicy.Snapshot{Publications: map[string]bool{"knowledge:kb-a": true}, Grants: map[string]bool{"knowledge:kb-a": true}})
	for _, scope := range []interfaces.ResourceFileScope{
		{TenantID: 7},
		{TenantID: 7, KnowledgeBaseID: "kb-a"},
		{TenantID: 7, MessageID: "msg-b", SessionID: "s-b"},
		{TenantID: 7, EmbedChannelID: "channel-a"},
	} {
		require.ErrorIs(t, r.AuthorizeFileRead(shared, ref, scope), errFileAccessDenied)
		require.ErrorIs(t, r.AuthorizeFileRead(shared, "local://7/legacy.doc", scope), errFileAccessDenied)
	}
	scope := interfaces.ResourceFileScope{TenantID: 7}
	require.NoError(t, r.AuthorizeFileRead(fileUserContext("alice"), ref, scope))
	admin := context.WithValue(fileUserContext("admin"), types.TenantRoleContextKey, types.TenantRoleAdmin)
	require.NoError(t, r.AuthorizeFileRead(admin, ref, scope))
	systemAdmin := context.WithValue(fileUserContext("system"), types.SystemAdminContextKey, true)
	require.NoError(t, r.AuthorizeFileRead(systemAdmin, ref, scope))
	crossTenant := context.WithValue(systemAdmin, types.TenantIDContextKey, uint64(8))
	require.ErrorIs(t, r.AuthorizeFileRead(crossTenant, ref, scope), errFileAccessDenied)
	denied := accesspolicy.WithSnapshot(fileUserContext("alice"), &accesspolicy.Snapshot{Restricted: true, Permissions: map[string]bool{}})
	require.ErrorIs(t, r.AuthorizeFileRead(denied, ref, interfaces.ResourceFileScope{TenantID: 7, KnowledgeBaseID: "kb-a"}), errFileAccessDenied)

	// A legacy registration without a source relation is still an original.
	require.NoError(t, db.Exec("DELETE FROM resource_bindings WHERE relation=?", types.ResourceRelationSourceFile).Error)
	for _, location := range []string{ref, "local://7/exports/private.png"} {
		require.NoError(t, db.Exec("UPDATE knowledges SET file_path=? WHERE id='doc-a'", location).Error)
		require.ErrorIs(t, r.AuthorizeFileRead(shared, ref, scope), errFileAccessDenied)
		require.NoError(t, r.AuthorizeFileRead(fileUserContext("alice"), ref, scope))
	}
}

func TestResourceFileAccessPrivateMessages(t *testing.T) {
	r, db, ref := fileAccessFixture(t)
	require.NoError(t, db.Exec("INSERT INTO sessions (id,tenant_id,user_id) VALUES ('s-a',7,'alice'),('s-b',7,'bob'),('s-legacy',7,'')").Error)
	require.NoError(t, db.Exec("INSERT INTO messages (id,session_id) VALUES ('msg-a','s-a'),('msg-b','s-b'),('msg-legacy','s-legacy')").Error)
	binding := &types.ResourceBinding{ResourceID: "file-1", TenantID: 7, OwnerType: types.ResourceOwnerMessage, OwnerID: "msg-a"}
	require.NoError(t, r.CreateBinding(context.Background(), binding))
	scope := interfaces.ResourceFileScope{TenantID: 7}
	require.NoError(t, r.AuthorizeFileRead(fileUserContext("alice"), ref, scope))
	require.Error(t, r.AuthorizeFileRead(fileUserContext("bob"), ref, scope))
	require.Error(t, r.AuthorizeFileRead(fileUserContext("alice"), ref, interfaces.ResourceFileScope{TenantID: 7, MessageID: "msg-b", SessionID: "s-b"}))
	require.NoError(t, r.AuthorizeFileRead(fileUserContext("alice"), ref, interfaces.ResourceFileScope{TenantID: 7, MessageID: "msg-a", SessionID: "s-a"}))
	require.Error(t, r.AuthorizeFileRead(fileUserContext("alice"), ref, interfaces.ResourceFileScope{TenantID: 7, KnowledgeBaseID: "kb-a"}))
	require.NoError(t, db.Model(&types.ResourceBinding{}).Where("id = ?", binding.ID).Update("owner_id", "msg-legacy").Error)
	require.Error(t, r.AuthorizeFileRead(fileUserContext("bob"), ref, scope))
}

func TestResourceFileAccessEmbedKnowledgeScope(t *testing.T) {
	r, db, ref := fileAccessFixture(t)
	require.NoError(t, db.Exec("INSERT INTO knowledge_bases (id,tenant_id,creator_id) VALUES ('kb-a',7,'alice')").Error)
	require.NoError(t, db.Exec("INSERT INTO knowledges (id,tenant_id,knowledge_base_id) VALUES ('doc-a',7,'kb-a')").Error)
	require.NoError(t, r.CreateBinding(context.Background(), &types.ResourceBinding{ResourceID: "file-1", TenantID: 7, OwnerType: types.ResourceOwnerKnowledge, OwnerID: "doc-a"}))
	require.NoError(t, db.Exec("INSERT INTO custom_agents (id,tenant_id,enabled,config) VALUES ('agent-a',7,TRUE,?)", `{"kb_selection_mode":"selected","knowledge_bases":["kb-a"]}`).Error)
	require.NoError(t, db.Exec("INSERT INTO embed_channels (id,tenant_id,agent_id,enabled) VALUES ('channel-a',7,'agent-a',TRUE)").Error)
	scope := interfaces.ResourceFileScope{TenantID: 7, EmbedChannelID: "channel-a"}
	require.NoError(t, r.AuthorizeFileRead(fileUserContext("visitor"), ref, scope))
	for _, config := range []string{`{"kb_selection_mode":"none"}`, `{"kb_selection_mode":"selected","knowledge_bases":["kb-other"]}`, `{"kb_selection_mode":"selected","knowledge_bases":[]}`} {
		require.NoError(t, db.Exec("UPDATE custom_agents SET config=?", config).Error)
		require.Error(t, r.AuthorizeFileRead(fileUserContext("visitor"), ref, scope))
	}
}

func TestResourceFileAccessMessageKnowledgeReferences(t *testing.T) {
	r, db, ref := fileAccessFixture(t)
	require.NoError(t, db.Exec("INSERT INTO knowledges (id,tenant_id,knowledge_base_id) VALUES ('doc-a',7,'kb-a')").Error)
	require.NoError(t, r.CreateBinding(context.Background(), &types.ResourceBinding{ResourceID: "file-1", TenantID: 7, OwnerType: types.ResourceOwnerKnowledge, OwnerID: "doc-a"}))
	require.NoError(t, db.Exec("INSERT INTO sessions (id,tenant_id,user_id) VALUES ('s-a',7,'alice')").Error)
	require.NoError(t, db.Exec("INSERT INTO messages (id,session_id,knowledge_references) VALUES ('msg-a','s-a',?)", []byte(`[{"knowledge_id":"doc-a"}]`)).Error)
	scope := interfaces.ResourceFileScope{TenantID: 7, MessageID: "msg-a", SessionID: "s-a"}
	require.NoError(t, r.AuthorizeFileRead(fileUserContext("alice"), ref, scope))
	require.Error(t, r.AuthorizeFileRead(fileUserContext("bob"), ref, scope))
	require.NoError(t, db.Exec("UPDATE messages SET knowledge_references=?", []byte(`[{"knowledge_id":"doc-unrelated"}]`)).Error)
	require.Error(t, r.AuthorizeFileRead(fileUserContext("alice"), ref, scope))
}

func TestResourceFileAccessEmbedPrivateSession(t *testing.T) {
	r, db, ref := fileAccessFixture(t)
	principal := types.EmbedSessionPrincipal(7, "channel-a", "s-a")
	require.NoError(t, db.Exec("INSERT INTO sessions (id,tenant_id,user_id) VALUES ('s-a',7,?)", principal.StorageID()).Error)
	require.NoError(t, db.Exec("INSERT INTO messages (id,session_id) VALUES ('msg-a','s-a')").Error)
	require.NoError(t, r.CreateBinding(context.Background(), &types.ResourceBinding{ResourceID: "file-1", TenantID: 7, OwnerType: types.ResourceOwnerMessage, OwnerID: "msg-a"}))
	ctx := types.WithPrincipal(fileUserContext("embed"), principal)
	scope := interfaces.ResourceFileScope{TenantID: 7, EmbedChannelID: "channel-a", SessionID: "s-a"}
	require.NoError(t, r.AuthorizeFileRead(ctx, ref, scope))
	require.Error(t, r.AuthorizeFileRead(fileUserContext("embed"), ref, scope))
	scope.SessionID = "s-other"
	require.Error(t, r.AuthorizeFileRead(ctx, ref, scope))
}