package kb

import (
	"context"
	"errors"
	"strings"
	"time"

	"github.com/tidbcloud/AgentTeam/internal/agsbridge"
	"github.com/tidbcloud/AgentTeam/internal/identity/trusted"
	"github.com/tidbcloud/AgentTeam/internal/service/memory"
	"github.com/tidbcloud/AgentTeam/internal/service/resource/repos"
	"github.com/tidbcloud/AgentTeam/internal/store/kb"
	"gorm.io/gorm"
)

func (s *KBService) ListAgentAccessibleKnowledgeBases(ctx context.Context, workspaceID, subject, agentID string, page, perPage int) ([]AgentAccessibleKnowledgeBase, error) {
	if err := s.validateScope(workspaceID, subject); err != nil {
		return nil, err
	}
	agentID = strings.TrimSpace(agentID)
	if err := s.requireOwnedAgent(ctx, workspaceID, subject, agentID); err != nil {
		return nil, err
	}
	page, perPage = normalizeKBPagination(page, perPage)
	rows, err := s.store.ListAgentAccessibleKnowledgeBases(ctx, strings.TrimSpace(workspaceID), agentID, false, page, perPage)
	if err != nil {
		return nil, err
	}
	out := make([]AgentAccessibleKnowledgeBase, 0, len(rows))
	canAdministerWorkspace := s.canAdministerWorkspaceForRead(ctx, workspaceID, subject)
	for _, row := range rows {
		if strings.TrimSpace(row.KnowledgeBase.PublicationState) == KBPublicationDraft &&
			!s.knowledgeBaseVisibleToSubjectWithWorkspaceAuthority(ctx, row.KnowledgeBase, subject, canAdministerWorkspace) {
			continue
		}
		out = append(out, AgentAccessibleKnowledgeBase{
			ID:              row.KnowledgeBase.ID,
			WorkspaceID:     row.KnowledgeBase.WorkspaceID,
			Name:            row.KnowledgeBase.Name,
			Description:     row.KnowledgeBase.Description,
			AccessLevel:     agentKBAccessLevelFromGrant(row.Grant),
			PermissionSetAt: row.Grant.UpdatedAt,
			UpdatedAt:       row.KnowledgeBase.UpdatedAt,
		})
	}
	return out, nil
}

func (s *KBService) GrantAgentKnowledgeBaseAccess(ctx context.Context, workspaceID, subject, agentID, kbID, accessLevel string) (AgentKnowledgeBaseGrantResult, error) {
	if err := s.validateScope(workspaceID, subject); err != nil {
		return AgentKnowledgeBaseGrantResult{}, err
	}
	agentID = strings.TrimSpace(agentID)
	kbID = strings.TrimSpace(kbID)
	if agentID == "" {
		return AgentKnowledgeBaseGrantResult{}, invalidArgument("agent id is required")
	}
	if kbID == "" {
		return AgentKnowledgeBaseGrantResult{}, invalidArgument("knowledge base id is required")
	}
	if err := s.requireOwnedAgent(ctx, workspaceID, subject, agentID); err != nil {
		return AgentKnowledgeBaseGrantResult{}, err
	}
	normalizedAccessLevel, err := normalizeAgentKBAccessLevelInput(accessLevel)
	if err != nil {
		return AgentKnowledgeBaseGrantResult{}, err
	}
	base, err := s.store.GetKnowledgeBase(ctx, strings.TrimSpace(workspaceID), kbID)
	if err != nil {
		return AgentKnowledgeBaseGrantResult{}, mapStoreNotFound(err)
	}
	s.decorateKnowledgeBaseAccess(ctx, &base, subject)
	if !base.CanManage {
		return AgentKnowledgeBaseGrantResult{}, forbidden("knowledge base requires manage access")
	}
	if !humanCanGrantAgentKBAccessLevel(base, normalizedAccessLevel) {
		return AgentKnowledgeBaseGrantResult{}, forbidden("knowledge base access level exceeds current user access")
	}
	now := s.now()
	grant, err := s.store.UpsertKnowledgeBaseGrant(ctx, kb.KnowledgeBaseGrant{
		ID:                      newKBID("kbg"),
		WorkspaceID:             strings.TrimSpace(workspaceID),
		KnowledgeBaseID:         kbID,
		SubjectType:             "agent",
		SubjectID:               agentID,
		Role:                    normalizeKBGrantRole(normalizedAccessLevel, "agent"),
		RepositoryAccessProfile: repositoryAccessProfileForAgentKBAccessLevel(normalizedAccessLevel),
		CreatedBy:               strings.TrimSpace(subject),
		UpdatedBy:               strings.TrimSpace(subject),
		CreatedAt:               now,
		UpdatedAt:               now,
	})
	if err != nil {
		return AgentKnowledgeBaseGrantResult{}, err
	}
	if err := s.reconcileKnowledgeBaseGrantRepositoryAccess(ctx, base, grant, subject); err != nil {
		return AgentKnowledgeBaseGrantResult{}, err
	}
	if err := s.enqueueAgentKnowledgeBaseAccessChange(ctx, strings.TrimSpace(workspaceID), agentID, now); err != nil {
		return AgentKnowledgeBaseGrantResult{}, err
	}
	return AgentKnowledgeBaseGrantResult{WorkspaceID: strings.TrimSpace(workspaceID), AgentID: agentID, KnowledgeBaseID: kbID, AccessLevel: normalizedAccessLevel, CreatedBy: strings.TrimSpace(subject), UpdatedBy: strings.TrimSpace(subject), CreatedAt: now, UpdatedAt: now}, nil
}

