package knowledgeqa

import (
	"github.com/tidbcloud/AgentTeam/internal/knowledgeqa/protocol"
	"github.com/tidbcloud/AgentTeam/internal/store/model"

	"context"
	"errors"
	"fmt"
	"strings"
	"time"

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

func (store *Store) PutRootHandles(ctx context.Context, handles []protocol.RootHandle) error {
	if store == nil || store.runtimeDB == nil {
		return fmt.Errorf("Knowledge Q&A root handle store is not configured")
	}
	if len(handles) == 0 {
		return nil
	}
	rows := make([]RootHandleRecord, 0, len(handles))
	seen := make(map[string]struct{}, len(handles))
	for _, handle := range handles {
		row, err := recordFromHandle(handle)
		if err != nil {
			return err
		}
		if _, exists := seen[row.ID]; exists {
			return fmt.Errorf("Knowledge Q&A root handle id is duplicated")
		}
		seen[row.ID] = struct{}{}
		rows = append(rows, row)
	}
	return store.runtimeDB.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
		if err := tx.Create(&rows).Error; err != nil {
			return fmt.Errorf("persist Knowledge Q&A root handles: %w", err)
		}
		return nil
	})
}

// LoadRootHandles loads one explicit execution's handles through a single
// authority-scoped query. Missing or foreign IDs remain indistinguishable to
// the service authorization boundary.
func (store *Store) LoadRootHandles(
	ctx context.Context,
	principal protocol.Principal,
	handleIDs []string,
) ([]protocol.RootHandle, error) {
	if store == nil || store.runtimeDB == nil {
		return nil, fmt.Errorf("Knowledge Q&A root handle store is not configured")
	}
	principal = normalizeAccountingPrincipal(principal)
	handleIDs = normalizeHandleIDs(handleIDs)
	if principal.OrganizationID == "" || principal.WorkspaceID == "" ||
		principal.RequesterID == "" || principal.TurnID == "" || len(handleIDs) == 0 {
		return nil, fmt.Errorf("Knowledge Q&A root handle authority is incomplete")
	}
	var rows []RootHandleRecord
	if err := store.runtimeDB.WithContext(ctx).
		Where(`organization_id = ? AND workspace_id = ? AND requester_id = ? AND turn_id = ?`,
			principal.OrganizationID, principal.WorkspaceID, principal.RequesterID, principal.TurnID).
		Where("id IN ?", handleIDs).
		Find(&rows).Error; err != nil {
		return nil, fmt.Errorf("load Knowledge Q&A execution root handles: %w", err)
	}
	handles := make([]protocol.RootHandle, 0, len(rows))
	for _, row := range rows {
		handles = append(handles, handleFromRecord(row))
	}
	return handles, nil
}

// DeleteRootHandles removes only the capabilities consumed by one execution.
// Retries share a turn ID, so turn-wide deletion could revoke a newer attempt's
// handles while a canceled prior attempt is still unwinding.
func (store *Store) DeleteRootHandles(ctx context.Context, principal protocol.Principal, handleIDs []string) error {
	if store == nil || store.runtimeDB == nil {
		return fmt.Errorf("Knowledge Q&A root handle store is not configured")
	}
	principal.OrganizationID = strings.TrimSpace(principal.OrganizationID)
	principal.WorkspaceID = strings.TrimSpace(principal.WorkspaceID)
	principal.RequesterID = strings.TrimSpace(principal.RequesterID)
	principal.TurnID = strings.TrimSpace(principal.TurnID)
	if principal.OrganizationID == "" || principal.WorkspaceID == "" ||
		principal.RequesterID == "" || principal.TurnID == "" {
		return fmt.Errorf("Knowledge Q&A root handle principal is incomplete")
	}
	handleIDs = normalizeHandleIDs(handleIDs)
	if len(handleIDs) == 0 {
		return nil
	}
	if err := store.runtimeDB.WithContext(ctx).
		Where("organization_id = ? AND workspace_id = ? AND requester_id = ? AND turn_id = ?",
			principal.OrganizationID, principal.WorkspaceID, principal.RequesterID, principal.TurnID).
		Where("id IN ?", handleIDs).
		Delete(&RootHandleRecord{}).Error; err != nil {
		return fmt.Errorf("delete Knowledge Q&A execution root handles: %w", err)
	}
	return nil
}

