package knowledgeqa

import (
	"context"
	"fmt"
	"strings"
	"time"

	"github.com/tidbcloud/AgentTeam/internal/coredb"
	"github.com/tidbcloud/AgentTeam/internal/store/model"
)

// ListPendingAccountingRecoveries returns a bounded due batch. Its index starts
// with recovery status and next-at, so cost does not grow with unrelated turns.
func (store *Store) ListPendingAccountingRecoveries(ctx context.Context, now time.Time, limit int) ([]AccountingRecovery, error) {
	if store == nil || store.db == nil {
		return nil, fmt.Errorf("knowledge Q&A store is not configured")
	}
	if limit <= 0 {
		limit = 32
	}
	rows, err := store.db.Query(ctx, `
SELECT attempt.workspace_id, attempt.organization_id, turn.requester_user_id,
       attempt.turn_id, attempt.id, attempt.accounting_recovery_attempts,
       attempt.started_at, attempt.accounting_recovery_deadline_at
FROM knowledge_qa_turn_attempts AS attempt
JOIN knowledge_qa_turns AS turn ON turn.id=attempt.turn_id
WHERE attempt.accounting_recovery_status='pending'
  AND attempt.accounting_recovery_next_at IS NOT NULL
  AND attempt.accounting_recovery_next_at <= $1
  AND attempt.accounting_recovery_deadline_at IS NOT NULL
ORDER BY attempt.accounting_recovery_next_at ASC, attempt.id ASC
LIMIT $2
`, now.UTC(), limit)
	if err != nil {
		return nil, fmt.Errorf("list pending Knowledge Q&A accounting recoveries: %w", err)
	}
	defer rows.Close()
	result := make([]AccountingRecovery, 0)
	for rows.Next() {
		var item AccountingRecovery
		if err := rows.Scan(
			&item.WorkspaceID, &item.OrganizationID, &item.RequesterID,
			&item.TurnID, &item.AttemptID, &item.Attempts,
			&item.StartedAt, &item.DeadlineAt,
		); err != nil {
			return nil, fmt.Errorf("scan pending Knowledge Q&A accounting recovery: %w", err)
		}
		result = append(result, item)
	}
	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("iterate pending Knowledge Q&A accounting recoveries: %w", err)
	}
	return result, nil
}

// ResolveAccountingRecovery merges an authoritative full checkpoint into one
// terminal attempt and closes its recovery state without reopening the turn.
func (store *Store) ResolveAccountingRecovery(ctx context.Context, resolution AccountingRecoveryResolution) error {
	if store == nil || store.db == nil {
		return fmt.Errorf("knowledge Q&A store is not configured")
	}
	resolution.WorkspaceID = strings.TrimSpace(resolution.WorkspaceID)
	resolution.TurnID = strings.TrimSpace(resolution.TurnID)
	resolution.AttemptID = strings.TrimSpace(resolution.AttemptID)
	resolution.Now = normalizeNow(resolution.Now)
	if resolution.WorkspaceID == "" || resolution.TurnID == "" || resolution.AttemptID == "" {
		return fmt.Errorf("Knowledge Q&A accounting recovery resolution is incomplete")
	}
	return coredb.WithTx(ctx, store.db, func(tx coredb.Tx) error {
		var status, recoveryStatus, settlementRef string
		var usageJSON, settlementRefsJSON []byte
		if err := tx.QueryRow(ctx, `
SELECT status, accounting_recovery_status, usage_json, settlement_ref, settlement_refs_json
FROM knowledge_qa_turn_attempts
WHERE workspace_id=$1 AND turn_id=$2 AND id=$3
FOR UPDATE
`, resolution.WorkspaceID, resolution.TurnID, resolution.AttemptID).
			Scan(&status, &recoveryStatus, &usageJSON, &settlementRef, &settlementRefsJSON); err != nil {
			if coredb.IsNotFound(err) {
				return ErrNotFound
			}
			return fmt.Errorf("lock Knowledge Q&A accounting recovery: %w", err)
		}
		if recoveryStatus == "resolved" {
			return nil
		}
		if recoveryStatus != "pending" || !isTerminalTurnStatus(status) {
			return ErrInvalidState
		}
		usageJSON, err := model.UnwrapTiDBOpaqueJSONText("Knowledge Q&A existing usage", usageJSON)
		if err != nil {
			return err
		}
		settlementRefsJSON, err = model.UnwrapTiDBOpaqueJSONText(
			"Knowledge Q&A existing settlement references",
			settlementRefsJSON,
		)
		if err != nil {
			return err
		}
		var existingUsage map[string]any
		if err := model.UnmarshalJSONObject("Knowledge Q&A existing usage", usageJSON, &existingUsage); err != nil {
			return err
		}
		var existingRefs []string
		if err := decodeStringList(settlementRefsJSON, &existingRefs); err != nil {
			return err
		}
		if strings.TrimSpace(settlementRef) != "" {
			existingRefs = append(existingRefs, settlementRef)
		}
		mergedUsage := mergeRecoveredUsage(existingUsage, resolution.Usage)
		mergedRefs := normalizeStringSet(append(existingRefs, resolution.SettlementRefs...))
		if recoveredLatest := strings.TrimSpace(resolution.SettlementRef); recoveredLatest != "" {
			mergedRefs = normalizeStringSet(append(mergedRefs, recoveredLatest))
		}
		latest := ""
		if len(mergedRefs) > 0 {
			latest = mergedRefs[len(mergedRefs)-1]
		}
		mergedUsageJSON, err := model.MarshalJSONObject("Knowledge Q&A recovered usage", nonNilMap(mergedUsage))
		if err != nil {
			return err
		}
		mergedRefsJSON, err := model.MarshalJSONObject("Knowledge Q&A recovered settlement references", mergedRefs)
		if err != nil {
			return err
		}
		if err := validatePersistedSettlementReferences(mergedRefs); err != nil {
			return err
		}
		if err := validatePersistedAccountingDocumentSizes(
			mergedUsageJSON, mergedRefsJSON, len(mergedRefs),
		); err != nil {
			return err
		}
		if _, err := tx.Exec(ctx, `
UPDATE knowledge_qa_turn_attempts
SET usage_json=CAST($2 AS JSON), settlement_ref=$3,
    settlement_refs_json=CAST($4 AS JSON),
    accounting_recovery_status='resolved', accounting_recovery_next_at=NULL,
    updated_at=$5
WHERE id=$1 AND accounting_recovery_status='pending'
`, resolution.AttemptID, mergedUsageJSON, latest, mergedRefsJSON, resolution.Now); err != nil {
			return fmt.Errorf("resolve Knowledge Q&A accounting recovery: %w", err)
		}
		return nil
	})
}