func (s *KBService) SetAgentKnowledgeBaseAccessForKnowledgeBase(ctx context.Context, workspaceID, subject, agentID, kbID, accessLevel string) (AgentKnowledgeBaseGrantResult, error) {
	if strings.TrimSpace(accessLevel) != "none" {
		return s.GrantAgentKnowledgeBaseAccess(ctx, workspaceID, subject, agentID, kbID, accessLevel)
	}
	if err := s.validateScope(workspaceID, subject); err != nil {
		return AgentKnowledgeBaseGrantResult{}, err
	}
	workspaceID = strings.TrimSpace(workspaceID)
	subject = strings.TrimSpace(subject)
	agentID = strings.TrimSpace(agentID)
	kbID = strings.TrimSpace(kbID)
	if err := s.requireOwnedAgent(ctx, workspaceID, subject, agentID); err != nil {
		return AgentKnowledgeBaseGrantResult{}, err
	}
	grant, err := s.store.GetKnowledgeBaseGrant(ctx, workspaceID, kbID, "agent", agentID)
	if err != nil {
		if errors.Is(err, gorm.ErrRecordNotFound) {
			return AgentKnowledgeBaseGrantResult{WorkspaceID: workspaceID, AgentID: agentID, KnowledgeBaseID: kbID, AccessLevel: "none"}, nil
		}
		return AgentKnowledgeBaseGrantResult{}, err
	}
	if err := s.reconcileDesiredRepositoryAccessRows(ctx, grant, nil, "", kb.KnowledgeBaseRepositoryAccessProfileRead, subject, resourceRepositoryActorForAgent(agentID)); err != nil {
		return AgentKnowledgeBaseGrantResult{}, err
	}
	if err := s.blockedRepositoryAccessRevokeError(ctx, workspaceID, grant.ID); err != nil {
		return AgentKnowledgeBaseGrantResult{}, err
	}
	if err := s.store.DeleteKnowledgeBaseGrantByID(ctx, workspaceID, kbID, grant.ID); err != nil {
		return AgentKnowledgeBaseGrantResult{}, mapStoreNotFound(err)
	}
	now := s.now()
	if err := s.enqueueAgentKnowledgeBaseAccessChange(ctx, workspaceID, agentID, now); err != nil {
		return AgentKnowledgeBaseGrantResult{}, err
	}
	return AgentKnowledgeBaseGrantResult{WorkspaceID: workspaceID, AgentID: agentID, KnowledgeBaseID: kbID, AccessLevel: "none", UpdatedBy: subject, UpdatedAt: now}, nil
}