func normalizeHandleIDs(input []string) []string {
	result := make([]string, 0, len(input))
	seen := make(map[string]struct{}, len(input))
	for _, value := range input {
		value = strings.TrimSpace(value)
		if value == "" {
			continue
		}
		if _, exists := seen[value]; exists {
			continue
		}
		seen[value] = struct{}{}
		result = append(result, value)
	}
	return result
}

func (store *Store) DeleteExpiredRootHandles(ctx context.Context, now time.Time, limit int, workspaceScope ...[]string) (int64, error) {
	if store == nil || store.runtimeDB == nil {
		return 0, fmt.Errorf("Knowledge Q&A root handle store is not configured")
	}
	if limit <= 0 {
		limit = 1000
	}
	var ids []string
	query := store.runtimeDB.WithContext(ctx).Model(&RootHandleRecord{}).
		Where("expires_at <= ?", now.UTC())
	query = applyWorkspaceScope(query, workspaceScope)
	if err := query.
		Order("expires_at ASC, id ASC").
		Limit(limit).
		Pluck("id", &ids).Error; err != nil {
		return 0, fmt.Errorf("list expired Knowledge Q&A root handles: %w", err)
	}
	if len(ids) == 0 {
		return 0, nil
	}
	result := store.runtimeDB.WithContext(ctx).Where("id IN ?", ids).Delete(&RootHandleRecord{})
	if result.Error != nil {
		return 0, fmt.Errorf("delete expired Knowledge Q&A root handles: %w", result.Error)
	}
	return result.RowsAffected, nil
}

// PutAccountingCheckpoint advances one attempt's content-free accounting
// aggregate. Stale generations are ignored, equal generations must be exact
// replays, requester authority can never be rebound, and terminal checkpoints
// are immutable.
func (store *Store) PutAccountingCheckpoint(ctx context.Context, checkpoint protocol.AccountingCheckpoint) error {
	if store == nil || store.runtimeDB == nil {
		return fmt.Errorf("Knowledge Q&A accounting checkpoint store is not configured")
	}
	row, err := accountingRecordFromCheckpoint(checkpoint)
	if err != nil {
		return err
	}
	return store.runtimeDB.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
		if err := tx.Clauses(clause.OnConflict{DoNothing: true}).Create(&row).Error; err != nil {
			return fmt.Errorf("create Knowledge Q&A accounting checkpoint: %w", err)
		}
		var current AccountingCheckpointRecord
		if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).
			Where("attempt_id = ?", row.AttemptID).Take(&current).Error; err != nil {
			return fmt.Errorf("lock Knowledge Q&A accounting checkpoint: %w", err)
		}
		if !sameAccountingAuthority(current, row) {
			return fmt.Errorf("Knowledge Q&A accounting checkpoint authority conflicts with the existing attempt")
		}
		if row.Generation < current.Generation {
			return nil
		}
		if row.Generation == current.Generation {
			if row.UsageJSON != current.UsageJSON || row.SettlementRefsJSON != current.SettlementRefsJSON {
				return fmt.Errorf("Knowledge Q&A accounting checkpoint generation conflicts with the existing aggregate")
			}
			if current.State == protocol.AccountingCheckpointTerminal || row.State == current.State {
				return nil
			}
			return tx.Model(&AccountingCheckpointRecord{}).
				Where("attempt_id = ?", row.AttemptID).
				Updates(map[string]any{
					"state":      protocol.AccountingCheckpointTerminal,
					"updated_at": row.UpdatedAt,
				}).Error
		}
		if current.State == protocol.AccountingCheckpointTerminal {
			return fmt.Errorf("Knowledge Q&A terminal accounting checkpoint cannot advance")
		}
		if err := validateAccountingProgression(current, row); err != nil {
			return err
		}
		if err := tx.Model(&AccountingCheckpointRecord{}).
			Where("attempt_id = ?", row.AttemptID).
			Updates(map[string]any{
				"generation":           row.Generation,
				"state":                row.State,
				"usage_json":           row.UsageJSON,
				"settlement_refs_json": row.SettlementRefsJSON,
				"updated_at":           row.UpdatedAt,
			}).Error; err != nil {
			return fmt.Errorf("advance Knowledge Q&A accounting checkpoint: %w", err)
		}
		return nil
	})
}

