package repository

import (
	"context"
	"crypto/sha256"
	"encoding/binary"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"strconv"
	"strings"
	"time"

	"github.com/Tencent/WeKnora/internal/models/requestbudget"
	"github.com/Tencent/WeKnora/internal/types"
	"gorm.io/gorm"
	"gorm.io/gorm/clause"
)

var ErrWikiRecoveryBudgetStopped = requestbudget.ErrStopped

type WikiRecoveryBudgetRepository struct{ DB *gorm.DB }

func recoveryScope(db *gorm.DB, expected types.WikiRecoveryBudget) *gorm.DB {
	return db.Model(&types.WikiRecoveryBudget{}).Where(
		"tenant_id = ? AND dead_letter_id = ? AND attempt_number = ? AND pending_op_id = ? AND knowledge_base_id = ? AND knowledge_id = ? AND model_id = ?",
		expected.TenantID, expected.DeadLetterID, expected.AttemptNumber, expected.PendingOpID,
		expected.KnowledgeBaseID, expected.KnowledgeID, expected.ModelID)
}

func (r WikiRecoveryBudgetRepository) Load(ctx context.Context, tenant uint64, identity types.WikiRecoveryIdentity) (*types.WikiRecoveryBudget, error) {
	if r.DB == nil || tenant == 0 || !identity.Valid() {
		return nil, fmt.Errorf("%w: incomplete recovery identity", ErrWikiRecoveryBudgetStopped)
	}
	var row types.WikiRecoveryBudget
	err := r.DB.WithContext(ctx).Where("tenant_id = ? AND dead_letter_id = ? AND attempt_number = ? AND pending_op_id = ?",
		tenant, identity.DeadLetterID, identity.Attempt, identity.PendingOpID).Take(&row).Error
	if err != nil {
		return nil, fmt.Errorf("%w: load budget: %v", ErrWikiRecoveryBudgetStopped, err)
	}
	return &row, nil
}

func (r WikiRecoveryBudgetRepository) Check(ctx context.Context, expected types.WikiRecoveryBudget) error {
	return r.checkOrReserve(ctx, expected, "", "", false)
}

func (r WikiRecoveryBudgetRepository) Reserve(ctx context.Context, expected types.WikiRecoveryBudget, modelID, phase string) error {
	return r.checkOrReserve(ctx, expected, modelID, phase, true)
}

// Row locking serializes the final slot across workers and processes. A failed
// send never reaches this repository again to refund its reservation.
func (r WikiRecoveryBudgetRepository) checkOrReserve(ctx context.Context, expected types.WikiRecoveryBudget, modelID, phase string, reserve bool) error {
	if r.DB == nil {
		return fmt.Errorf("%w: budget database unavailable", ErrWikiRecoveryBudgetStopped)
	}
	var stopped error
	err := r.DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
		var row types.WikiRecoveryBudget
		if err := recoveryScope(tx, expected).Clauses(clause.Locking{Strength: "UPDATE"}).Take(&row).Error; err != nil {
			return err
		}
		reason := ""
		now := time.Now().UTC()
		switch {
		case row.Status != "active":
			reason = row.StopReason
			if reason == "" {
				reason = row.Status
			}
		case row.MaxProviderRequests <= 0 || row.DeadlineAt.IsZero():
			reason = "invalid_budget"
		case !now.Before(row.DeadlineAt):
			reason = "deadline"
		case reserve && row.ReservedProviderRequests >= row.MaxProviderRequests:
			reason = "request_limit"
		case row.SourceFingerprint != expected.SourceFingerprint || row.ModelFingerprint != expected.ModelFingerprint ||
			row.ReviewedRelease != expected.ReviewedRelease || row.ReviewedAPIImage != expected.ReviewedAPIImage ||
			row.MaxProviderRequests != expected.MaxProviderRequests || !row.DeadlineAt.Equal(expected.DeadlineAt) ||
			row.EmbeddingModelID != expected.EmbeddingModelID:
			reason = "budget_identity_changed"
		case reserve && modelID != row.ModelID && (row.EmbeddingModelID == "" || modelID != row.EmbeddingModelID):
			reason = "model_mismatch"
		}
		if reason == "" {
			if err := validateRecoveryReceipt(tx, row); err != nil {
				reason = "review_receipt_mismatch"
			}
		}
		if reason == "" {
			source, model, err := WikiRecoveryFingerprints(tx, row)
			if err != nil {
				return err
			}
			if source != row.SourceFingerprint {
				reason = "source_changed"
			}
			if model != row.ModelFingerprint {
				reason = "model_configuration_changed"
			}
		}
		if reason != "" {
			stopped = fmt.Errorf("%w: %s", ErrWikiRecoveryBudgetStopped, reason)
			if row.Status == "active" {
				return recoveryScope(tx, expected).Updates(map[string]any{"status": "stopped_budget", "stop_reason": reason, "last_phase": phase, "updated_at": now}).Error
			}
			return nil
		}
		if !reserve {
			return nil
		}
		result := recoveryScope(tx, expected).Where("status = ? AND reserved_provider_requests < max_provider_requests AND deadline_at > ?", "active", now).
			Updates(map[string]any{"reserved_provider_requests": gorm.Expr("reserved_provider_requests + 1"), "last_phase": phase, "updated_at": now})
		if result.Error != nil {
			return result.Error
		}
		if result.RowsAffected != 1 {
			return errors.New("reservation lost its active budget")
		}
		return nil
	})
	if err != nil {
		return fmt.Errorf("%w: budget transaction: %v", ErrWikiRecoveryBudgetStopped, err)
	}
	return stopped
}