func (s *KBService) SetAgentKnowledgeBaseAccess(ctx context.Context, workspaceID, subject, agentID string, input AgentKnowledgeBaseBulkAccessMutation) (AgentKnowledgeBaseBulkAccessResult, error) {
	if err := s.validateScope(workspaceID, subject); err != nil {
		return AgentKnowledgeBaseBulkAccessResult{}, err
	}
	workspaceID = strings.TrimSpace(workspaceID)
	subject = strings.TrimSpace(subject)
	agentID = strings.TrimSpace(agentID)
	if err := s.requireOwnedAgent(ctx, workspaceID, subject, agentID); err != nil {
		return AgentKnowledgeBaseBulkAccessResult{}, err
	}
	requested, err := normalizeAgentKBAccessSelections(input)
	if err != nil {
		return AgentKnowledgeBaseBulkAccessResult{}, err
	}
	bases, err := s.store.ListKnowledgeBases(ctx, workspaceID, "", false, 1, 10000)
	if err != nil {
		return AgentKnowledgeBaseBulkAccessResult{}, err
	}
	if err := s.decorateKnowledgeBaseAccessList(ctx, bases, subject); err != nil {
		return AgentKnowledgeBaseBulkAccessResult{}, err
	}
	visible := make(map[string]kb.KnowledgeBase, len(bases))
	for _, base := range bases {
		if strings.TrimSpace(base.Role) != "" {
			visible[base.ID] = base
		}
	}
	for id, accessLevel := range requested {
		base, ok := visible[id]
		if !ok {
			return AgentKnowledgeBaseBulkAccessResult{}, forbidden("knowledge base requires view access")
		}
		if !humanCanGrantAgentKBAccessLevel(base, accessLevel) {
			return AgentKnowledgeBaseBulkAccessResult{}, forbidden("knowledge base access level exceeds current user access")
		}
	}
	current, err := s.store.ListAgentAccessibleKnowledgeBases(ctx, workspaceID, agentID, false, 1, 10000)
	if err != nil {
		return AgentKnowledgeBaseBulkAccessResult{}, err
	}
	currentGrantByKB := make(map[string]kb.KnowledgeBaseGrant, len(current))
	for _, row := range current {
		if _, ok := visible[row.KnowledgeBase.ID]; ok {
			currentGrantByKB[row.KnowledgeBase.ID] = row.Grant
		}
	}
	now := s.now()
	changed := false
	result := AgentKnowledgeBaseBulkAccessResult{AgentID: agentID, KnowledgeBaseIDs: make([]string, 0, len(requested))}
	for id, accessLevel := range requested {
		result.KnowledgeBaseIDs = append(result.KnowledgeBaseIDs, id)
		existing, exists := currentGrantByKB[id]
		grant, created, updated, err := s.upsertAgentKnowledgeBaseGrantSelection(ctx, workspaceID, subject, agentID, id, accessLevel, existing, exists, now)
		if err != nil {
			return AgentKnowledgeBaseBulkAccessResult{}, err
		}
		if !updated {
			continue
		}
		currentGrantByKB[id] = grant
		changed = true
		if created {
			result.GrantedKnowledgeBaseIDs = append(result.GrantedKnowledgeBaseIDs, id)
		}
	}
	for id, grant := range currentGrantByKB {
		if _, keep := requested[id]; keep {
			continue
		}
		if err := s.reconcileDesiredRepositoryAccessRows(ctx, grant, nil, "", kb.KnowledgeBaseRepositoryAccessProfileRead, subject, resourceRepositoryActorForAgent(grant.SubjectID)); err != nil {
			return AgentKnowledgeBaseBulkAccessResult{}, err
		}
		if err := s.blockedRepositoryAccessRevokeError(ctx, workspaceID, grant.ID); err != nil {
			return AgentKnowledgeBaseBulkAccessResult{}, err
		}
		if err := s.store.DeleteKnowledgeBaseGrantByID(ctx, workspaceID, id, grant.ID); err != nil {
			return AgentKnowledgeBaseBulkAccessResult{}, mapStoreNotFound(err)
		}
		changed = true
		result.RevokedKnowledgeBaseIDs = append(result.RevokedKnowledgeBaseIDs, id)
	}
	for id := range requested {
		grant := currentGrantByKB[id]
		if err := s.reconcileKnowledgeBaseGrantRepositoryAccess(ctx, visible[id], grant, subject); err != nil {
			return AgentKnowledgeBaseBulkAccessResult{}, err
		}
	}
	if changed {
		if err := s.enqueueAgentKnowledgeBaseAccessChange(ctx, workspaceID, agentID, now); err != nil {
			return AgentKnowledgeBaseBulkAccessResult{}, err
		}
	}
	return result, nil
}

func (s *KBService) upsertAgentKnowledgeBaseGrantSelection(ctx context.Context, workspaceID, subject, agentID, kbID, accessLevel string, existing kb.KnowledgeBaseGrant, exists bool, now time.Time) (kb.KnowledgeBaseGrant, bool, bool, error) {
	repositoryAccessProfile := repositoryAccessProfileForAgentKBAccessLevel(accessLevel)
	if exists {
		if kb.NormalizeKnowledgeBaseRepositoryAccessProfile(existing.RepositoryAccessProfile) == repositoryAccessProfile {
			return existing, false, false, nil
		}
		grant, err := s.store.UpsertKnowledgeBaseGrant(ctx, kb.KnowledgeBaseGrant{
			ID:                      existing.ID,
			WorkspaceID:             workspaceID,
			KnowledgeBaseID:         kbID,
			SubjectType:             "agent",
			SubjectID:               agentID,
			Role:                    "view",
			RepositoryAccessProfile: repositoryAccessProfile,
			CreatedBy:               existing.CreatedBy,
			UpdatedBy:               subject,
			CreatedAt:               existing.CreatedAt,
			UpdatedAt:               now,
		})
		return grant, false, true, err
	}
	grant, err := s.store.UpsertKnowledgeBaseGrant(ctx, kb.KnowledgeBaseGrant{
		ID:                      newKBID("kbg"),
		WorkspaceID:             workspaceID,
		KnowledgeBaseID:         kbID,
		SubjectType:             "agent",
		SubjectID:               agentID,
		Role:                    "view",
		RepositoryAccessProfile: repositoryAccessProfile,
		CreatedBy:               subject,
		UpdatedBy:               subject,
		CreatedAt:               now,
		UpdatedAt:               now,
	})
	return grant, true, true, err
}