// LoadAccountingCheckpoint returns only an unexpired checkpoint whose full
// organization, workspace, requester, turn, and attempt authority matches.
func (store *Store) LoadAccountingCheckpoint(
	ctx context.Context,
	principal protocol.Principal,
	attemptID string,
	now time.Time,
) (protocol.AccountingCheckpoint, bool, error) {
	if store == nil || store.runtimeDB == nil {
		return protocol.AccountingCheckpoint{}, false,
			fmt.Errorf("Knowledge Q&A accounting checkpoint store is not configured")
	}
	principal = normalizeAccountingPrincipal(principal)
	attemptID = strings.TrimSpace(attemptID)
	if principal.OrganizationID == "" || principal.WorkspaceID == "" ||
		principal.RequesterID == "" || principal.TurnID == "" || attemptID == "" {
		return protocol.AccountingCheckpoint{}, false,
			fmt.Errorf("Knowledge Q&A accounting recovery authority is incomplete")
	}
	var row AccountingCheckpointRecord
	err := store.runtimeDB.WithContext(ctx).
		Where(`attempt_id = ? AND organization_id = ? AND workspace_id = ?
AND requester_id = ? AND turn_id = ? AND expires_at > ?`,
			attemptID, principal.OrganizationID, principal.WorkspaceID,
			principal.RequesterID, principal.TurnID, now.UTC()).
		Take(&row).Error
	if errors.Is(err, gorm.ErrRecordNotFound) {
		return protocol.AccountingCheckpoint{}, false, nil
	}
	if err != nil {
		return protocol.AccountingCheckpoint{}, false,
			fmt.Errorf("load Knowledge Q&A accounting checkpoint: %w", err)
	}
	checkpoint, err := accountingCheckpointFromRecord(row)
	if err != nil {
		return protocol.AccountingCheckpoint{}, false, err
	}
	return checkpoint, true, nil
}

// DeleteExpiredAccountingCheckpoints deletes one bounded expiry-ordered batch.
func (store *Store) DeleteExpiredAccountingCheckpoints(ctx context.Context, now time.Time, limit int, workspaceScope ...[]string) (int64, error) {
	if store == nil || store.runtimeDB == nil {
		return 0, fmt.Errorf("Knowledge Q&A accounting checkpoint store is not configured")
	}
	if limit <= 0 {
		limit = 1000
	}
	var ids []string
	query := store.runtimeDB.WithContext(ctx).Model(&AccountingCheckpointRecord{}).
		Where("expires_at <= ?", now.UTC())
	query = applyWorkspaceScope(query, workspaceScope)
	if err := query.Order("expires_at ASC, attempt_id ASC").
		Limit(limit).Pluck("attempt_id", &ids).Error; err != nil {
		return 0, fmt.Errorf("list expired Knowledge Q&A accounting checkpoints: %w", err)
	}
	if len(ids) == 0 {
		return 0, nil
	}
	result := store.runtimeDB.WithContext(ctx).Where("attempt_id IN ?", ids).Delete(&AccountingCheckpointRecord{})
	if result.Error != nil {
		return 0, fmt.Errorf("delete expired Knowledge Q&A accounting checkpoints: %w", result.Error)
	}
	return result.RowsAffected, nil
}