// RecordAccountingRecoveryRetry advances bounded retry state or expires a job
// after the MeshX checkpoint recovery window closes.
func (store *Store) RecordAccountingRecoveryRetry(
	ctx context.Context,
	attemptID string,
	nextAttemptAt time.Time,
	expired bool,
) error {
	if store == nil || store.db == nil {
		return fmt.Errorf("knowledge Q&A store is not configured")
	}
	attemptID = strings.TrimSpace(attemptID)
	if attemptID == "" {
		return fmt.Errorf("Knowledge Q&A accounting recovery attempt is required")
	}
	status := "pending"
	var next any = nextAttemptAt.UTC()
	if expired {
		status = "expired"
		next = nil
	}
	if _, err := store.db.Exec(ctx, `
UPDATE knowledge_qa_turn_attempts
SET accounting_recovery_status=$2,
    accounting_recovery_attempts=accounting_recovery_attempts+1,
    accounting_recovery_next_at=$3
WHERE id=$1 AND accounting_recovery_status='pending'
`, attemptID, status, next); err != nil {
		return fmt.Errorf("advance Knowledge Q&A accounting recovery retry: %w", err)
	}
	return nil
}

func markAccountingRecoveryPendingTx(ctx context.Context, tx coredb.Tx, attemptID string, now time.Time) error {
	if _, err := tx.Exec(ctx, `
UPDATE knowledge_qa_turn_attempts
SET accounting_recovery_status='pending',
    accounting_recovery_next_at=COALESCE(accounting_recovery_next_at,$2),
    accounting_recovery_deadline_at=COALESCE(accounting_recovery_deadline_at,$3)
WHERE id=$1 AND accounting_recovery_status IN ('not_required','pending')
`, attemptID, now, now.Add(AccountingRecoveryWindow)); err != nil {
		return fmt.Errorf("mark Knowledge Q&A accounting recovery pending: %w", err)
	}
	return nil
}

func resolveAccountingRecoveryTx(ctx context.Context, tx coredb.Tx, attemptID string, now time.Time) error {
	if _, err := tx.Exec(ctx, `
UPDATE knowledge_qa_turn_attempts
SET accounting_recovery_status='resolved', accounting_recovery_next_at=NULL,
    updated_at=$2
WHERE id=$1 AND accounting_recovery_status='pending'
`, attemptID, now); err != nil {
		return fmt.Errorf("resolve Knowledge Q&A accounting recovery from terminal result: %w", err)
	}
	return nil
}

func mergeRecoveredUsage(existing, recovered map[string]any) map[string]any {
	result := make(map[string]any, len(existing)+len(recovered))
	for key, value := range existing {
		result[key] = value
	}
	for key, recoveredValue := range recovered {
		if recoveredMap, ok := recoveredValue.(map[string]any); ok {
			existingMap, _ := result[key].(map[string]any)
			result[key] = mergeRecoveredUsage(existingMap, recoveredMap)
			continue
		}
		recoveredNumber, recoveredNumeric := recoveredValue.(float64)
		existingNumber, existingNumeric := result[key].(float64)
		if recoveredNumeric && existingNumeric && existingNumber > recoveredNumber {
			continue
		}
		result[key] = recoveredValue
	}
	return result
}