// StopAndArchive removes only rows carrying this exact recovery identity. The
// source documents, page drafts, history and validated checkpoints are retained.
func (r WikiRecoveryBudgetRepository) StopAndArchive(ctx context.Context, expected types.WikiRecoveryBudget, reason, phase string) error {
	return r.DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
		var budget types.WikiRecoveryBudget
		if err := recoveryScope(tx, expected).Clauses(clause.Locking{Strength: "UPDATE"}).Take(&budget).Error; err != nil {
			return err
		}
		if budget.Status == "completed" {
			return nil
		}
		var rows []types.TaskPendingOp
		if err := tx.Where("tenant_id = ? AND scope_id = ? AND task_type IN ?", budget.TenantID, budget.KnowledgeBaseID,
			[]string{types.TypeWikiIngest, types.TypeWikiFinalize}).Find(&rows).Error; err != nil {
			return err
		}
		var retained []types.TaskPendingOp
		var ids []int64
		for _, row := range rows {
			var identity types.WikiRecoveryIdentity
			payloadMatches := json.Unmarshal(row.Payload, &identity) == nil && identity == budget.Identity()
			keyMatches := row.TaskType == types.TypeWikiFinalize && strings.HasPrefix(row.DedupKey, "recovery:"+budget.Identity().Key()+":")
			if (payloadMatches || keyMatches) && (row.TaskType != types.TypeWikiIngest || row.ID == budget.PendingOpID) {
				retained = append(retained, row)
				ids = append(ids, row.ID)
			}
		}
		payload, err := json.Marshal(retained)
		if err != nil {
			return err
		}
		var progress []struct {
			Fingerprint string
			Stage       string
			ItemKey     string
			UpdatedAt   time.Time
		}
		if tx.Migrator().HasTable("wiki_ingest_checkpoints") {
			if err := tx.Table("wiki_ingest_checkpoints").Select("fingerprint, stage, item_key, updated_at").
				Where("tenant_id = ? AND knowledge_base_id = ? AND knowledge_id = ?", budget.TenantID, budget.KnowledgeBaseID, budget.KnowledgeID).Find(&progress).Error; err != nil {
				return err
			}
		}
		progressJSON, err := json.Marshal(map[string]any{
			"source_fingerprint":   budget.SourceFingerprint,
			"validation_state":     "retained_metadata_requires_fingerprint_revalidation",
			"retained_checkpoints": progress,
		})
		if err != nil {
			return err
		}
		if budget.StopReason != "" {
			reason = budget.StopReason
		}
		updates := map[string]any{"status": "stopped_budget", "stop_reason": reason, "updated_at": time.Now().UTC()}
		if phase != "" && budget.LastPhase == "" {
			updates["last_phase"] = phase
		}
		if len(retained) > 0 {
			updates["paused_payload"] = string(payload)
			updates["verified_progress"] = string(progressJSON)
		}
		if err := recoveryScope(tx, expected).Updates(updates).Error; err != nil {
			return err
		}
		if len(ids) > 0 {
			return tx.Where("tenant_id = ? AND id IN ?", budget.TenantID, ids).Delete(&types.TaskPendingOp{}).Error
		}
		return nil
	})
}