func applyWorkspaceScope(query *gorm.DB, scopes [][]string) *gorm.DB {
	if len(scopes) == 0 || scopes[0] == nil {
		return query
	}
	workspaceIDs := make([]string, 0, len(scopes[0]))
	seen := make(map[string]struct{}, len(scopes[0]))
	for _, workspaceID := range scopes[0] {
		workspaceID = strings.TrimSpace(workspaceID)
		if workspaceID == "" {
			continue
		}
		if _, ok := seen[workspaceID]; ok {
			continue
		}
		seen[workspaceID] = struct{}{}
		workspaceIDs = append(workspaceIDs, workspaceID)
	}
	if len(workspaceIDs) == 0 {
		return query.Where("1 = 0")
	}
	return query.Where("workspace_id IN ?", workspaceIDs)
}

func accountingRecordFromCheckpoint(checkpoint protocol.AccountingCheckpoint) (AccountingCheckpointRecord, error) {
	checkpoint.Principal = normalizeAccountingPrincipal(checkpoint.Principal)
	checkpoint.AttemptID = strings.TrimSpace(checkpoint.AttemptID)
	settlementRefs, err := normalizeAccountingSettlementRefs(checkpoint.SettlementRefs)
	if err != nil {
		return AccountingCheckpointRecord{}, err
	}
	checkpoint.SettlementRefs = settlementRefs
	if checkpoint.AttemptID == "" || checkpoint.Principal.OrganizationID == "" ||
		checkpoint.Principal.WorkspaceID == "" || checkpoint.Principal.RequesterID == "" ||
		checkpoint.Principal.TurnID == "" || checkpoint.Generation < 0 ||
		(checkpoint.State != protocol.AccountingCheckpointRunning &&
			checkpoint.State != protocol.AccountingCheckpointTerminal) ||
		checkpoint.ExpiresAt.IsZero() || checkpoint.UpdatedAt.IsZero() {
		return AccountingCheckpointRecord{}, fmt.Errorf("Knowledge Q&A accounting checkpoint is incomplete")
	}
	usageJSON, err := model.MarshalJSONObject("Knowledge Q&A accounting usage", nonNilAccountingUsage(checkpoint.Usage))
	if err != nil {
		return AccountingCheckpointRecord{}, fmt.Errorf("Knowledge Q&A accounting usage is invalid or too large")
	}
	settlementRefsJSON, err := model.MarshalJSONObject("Knowledge Q&A settlement lineage", checkpoint.SettlementRefs)
	if err != nil {
		return AccountingCheckpointRecord{}, fmt.Errorf("Knowledge Q&A settlement lineage is invalid")
	}
	if err := validateAccountingDocumentSizes(
		usageJSON, settlementRefsJSON, len(checkpoint.SettlementRefs),
	); err != nil {
		return AccountingCheckpointRecord{}, err
	}
	return AccountingCheckpointRecord{
		AttemptID: checkpoint.AttemptID, OrganizationID: checkpoint.Principal.OrganizationID,
		WorkspaceID: checkpoint.Principal.WorkspaceID, RequesterID: checkpoint.Principal.RequesterID,
		TurnID: checkpoint.Principal.TurnID, Generation: checkpoint.Generation, State: checkpoint.State,
		UsageJSON: string(usageJSON), SettlementRefsJSON: settlementLineageJSON(settlementRefsJSON),
		ExpiresAt: checkpoint.ExpiresAt.UTC(), CreatedAt: checkpoint.UpdatedAt.UTC(), UpdatedAt: checkpoint.UpdatedAt.UTC(),
	}, nil
}

func validateAccountingDocumentSizes(usageJSON, settlementRefsJSON []byte, settlementRefCount int) error {
	if len(usageJSON) > protocol.MaxAccountingUsageJSONBytes {
		return fmt.Errorf("Knowledge Q&A accounting usage is invalid or too large")
	}
	if settlementRefCount > protocol.MaxAccountingSettlementReferences ||
		len(settlementRefsJSON) > protocol.MaxAccountingSettlementLineageJSONBytes {
		return fmt.Errorf("Knowledge Q&A settlement lineage exceeds its storage boundary")
	}
	return nil
}

