package service

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"strconv"
	"strings"
	"sync"
	"time"

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

var ErrWikiRecoveryBudgetStopped = requestbudget.ErrStopped

// Check before configuration short-circuits so a disabled/deleted source or
// expired deadline still settles durably without waiting for a model call.
func (s *wikiIngestService) checkWikiRecoveryBudgets(ctx context.Context, payload WikiIngestPayload) error {
	if s.checkpointDB == nil || !s.checkpointDB.Migrator().HasTable(&types.WikiRecoveryBudget{}) {
		return nil
	}
	var rows []types.WikiRecoveryBudget
	if err := s.checkpointDB.WithContext(ctx).Where("tenant_id = ? AND knowledge_base_id = ? AND status IN ?", payload.TenantID, payload.KnowledgeBaseID, []string{"active", "stopped_budget"}).Find(&rows).Error; err != nil {
		return err
	}
	repo := apprepo.WikiRecoveryBudgetRepository{DB: s.checkpointDB}
	for _, row := range rows {
		if err := repo.Check(ctx, row); err != nil {
			if archiveErr := repo.StopAndArchive(ctx, row, "guard_rejected", "preflight"); archiveErr != nil {
				return errors.Join(err, archiveErr)
			}
		}
	}
	return nil
}

func (s *wikiIngestService) rejectOrdinaryRecoveryOverlap(ctx context.Context, payload WikiIngestPayload, ops []WikiPendingOp) error {
	if s.checkpointDB == nil || !s.checkpointDB.Migrator().HasTable(&types.WikiRecoveryBudget{}) {
		return nil
	}
	ids := make([]string, 0, len(ops))
	for _, op := range ops {
		ids = append(ids, op.KnowledgeID)
	}
	var count int64
	if err := s.checkpointDB.WithContext(ctx).Model(&types.WikiRecoveryBudget{}).
		Where("tenant_id = ? AND knowledge_base_id = ? AND knowledge_id IN ? AND status IN ?", payload.TenantID, payload.KnowledgeBaseID, ids, []string{"active", "stopped_budget"}).Count(&count).Error; err != nil {
		return err
	}
	if count != 0 {
		return fmt.Errorf("%w: ordinary operation overlaps controlled recovery", ErrWikiRecoveryBudgetStopped)
	}
	return nil
}

type wikiRecoveryGuard struct {
	repo    apprepo.WikiRecoveryBudgetRepository
	row     types.WikiRecoveryBudget
	cancel  context.CancelCauseFunc
	mu      sync.Mutex
	stopped error
}

func (g *wikiRecoveryGuard) Key() string {
	return fmt.Sprintf("%d:%s", g.row.TenantID, g.row.Identity().Key())
}

func (g *wikiRecoveryGuard) fail(err error) error {
	if err != nil {
		g.mu.Lock()
		if g.stopped == nil {
			g.stopped = err
		}
		err = g.stopped
		g.mu.Unlock()
		g.cancel(err)
	}
	return err
}

func (g *wikiRecoveryGuard) Check(ctx context.Context) error {
	g.mu.Lock()
	stopped := g.stopped
	g.mu.Unlock()
	if stopped != nil {
		return stopped
	}
	if !time.Now().Before(g.row.DeadlineAt) {
		cleanup, cancel := wikiIngestCleanupContext(ctx)
		defer cancel()
		return g.fail(g.repo.Check(cleanup, g.row))
	}
	if err := ctx.Err(); err != nil {
		return err
	}
	return g.fail(g.repo.Check(ctx, g.row))
}

func (g *wikiRecoveryGuard) Reserve(ctx context.Context, model, phase string) error {
	if err := g.Check(ctx); err != nil {
		return err
	}
	return g.fail(g.repo.Reserve(ctx, g.row, model, phase))
}

func wikiRecoveryFromContext(ctx context.Context) *wikiRecoveryGuard {
	guard, _ := requestbudget.FromContext(ctx).(*wikiRecoveryGuard)
	return guard
}

func wikiRecoveryIdentityFromContext(ctx context.Context) types.WikiRecoveryIdentity {
	if guard := wikiRecoveryFromContext(ctx); guard != nil {
		return guard.row.Identity()
	}
	return types.WikiRecoveryIdentity{}
}