func (s *KBService) listOwnedAgents(ctx context.Context, workspaceID, subject string) ([]memory.AgentOwnership, error) {
	if s.agentOwnershipSource == nil {
		return nil, invalidArgument("agent ownership source is not configured")
	}
	agents, err := s.agentOwnershipSource.ListOwnedAgents(ctx, strings.TrimSpace(workspaceID), strings.TrimSpace(subject))
	if err != nil {
		return nil, err
	}
	out := make([]memory.AgentOwnership, 0, len(agents))
	seen := make(map[string]struct{}, len(agents))
	for _, agent := range agents {
		id := strings.TrimSpace(agent.ID)
		if id == "" {
			continue
		}
		if _, ok := seen[id]; ok {
			continue
		}
		seen[id] = struct{}{}
		agent.ID = id
		agent.Name = strings.TrimSpace(agent.Name)
		agent.DisplayName = strings.TrimSpace(agent.DisplayName)
		out = append(out, agent)
	}
	return out, nil
}

func (s *KBService) requireOwnedAgent(ctx context.Context, workspaceID, subject, agentID string) error {
	agentID = strings.TrimSpace(agentID)
	if agentID == "" {
		return invalidArgument("agent id is required")
	}
	agents, err := s.listOwnedAgents(ctx, workspaceID, subject)
	if err != nil {
		return err
	}
	for _, agent := range agents {
		if agent.ID == agentID {
			return nil
		}
	}
	return forbidden("agent must be owned by the current user")
}

func resourceRepositoryActorForAgent(agentID string) repos.ResourceRepositoryActor {
	agentID = strings.TrimSpace(agentID)
	login := agsbridge.LoginForAgentTeam(trusted.Identity{Subject: agentID})
	return repos.ResourceRepositoryActor{
		Provider: agsbridge.Provider,
		Subject:  agentID,
		Login:    login,
		Name:     agentID,
	}
}

func normalizeAgentKBAccessLevelInput(accessLevel string) (string, error) {
	switch strings.TrimSpace(accessLevel) {
	case "", "read":
		return "read", nil
	case "write":
		return "write", nil
	default:
		return "", invalidArgument("agent knowledge base access level must be read or write")
	}
}

func normalizeAgentKBAccessSelections(input AgentKnowledgeBaseBulkAccessMutation) (map[string]string, error) {
	requested := make(map[string]string, len(input.KnowledgeBases)+len(input.KnowledgeBaseIDs))
	if len(input.KnowledgeBases) > 0 {
		for _, item := range input.KnowledgeBases {
			id := strings.TrimSpace(item.KnowledgeBaseID)
			if id == "" {
				continue
			}
			accessLevel, err := normalizeAgentKBAccessLevelInput(item.AccessLevel)
			if err != nil {
				return nil, err
			}
			requested[id] = accessLevel
		}
		return requested, nil
	}
	for _, id := range input.KnowledgeBaseIDs {
		id = strings.TrimSpace(id)
		if id != "" {
			requested[id] = "read"
		}
	}
	return requested, nil
}

func agentKBAccessLevelFromGrant(grant kb.KnowledgeBaseGrant) string {
	switch kb.NormalizeKnowledgeBaseRepositoryAccessProfile(grant.RepositoryAccessProfile) {
	case kb.KnowledgeBaseRepositoryAccessProfileWrite, kb.KnowledgeBaseRepositoryAccessProfileManage:
		return "write"
	default:
		return "read"
	}
}

func repositoryAccessProfileForAgentKBAccessLevel(accessLevel string) string {
	if strings.TrimSpace(accessLevel) == "write" {
		return kb.KnowledgeBaseRepositoryAccessProfileWrite
	}
	return kb.KnowledgeBaseRepositoryAccessProfileRead
}

func humanCanGrantAgentKBAccessLevel(base kb.KnowledgeBase, accessLevel string) bool {
	role := strings.TrimSpace(base.Role)
	if role == "" {
		return false
	}
	if strings.TrimSpace(accessLevel) != "write" {
		return true
	}
	if roleCanEdit(role) {
		return true
	}
	return repositoryAccessPermissionRank(
		repositoryPermissionForAccessProfile(base.RepositoryAccessProfile),
	) >= repositoryAccessPermissionRank(
		repositoryAccessProfileForAgentKBAccessLevel("write"),
	)
}