func accountingCheckpointFromRecord(row AccountingCheckpointRecord) (protocol.AccountingCheckpoint, error) {
	var usage map[string]any
	if err := model.UnmarshalJSONString("Knowledge Q&A accounting usage", row.UsageJSON, &usage); err != nil {
		return protocol.AccountingCheckpoint{}, fmt.Errorf("decode Knowledge Q&A accounting usage: %w", err)
	}
	var settlementRefs []string
	if err := model.UnmarshalJSONString("Knowledge Q&A settlement lineage", string(row.SettlementRefsJSON), &settlementRefs); err != nil {
		return protocol.AccountingCheckpoint{}, fmt.Errorf("decode Knowledge Q&A settlement lineage: %w", err)
	}
	return protocol.AccountingCheckpoint{
		Principal: protocol.Principal{
			OrganizationID: row.OrganizationID, WorkspaceID: row.WorkspaceID,
			RequesterID: row.RequesterID, TurnID: row.TurnID,
		},
		AttemptID: row.AttemptID, Generation: row.Generation, State: row.State,
		Usage: usage, SettlementRefs: settlementRefs,
		ExpiresAt: row.ExpiresAt, UpdatedAt: row.UpdatedAt,
	}, nil
}

func validateAccountingProgression(current, next AccountingCheckpointRecord) error {
	var currentUsage, nextUsage map[string]any
	if err := model.UnmarshalJSONString("existing Knowledge Q&A accounting usage", current.UsageJSON, &currentUsage); err != nil {
		return fmt.Errorf("decode existing Knowledge Q&A accounting usage: %w", err)
	}
	if err := model.UnmarshalJSONString("next Knowledge Q&A accounting usage", next.UsageJSON, &nextUsage); err != nil {
		return fmt.Errorf("decode next Knowledge Q&A accounting usage: %w", err)
	}
	if !accountingUsageContains(nextUsage, currentUsage) {
		return fmt.Errorf("Knowledge Q&A accounting usage cannot regress")
	}
	var currentRefs, nextRefs []string
	if err := model.UnmarshalJSONString("existing Knowledge Q&A settlement lineage", string(current.SettlementRefsJSON), &currentRefs); err != nil {
		return fmt.Errorf("decode existing Knowledge Q&A settlement lineage: %w", err)
	}
	if err := model.UnmarshalJSONString("next Knowledge Q&A settlement lineage", string(next.SettlementRefsJSON), &nextRefs); err != nil {
		return fmt.Errorf("decode next Knowledge Q&A settlement lineage: %w", err)
	}
	if len(nextRefs) < len(currentRefs) {
		return fmt.Errorf("Knowledge Q&A settlement lineage cannot regress")
	}
	for index := range currentRefs {
		if nextRefs[index] != currentRefs[index] {
			return fmt.Errorf("Knowledge Q&A settlement lineage cannot be reordered")
		}
	}
	return nil
}

func accountingUsageContains(current, previous map[string]any) bool {
	for key, previousValue := range previous {
		currentValue, exists := current[key]
		if !exists {
			return false
		}
		previousNumber, previousNumeric := previousValue.(float64)
		currentNumber, currentNumeric := currentValue.(float64)
		if previousNumeric {
			if !currentNumeric || currentNumber < previousNumber {
				return false
			}
			continue
		}
		previousMap, previousNested := previousValue.(map[string]any)
		if previousNested {
			currentMap, currentNested := currentValue.(map[string]any)
			if !currentNested || !accountingUsageContains(currentMap, previousMap) {
				return false
			}
		}
	}
	return true
}

func sameAccountingAuthority(left, right AccountingCheckpointRecord) bool {
	return left.OrganizationID == right.OrganizationID && left.WorkspaceID == right.WorkspaceID &&
		left.RequesterID == right.RequesterID && left.TurnID == right.TurnID
}