func (s *wikiIngestService) bindWikiRecovery(ctx context.Context, payload WikiIngestPayload, identity types.WikiRecoveryIdentity, knowledgeID, modelID string) (context.Context, context.CancelFunc, error) {
	if !identity.Present() {
		return ctx, func() {}, nil
	}
	repo := apprepo.WikiRecoveryBudgetRepository{DB: s.checkpointDB}
	row, err := repo.Load(ctx, payload.TenantID, identity)
	if err != nil {
		return ctx, func() {}, err
	}
	if row.KnowledgeBaseID != payload.KnowledgeBaseID || (knowledgeID != "" && row.KnowledgeID != knowledgeID) || row.ModelID != modelID {
		return ctx, func() {}, fmt.Errorf("%w: task scope mismatch", ErrWikiRecoveryBudgetStopped)
	}
	// Deployment must explicitly attest its immutable release and image. A
	// controlled op cannot execute in an unreviewed local or stale API build.
	if row.ReviewedRelease == "" || row.ReviewedAPIImage == "" ||
		os.Getenv("WEKNORA_REVIEWED_RELEASE") != row.ReviewedRelease || os.Getenv("WEKNORA_REVIEWED_API_IMAGE") != row.ReviewedAPIImage {
		return ctx, func() {}, fmt.Errorf("%w: reviewed runtime mismatch", ErrWikiRecoveryBudgetStopped)
	}
	ctx, deadlineCancel := context.WithDeadlineCause(ctx, row.DeadlineAt, fmt.Errorf("%w: deadline", ErrWikiRecoveryBudgetStopped))
	ctx, cancel := context.WithCancelCause(ctx)
	guard := &wikiRecoveryGuard{repo: repo, row: *row, cancel: cancel}
	ctx = requestbudget.WithGuard(ctx, guard)
	if err := guard.Check(ctx); err != nil {
		deadlineCancel()
		cancel(err)
		return ctx, func() {}, err
	}
	return ctx, func() { deadlineCancel(); cancel(nil) }, nil
}

func (s *wikiIngestService) settleWikiRecoveryStop(ctx context.Context, identity types.WikiRecoveryIdentity, tenant uint64, err error) error {
	if !identity.Present() || err == nil || errors.Is(err, errWikiCheckpointYield) {
		return err
	}
	reason := wikiRecoveryFailureReason(ctx, err)
	if !errors.Is(err, ErrWikiRecoveryBudgetStopped) {
		err = fmt.Errorf("%w: recovery phase failed: %w", ErrWikiRecoveryBudgetStopped, err)
	}
	cleanup, cancel := wikiIngestCleanupContext(ctx)
	defer cancel()
	repo := apprepo.WikiRecoveryBudgetRepository{DB: s.checkpointDB}
	row, loadErr := repo.Load(cleanup, tenant, identity)
	if loadErr != nil {
		return errors.Join(err, loadErr)
	}
	phase := row.LastPhase
	if phase == "" {
		phase = requestbudget.Phase(ctx)
	}
	if phase == "" {
		phase = "recovery"
	}
	return errors.Join(err, repo.StopAndArchive(cleanup, *row, reason, phase))
}

// Persist only a bounded category, never provider response bodies or secrets.
// Phase failures still stop the same recovery; they are not budget exhaustion.
func wikiRecoveryFailureReason(ctx context.Context, err error) string {
	if errors.Is(err, ErrWikiRecoveryBudgetStopped) || errors.Is(context.Cause(ctx), ErrWikiRecoveryBudgetStopped) {
		return "guard_rejected"
	}
	var outputErr *wikiOutputError
	if errors.As(err, &outputErr) {
		if outputErr.incomplete {
			return "model_output_incomplete"
		}
		return "model_output_invalid"
	}
	if errors.Is(err, context.DeadlineExceeded) {
		return "phase_timeout"
	}
	if errors.Is(err, context.Canceled) {
		return "phase_cancelled"
	}
	return "phase_failed"
}

func selectWikiFinalizeGroup(rows []*types.TaskPendingOp) ([]*types.TaskPendingOp, types.WikiRecoveryIdentity, error) {
	var first types.WikiRecoveryIdentity
	var selected []*types.TaskPendingOp
	identities := make([]types.WikiRecoveryIdentity, len(rows))
	for i, row := range rows {
		var identity types.WikiRecoveryIdentity
		var keyIdentity types.WikiRecoveryIdentity
		if strings.HasPrefix(row.DedupKey, "recovery:") {
			parts := strings.SplitN(row.DedupKey, ":", 6)
			if len(parts) != 6 {
				return nil, first, fmt.Errorf("%w: invalid finalize dedup identity", ErrWikiRecoveryBudgetStopped)
			}
			keyIdentity.DeadLetterID, _ = strconv.ParseInt(parts[1], 10, 64)
			keyIdentity.Attempt, _ = strconv.Atoi(parts[2])
			keyIdentity.PendingOpID, _ = strconv.ParseInt(parts[3], 10, 64)
			keyIdentity.Version, _ = strconv.Atoi(parts[4])
		}
		if len(row.Payload) > 0 {
			if err := json.Unmarshal(row.Payload, &identity); err != nil {
				return nil, keyIdentity, fmt.Errorf("%w: invalid finalize payload: %v", ErrWikiRecoveryBudgetStopped, err)
			}
		}
		if keyIdentity.Present() || identity.Present() {
			if !keyIdentity.Valid() || keyIdentity != identity {
				if keyIdentity.Valid() {
					identity = keyIdentity
				}
				return nil, identity, fmt.Errorf("%w: finalize payload and dedup identity mismatch", ErrWikiRecoveryBudgetStopped)
			}
			var payload wikiFinalizeRow
			if err := json.Unmarshal(row.Payload, &payload); err != nil {
				return nil, identity, fmt.Errorf("%w: invalid finalize content: %v", ErrWikiRecoveryBudgetStopped, err)
			}
			if (row.Op != wikiFinalizeOpSlug && row.Op != wikiFinalizeOpChange) ||
				(row.Op == wikiFinalizeOpSlug && payload.Slug == "") || (row.Op == wikiFinalizeOpChange && payload.Change == nil) {
				return nil, identity, fmt.Errorf("%w: incomplete finalize content", ErrWikiRecoveryBudgetStopped)
			}
		}
		identities[i] = identity
		// Controlled rows go first so an ordinary index update cannot absorb
		// this recovery's work and leave its own finalize lane unbudgeted.
		if !first.Present() && identity.Present() {
			first = identity
		}
	}
	for i, row := range rows {
		if identities[i] == first {
			selected = append(selected, row)
		}
	}
	if first.Present() && !first.Valid() {
		return nil, first, fmt.Errorf("%w: invalid finalize identity", ErrWikiRecoveryBudgetStopped)
	}
	return selected, first, nil
}

