package kb

import (
	"context"
	"fmt"
	"strings"

	"github.com/tidbcloud/AgentTeam/internal/store/content"
)

type kbSourceWriteRoute struct {
	native     bool
	generation uint64
}

func (s *KBService) beginKBSourceWrite(ctx context.Context, workspaceID, sourceID string) (kbSourceWriteRoute, error) {
	if s == nil || s.contentAuthority == nil {
		return kbSourceWriteRoute{}, nil
	}
	rows, err := s.contentAuthority.ListAuthoritiesForBindings(ctx, strings.TrimSpace(workspaceID), kbSourceContentDomain, []string{strings.TrimSpace(sourceID)})
	if err != nil {
		return kbSourceWriteRoute{}, fmt.Errorf("resolve KB source authority: %w", err)
	}
	authority, found := rows[strings.TrimSpace(sourceID)]
	if !found {
		return kbSourceWriteRoute{}, nil
	}
	switch strings.TrimSpace(authority.State) {
	case content.AuthorityStateLegacyActive:
		return kbSourceWriteRoute{generation: authority.Generation}, nil
	case content.AuthorityStateNativeActive:
		return kbSourceWriteRoute{native: true, generation: authority.Generation}, nil
	default:
		return kbSourceWriteRoute{}, codedConflict("kb-source-authority-invalid", "KB source authority is invalid", nil)
	}
}

func (s *KBService) ensureKBSourceWriteCurrent(ctx context.Context, workspaceID, sourceID string, route kbSourceWriteRoute) error {
	if s == nil || s.contentAuthority == nil || (route.generation == 0 && !route.native) {
		return nil
	}
	now, err := s.beginKBSourceWrite(ctx, workspaceID, sourceID)
	if err != nil {
		return err
	}
	if now.native != route.native || now.generation != route.generation {
		return codedConflict("kb-source-authority-changed", "KB source authority changed during mutation", nil)
	}
	return nil
}

// kbSourceAuthorityState resolves one exact source binding. Missing authority
// rows preserve the legacy default; persisted unknown states fail closed.
func (s *KBService) kbSourceAuthorityState(ctx context.Context, workspaceID, sourceID string) (string, error) {
	workspaceID = strings.TrimSpace(workspaceID)
	sourceID = strings.TrimSpace(sourceID)
	if workspaceID == "" || sourceID == "" {
		return "", invalidArgument("KB source authority identity is required")
	}
	native, err := s.kbSourceNativeBindings(ctx, workspaceID, []string{sourceID})
	if err != nil {
		return "", err
	}
	if _, ok := native[sourceID]; ok {
		return content.AuthorityStateNativeActive, nil
	}
	return content.AuthorityStateLegacyActive, nil
}

func (s *KBService) kbSourceNativeActive(ctx context.Context, workspaceID, sourceID string) (bool, error) {
	state, err := s.kbSourceAuthorityState(ctx, workspaceID, sourceID)
	return state == content.AuthorityStateNativeActive, err
}

func nativeKBSourceLegacyOperationConflict() error {
	return codedConflict("kb-source-native-active", "Legacy KB source operations are unavailable after native activation", nil)
}

// kbSourceNativeBindings resolves a caller-supplied exact set in bounded
// batches. Missing rows remain legacy; unknown persisted states fail closed.
func (s *KBService) kbSourceNativeBindings(ctx context.Context, workspaceID string, sourceIDs []string) (map[string]struct{}, error) {
	workspaceID = strings.TrimSpace(workspaceID)
	if workspaceID == "" {
		return nil, invalidArgument("KB source authority workspace is required")
	}
	keys := make([]string, 0, len(sourceIDs))
	seen := make(map[string]struct{}, len(sourceIDs))
	for _, sourceID := range sourceIDs {
		sourceID = strings.TrimSpace(sourceID)
		if sourceID == "" {
			return nil, invalidArgument("KB source authority identity is required")
		}
		if _, ok := seen[sourceID]; ok {
			continue
		}
		seen[sourceID] = struct{}{}
		keys = append(keys, sourceID)
	}
	native := make(map[string]struct{})
	if s == nil || s.contentAuthority == nil || len(keys) == 0 {
		return native, nil
	}
	for start := 0; start < len(keys); start += content.MaxAuthorityBindingBatch {
		end := min(start+content.MaxAuthorityBindingBatch, len(keys))
		authorities, err := s.contentAuthority.ListAuthoritiesForBindings(ctx, workspaceID, kbSourceContentDomain, keys[start:end])
		if err != nil {
			return nil, fmt.Errorf("resolve KB source authorities: %w", err)
		}
		for _, sourceID := range keys[start:end] {
			authority, ok := authorities[sourceID]
			if !ok {
				continue
			}
			switch strings.TrimSpace(authority.State) {
			case content.AuthorityStateLegacyActive:
			case content.AuthorityStateNativeActive:
				native[sourceID] = struct{}{}
			default:
				return nil, codedConflict("kb-source-authority-invalid", "KB source authority is invalid", nil)
			}
		}
	}
	return native, nil
}