func normalizeAccountingPrincipal(principal protocol.Principal) protocol.Principal {
	principal.OrganizationID = strings.TrimSpace(principal.OrganizationID)
	principal.WorkspaceID = strings.TrimSpace(principal.WorkspaceID)
	principal.RequesterID = strings.TrimSpace(principal.RequesterID)
	principal.TurnID = strings.TrimSpace(principal.TurnID)
	return principal
}

func normalizeAccountingSettlementRefs(input []string) ([]string, error) {
	result := make([]string, 0, len(input))
	seen := make(map[string]struct{}, len(input))
	for _, value := range input {
		value = strings.TrimSpace(value)
		if value == "" {
			continue
		}
		if len(value) > protocol.MaxAccountingSettlementReferenceBytes {
			return nil, fmt.Errorf("Knowledge Q&A settlement lineage exceeds its storage boundary")
		}
		if _, exists := seen[value]; exists {
			continue
		}
		seen[value] = struct{}{}
		result = append(result, value)
	}
	return result, nil
}

func nonNilAccountingUsage(usage map[string]any) map[string]any {
	if usage == nil {
		return map[string]any{}
	}
	return usage
}

func recordFromHandle(handle protocol.RootHandle) (RootHandleRecord, error) {
	row := RootHandleRecord{
		ID:              strings.TrimSpace(handle.ID),
		OrganizationID:  strings.TrimSpace(handle.OrganizationID),
		WorkspaceID:     strings.TrimSpace(handle.WorkspaceID),
		RequesterID:     strings.TrimSpace(handle.RequesterID),
		TurnID:          strings.TrimSpace(handle.TurnID),
		KnowledgeBaseID: strings.TrimSpace(handle.KnowledgeBaseID),
		SourceID:        strings.TrimSpace(handle.SourceID),
		SourceName:      strings.TrimSpace(handle.SourceName),
		RootKind:        strings.TrimSpace(handle.RootKind),
		RepoRole:        strings.TrimSpace(handle.RepoRole),
		Revision:        strings.TrimSpace(handle.Revision),
		RepositoryRef:   strings.TrimSpace(handle.RepositoryRef),
		DefaultRef:      strings.TrimSpace(handle.DefaultRef),
		ResolvedRef:     strings.TrimSpace(handle.ResolvedRef),
		ResolvedRefKind: strings.TrimSpace(handle.ResolvedRefKind),
		FreshnessStatus: strings.TrimSpace(handle.FreshnessStatus),
		ExpiresAt:       handle.ExpiresAt.UTC(),
		CreatedAt:       handle.CreatedAt.UTC(),
	}
	if row.ID == "" || row.OrganizationID == "" || row.WorkspaceID == "" ||
		row.RequesterID == "" || row.TurnID == "" || row.KnowledgeBaseID == "" ||
		row.RootKind == "" || row.RepoRole == "" || row.Revision == "" ||
		row.RepositoryRef == "" || row.ExpiresAt.IsZero() || row.CreatedAt.IsZero() {
		return RootHandleRecord{}, fmt.Errorf("Knowledge Q&A root handle is incomplete")
	}
	return row, nil
}

func handleFromRecord(row RootHandleRecord) protocol.RootHandle {
	return protocol.RootHandle{
		ID:              row.ID,
		OrganizationID:  row.OrganizationID,
		WorkspaceID:     row.WorkspaceID,
		RequesterID:     row.RequesterID,
		TurnID:          row.TurnID,
		KnowledgeBaseID: row.KnowledgeBaseID,
		SourceID:        row.SourceID,
		SourceName:      row.SourceName,
		RootKind:        row.RootKind,
		RepoRole:        row.RepoRole,
		Revision:        row.Revision,
		RepositoryRef:   row.RepositoryRef,
		DefaultRef:      row.DefaultRef,
		ResolvedRef:     row.ResolvedRef,
		ResolvedRefKind: row.ResolvedRefKind,
		FreshnessStatus: row.FreshnessStatus,
		ExpiresAt:       row.ExpiresAt,
		CreatedAt:       row.CreatedAt,
	}
}