func (r WikiRecoveryBudgetRepository) Complete(ctx context.Context, expected types.WikiRecoveryBudget, finalizeIDs []int64) (bool, error) {
	completed := false
	err := r.DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
		var row types.WikiRecoveryBudget
		if err := recoveryScope(tx, expected).Clauses(clause.Locking{Strength: "UPDATE"}).Take(&row).Error; err != nil {
			return err
		}
		if row.Status != "active" || !time.Now().Before(row.DeadlineAt) {
			return fmt.Errorf("%w: finalization no longer active", ErrWikiRecoveryBudgetStopped)
		}
		if len(finalizeIDs) == 0 {
			return errors.New("controlled recovery has no finalize receipt")
		}
		if err := validateRecoveryReceipt(tx, row); err != nil {
			return fmt.Errorf("%w: finalization receipt mismatch", ErrWikiRecoveryBudgetStopped)
		}
		var selected []types.TaskPendingOp
		if err := tx.Where("tenant_id = ? AND task_type = ? AND scope_id = ? AND id IN ?", row.TenantID, types.TypeWikiFinalize, row.KnowledgeBaseID, finalizeIDs).Find(&selected).Error; err != nil {
			return err
		}
		if len(selected) != len(finalizeIDs) {
			return fmt.Errorf("%w: finalization row scope mismatch", ErrWikiRecoveryBudgetStopped)
		}
		for _, op := range selected {
			var identity types.WikiRecoveryIdentity
			if json.Unmarshal(op.Payload, &identity) != nil || identity != row.Identity() {
				return fmt.Errorf("%w: finalization row identity mismatch", ErrWikiRecoveryBudgetStopped)
			}
		}
		if err := tx.Where("tenant_id = ? AND task_type = ? AND id IN ?", row.TenantID, types.TypeWikiFinalize, finalizeIDs).Delete(&types.TaskPendingOp{}).Error; err != nil {
			return err
		}
		var remaining []types.TaskPendingOp
		if err := tx.Where("tenant_id = ? AND scope_id = ? AND task_type IN ?", row.TenantID, row.KnowledgeBaseID, []string{types.TypeWikiIngest, types.TypeWikiFinalize}).Find(&remaining).Error; err != nil {
			return err
		}
		for _, op := range remaining {
			var identity types.WikiRecoveryIdentity
			if (json.Unmarshal(op.Payload, &identity) == nil && identity == row.Identity()) ||
				(op.TaskType == types.TypeWikiFinalize && strings.HasPrefix(op.DedupKey, "recovery:"+row.Identity().Key()+":")) {
				return nil
			}
		}
		if err := recoveryScope(tx, expected).Updates(map[string]any{"status": "completed", "last_phase": "finalize", "updated_at": time.Now().UTC()}).Error; err != nil {
			return err
		}
		completed = true
		return nil
	})
	return completed, err
}

// The receipt is the independent, persisted review. Reopening a process must
// not accept an edited budget row as a fresh immutable snapshot.
func validateRecoveryReceipt(db *gorm.DB, row types.WikiRecoveryBudget) error {
	var receipt struct {
		KnowledgeBaseID   string
		KnowledgeID       string
		ModelID           string
		DeployedFixReview string
	}
	if err := db.Table("wiki_dead_letter_recoveries").Select("knowledge_base_id, knowledge_id, model_id, deployed_fix_review").
		Where("tenant_id = ? AND dead_letter_id = ? AND attempt_number = ? AND pending_op_id = ?", row.TenantID, row.DeadLetterID, row.AttemptNumber, row.PendingOpID).Take(&receipt).Error; err != nil {
		return err
	}
	var review struct {
		TenantID            uint64    `json:"tenant_id"`
		DeadLetterID        int64     `json:"dead_letter_id"`
		AttemptNumber       int       `json:"max_attempt"`
		BudgetVersion       int       `json:"budget_version"`
		KnowledgeBaseID     string    `json:"knowledge_base_id"`
		KnowledgeID         string    `json:"knowledge_id"`
		ModelID             string    `json:"model_id"`
		EmbeddingModelID    string    `json:"embedding_model_id"`
		SourceFingerprint   string    `json:"source_fingerprint"`
		ModelFingerprint    string    `json:"model_fingerprint"`
		FixRelease          string    `json:"fix_release"`
		APIImageID          string    `json:"api_image_id"`
		MaxProviderRequests int64     `json:"max_provider_requests"`
		DeadlineAt          time.Time `json:"deadline_at"`
	}
	if err := json.Unmarshal([]byte(receipt.DeployedFixReview), &review); err != nil {
		return err
	}
	if receipt.KnowledgeBaseID != row.KnowledgeBaseID || receipt.KnowledgeID != row.KnowledgeID || receipt.ModelID != row.ModelID ||
		review.TenantID != row.TenantID || review.DeadLetterID != row.DeadLetterID || review.AttemptNumber < row.AttemptNumber || review.BudgetVersion != 1 ||
		review.KnowledgeBaseID != row.KnowledgeBaseID || review.KnowledgeID != row.KnowledgeID || review.ModelID != row.ModelID ||
		review.EmbeddingModelID != row.EmbeddingModelID || review.SourceFingerprint != row.SourceFingerprint || review.ModelFingerprint != row.ModelFingerprint ||
		review.FixRelease != row.ReviewedRelease || review.APIImageID != row.ReviewedAPIImage ||
		review.MaxProviderRequests != row.MaxProviderRequests || !review.DeadlineAt.Equal(row.DeadlineAt) {
		return errors.New("recovery differs from its review receipt")
	}
	return nil
}

