package kb

import (
	"github.com/tidbcloud/AgentTeam/internal/kbgovernance"

	"context"
	"fmt"
	"strings"

	"gorm.io/gorm"
	"gorm.io/gorm/clause"
)

func (s *Store) loadKnowledgeBasePublicationSnapshot(ctx context.Context, base KnowledgeBase, workingRevision string, lock bool) (PublicationSnapshotResult, error) {
	// Every caller supplies one already-authorized KB and an exact working
	// revision. Source hydration stays bounded by the KB product limit.
	query := func() *gorm.DB {
		result := s.db.WithContext(ctx)
		if lock {
			result = result.Clauses(clause.Locking{Strength: "UPDATE"})
		}
		return result
	}
	var sources []Source
	if err := query().Where("workspace_id = ? AND knowledge_base_id = ? AND deleted_at IS NULL", base.WorkspaceID, base.ID).Order("id ASC").Limit(101).Find(&sources).Error; err != nil {
		return PublicationSnapshotResult{}, err
	}
	if len(sources) > 100 {
		return PublicationSnapshotResult{}, fmt.Errorf("publication snapshot exceeds the source limit")
	}
	var tags []KnowledgeBaseTagAssignment
	if err := query().Where("workspace_id = ? AND knowledge_base_id = ?", base.WorkspaceID, base.ID).Order("tag_id ASC").Find(&tags).Error; err != nil {
		return PublicationSnapshotResult{}, err
	}
	result := PublicationSnapshotResult{SourceCount: len(sources), Snapshot: kbgovernance.PublicationSnapshot{
		KnowledgeBaseID: base.ID, WorkingRevision: strings.TrimSpace(workingRevision), IsActive: base.IsActive,
		Name: base.Name, Description: base.Description, GroupID: base.GroupID,
		TagIDs: make([]string, 0, len(tags)), Sources: make([]kbgovernance.PublicationSnapshotSource, 0, len(sources)),
	}}
	for _, tag := range tags {
		result.Snapshot.TagIDs = append(result.Snapshot.TagIDs, tag.TagID)
	}
	sourceIDs := make([]string, 0, len(sources))
	for _, source := range sources {
		sourceIDs = append(sourceIDs, source.ID)
	}
	var revisions []SourceRevision
	var artifacts []KnowledgeBaseSetupSourceArtifact
	if len(sourceIDs) > 0 {
		if err := query().Where("workspace_id = ? AND knowledge_base_id = ? AND source_id IN ? AND is_current = ?", base.WorkspaceID, base.ID, sourceIDs, true).Order("source_id ASC, updated_at DESC, id DESC").Find(&revisions).Error; err != nil {
			return PublicationSnapshotResult{}, err
		}
		if err := query().Where("workspace_id = ? AND knowledge_base_id = ? AND source_id IN ?", base.WorkspaceID, base.ID, sourceIDs).Find(&artifacts).Error; err != nil {
			return PublicationSnapshotResult{}, err
		}
	}
	revisionBySource := make(map[string]SourceRevision, len(revisions))
	for _, revision := range revisions {
		if _, exists := revisionBySource[revision.SourceID]; !exists {
			revisionBySource[revision.SourceID] = revision
		}
	}
	artifactBySourceKind := make(map[string]string, len(artifacts))
	for _, artifact := range artifacts {
		artifactBySourceKind[artifact.SourceID+":"+artifact.Kind] = artifact.State
	}
	for _, source := range sources {
		row := kbgovernance.PublicationSnapshotSource{
			ID: source.ID, SourceType: source.SourceType, Name: source.Name, Description: source.Description,
			DefaultRef: source.DefaultRef, IndexingMode: source.IndexingMode, Status: source.Status,
			SyncStatus: source.SyncStatus, IndexStatus: source.IndexStatus, Config: source.Config,
			WikiDisposition:      artifactBySourceKind[source.ID+":"+GenerationKindWiki],
			DashboardDisposition: artifactBySourceKind[source.ID+":dashboard"],
		}
		if revision, exists := revisionBySource[source.ID]; exists {
			row.RevisionID, row.SourceCommitSHA = revision.ID, revision.SourceCommitSHA
			row.WikiRevision, row.DashboardRevision = revision.PublishedWikiRevision, revision.UADashboardCommitSHA
		}
		if row.WikiRevision != "" {
			result.WikiReadyCount++
		} else if row.WikiDisposition == "skipped" {
			result.SkippedCount++
		}
		if row.DashboardRevision != "" {
			result.DashboardReadyCount++
		} else if row.DashboardDisposition == "skipped" {
			result.SkippedCount++
		}
		result.Snapshot.Sources = append(result.Snapshot.Sources, row)
	}
	encoded, digest, err := kbgovernance.EncodePublicationSnapshot(result.Snapshot)
	if err != nil {
		return PublicationSnapshotResult{}, err
	}
	result.JSON, result.Digest = encoded, digest
	return result, nil
}