func (s *wikiIngestService) rejectOrdinaryFinalizeRecoveryOverlap(ctx context.Context, payload WikiIngestPayload) error {
	if s.checkpointDB == nil || !s.checkpointDB.Migrator().HasTable(&types.WikiRecoveryBudget{}) {
		return nil
	}
	var count int64
	if err := s.checkpointDB.WithContext(ctx).Model(&types.WikiRecoveryBudget{}).
		Where("tenant_id = ? AND knowledge_base_id = ? AND status = ?", payload.TenantID, payload.KnowledgeBaseID, "active").Count(&count).Error; err != nil {
		return err
	}
	if count > 0 {
		return fmt.Errorf("%w: controlled recovery still owns index finalization", ErrWikiIngestConcurrent)
	}
	return nil
}

func (s *wikiIngestService) enqueueControlledWikiFinalize(ctx context.Context, payload WikiIngestPayload, slugs []string, titles map[string]string, changes []wikiFinalizeChange) error {
	guard := wikiRecoveryFromContext(ctx)
	if guard == nil {
		return fmt.Errorf("%w: missing recovery guard", ErrWikiRecoveryBudgetStopped)
	}
	if err := guard.Check(ctx); err != nil {
		return err
	}
	identity := guard.row.Identity()
	var rows []types.TaskPendingOp
	appendRow := func(op, key string, value wikiFinalizeRow) error {
		value.WikiRecoveryIdentity = identity
		body, err := json.Marshal(value)
		if err != nil {
			return err
		}
		rows = append(rows, types.TaskPendingOp{TenantID: payload.TenantID, TaskType: wikiFinalizeTaskType, Scope: wikiTaskScope, ScopeID: payload.KnowledgeBaseID, Op: op, DedupKey: "recovery:" + identity.Key() + ":" + key, Payload: body, EnqueuedAt: time.Now().UTC()})
		return nil
	}
	for _, slug := range slugs {
		if err := appendRow(wikiFinalizeOpSlug, slug, wikiFinalizeRow{Slug: slug, Title: titles[slug]}); err != nil {
			return err
		}
	}
	for i := range changes {
		if err := appendRow(wikiFinalizeOpChange, fmt.Sprintf("change:%d", i), wikiFinalizeRow{Change: &changes[i]}); err != nil {
			return err
		}
	}
	if len(changes) == 0 {
		return fmt.Errorf("%w: recovery has no completed document to finalize", ErrWikiRecoveryBudgetStopped)
	}
	err := s.checkpointDB.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
		var budget types.WikiRecoveryBudget
		if err := tx.Where("tenant_id = ? AND pending_op_id = ?", payload.TenantID, identity.PendingOpID).Clauses(clause.Locking{Strength: "UPDATE"}).Take(&budget).Error; err != nil {
			return err
		}
		if budget.Identity() != identity || budget.Status != "active" || !time.Now().Before(budget.DeadlineAt) {
			return fmt.Errorf("%w: finalize handoff no longer active", ErrWikiRecoveryBudgetStopped)
		}
		var op types.TaskPendingOp
		if err := tx.Where("tenant_id = ? AND id = ? AND task_type = ? AND scope_id = ?", payload.TenantID, identity.PendingOpID, wikiTaskType, payload.KnowledgeBaseID).Take(&op).Error; err != nil {
			return err
		}
		var saved types.WikiRecoveryIdentity
		if json.Unmarshal(op.Payload, &saved) != nil || saved != identity {
			return fmt.Errorf("%w: handoff op mismatch", ErrWikiRecoveryBudgetStopped)
		}
		if err := tx.Create(&rows).Error; err != nil {
			return err
		}
		return tx.Where("tenant_id = ? AND id = ?", payload.TenantID, identity.PendingOpID).Delete(&types.TaskPendingOp{}).Error
	})
	if err != nil {
		return err
	}
	s.scheduleFinalize(ctx, payload)
	return nil
}