// Fingerprints use unsigned 64-bit big-endian UTF-8 byte lengths followed by
// each field's bytes. SQL-side callers can reproduce this without JSON ordering
// or escaping ambiguities. Parameters remain encrypted and are only hashed.
func recoveryFingerprint(fields []string) string {
	h := sha256.New()
	var size [8]byte
	for _, field := range fields {
		binary.BigEndian.PutUint64(size[:], uint64(len(field)))
		h.Write(size[:])
		h.Write([]byte(field))
	}
	return hex.EncodeToString(h.Sum(nil))
}

func WikiRecoveryFingerprints(db *gorm.DB, budget types.WikiRecoveryBudget) (string, string, error) {
	var doc struct{ Title string }
	if err := db.Table("knowledges").Select("title").Where("tenant_id = ? AND id = ? AND knowledge_base_id = ? AND deleted_at IS NULL", budget.TenantID, budget.KnowledgeID, budget.KnowledgeBaseID).Take(&doc).Error; err != nil {
		return "", "", err
	}
	var kb struct {
		SummaryModelID   string
		EmbeddingModelID string
		WikiConfig       string
		IndexingStrategy string
	}
	if err := db.Table("knowledge_bases").Select("summary_model_id, embedding_model_id, wiki_config, indexing_strategy").Where("tenant_id = ? AND id = ? AND deleted_at IS NULL", budget.TenantID, budget.KnowledgeBaseID).Take(&kb).Error; err != nil {
		return "", "", err
	}
	var config struct {
		SynthesisModelID string `json:"synthesis_model_id"`
	}
	if err := json.Unmarshal([]byte(kb.WikiConfig), &config); err != nil {
		return "", "", err
	}
	if config.SynthesisModelID == "" {
		config.SynthesisModelID = kb.SummaryModelID
	}
	var strategy struct {
		WikiEnabled bool `json:"wiki_enabled"`
	}
	if err := json.Unmarshal([]byte(kb.IndexingStrategy), &strategy); err != nil {
		return "", "", err
	}
	if !strategy.WikiEnabled || config.SynthesisModelID != budget.ModelID || kb.EmbeddingModelID != budget.EmbeddingModelID {
		return "", "", errors.New("recovery knowledge base configuration changed")
	}
	var chunks []types.Chunk
	if err := db.Where("tenant_id = ? AND knowledge_id = ? AND knowledge_base_id = ? AND deleted_at IS NULL", budget.TenantID, budget.KnowledgeID, budget.KnowledgeBaseID).Order("id ASC").Find(&chunks).Error; err != nil {
		return "", "", err
	}
	if len(chunks) == 0 {
		return "", "", errors.New("recovery source has no chunks")
	}
	fields := []string{"wiki-recovery-source-v1", strconv.FormatUint(budget.TenantID, 10), budget.KnowledgeBaseID, budget.KnowledgeID, doc.Title}
	for _, c := range chunks {
		fields = append(fields, c.ID, c.Content, strconv.Itoa(c.ContentRevision), strconv.Itoa(c.ChunkIndex), strconv.Itoa(c.StartAt), strconv.Itoa(c.EndAt), c.ChunkType, c.ParentChunkID, c.ImageInfo, strconv.FormatBool(c.IsEnabled))
	}
	modelFields := []string{"wiki-recovery-model-v1", kb.WikiConfig, kb.IndexingStrategy}
	for _, id := range []string{budget.ModelID, budget.EmbeddingModelID} {
		if id == "" {
			continue
		}
		var m struct {
			ID         string
			Name       string
			Type       string
			Source     string
			Status     string
			Parameters string
		}
		if err := db.Table("models").Select("id, name, type, source, status, parameters").Where("tenant_id = ? AND id = ? AND deleted_at IS NULL", budget.TenantID, id).Take(&m).Error; err != nil {
			return "", "", err
		}
		if m.Status != "active" {
			return "", "", errors.New("recovery model is not active")
		}
		modelFields = append(modelFields, m.ID, m.Name, m.Type, m.Source, m.Status, m.Parameters)
	}
	return recoveryFingerprint(fields), recoveryFingerprint(modelFields), nil
}