package run

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

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

// AgentRecoveryState is the authoritative continuous recovery cycle for one
// (workspace, Agent). Legacy workflow_run_effects recovery columns remain a
// per-callback compatibility projection; this state owns the reset boundary,
// shared cap, and epoch fence.
type AgentRecoveryState struct {
	WorkspaceID   string
	AgentID       string
	RecoveryCount int
	RecoveryEpoch int64
	UpdatedAt     time.Time
}

// GetAgentRecoveryState returns the persisted Agent recovery cycle. A state row
// is created lazily by RecoveryTick or ResetAgentRecovery, so an absent row is
// the zero cycle rather than an error.
func (store *Store) GetAgentRecoveryState(ctx context.Context, workspaceID, agentID string) (AgentRecoveryState, error) {
	if err := store.Validate(); err != nil {
		return AgentRecoveryState{}, err
	}
	workspaceID, agentID = strings.TrimSpace(workspaceID), strings.TrimSpace(agentID)
	if workspaceID == "" || agentID == "" {
		return AgentRecoveryState{}, storeerr.BadRequest("workspaceId and agentId are required")
	}
	var state AgentRecoveryState
	err := store.db.QueryRow(ctx, `
SELECT workspace_id, agent_id, recovery_count, recovery_epoch, updated_at
FROM workflow_agent_recovery_states
WHERE workspace_id=$1 AND agent_id=$2`, workspaceID, agentID).Scan(
		&state.WorkspaceID, &state.AgentID, &state.RecoveryCount, &state.RecoveryEpoch, &state.UpdatedAt,
	)
	if coredb.IsNotFound(err) {
		return AgentRecoveryState{WorkspaceID: workspaceID, AgentID: agentID}, nil
	}
	if err != nil {
		return AgentRecoveryState{}, fmt.Errorf("load Workflow Agent recovery state: %w", err)
	}
	state.UpdatedAt = utc(state.UpdatedAt)
	return state, nil
}

// ResetAgentRecovery atomically begins a fresh Agent recovery epoch. The epoch
// advances even when recovery_count is already zero: a delayed RecoveryTick
// discovered in the previous cycle must not be able to increment the new
// cycle after a successful CLI action.
func (store *Store) ResetAgentRecovery(ctx context.Context, workspaceID, agentID string) error {
	if err := store.Validate(); err != nil {
		return err
	}
	workspaceID, agentID = strings.TrimSpace(workspaceID), strings.TrimSpace(agentID)
	if workspaceID == "" || agentID == "" {
		return storeerr.BadRequest("workspaceId and agentId are required")
	}
	now := utc(time.Now())
	store.runRowLockAttempting("reset")
	return coredb.WithTx(ctx, store.db, func(tx coredb.Tx) error {
		state, err := lockAgentRecoveryState(ctx, tx, workspaceID, agentID, now)
		if err != nil {
			return err
		}
		if err := store.runRowLockBarrier("reset"); err != nil {
			return err
		}
		if state.RecoveryEpoch == maxInt64 {
			return fmt.Errorf("Workflow Agent recovery epoch exhausted")
		}
		if _, err := tx.Exec(ctx, `
UPDATE workflow_agent_recovery_states
SET recovery_count=0, recovery_epoch=$3, updated_at=$4
WHERE workspace_id=$1 AND agent_id=$2`, workspaceID, agentID, state.RecoveryEpoch+1, now); err != nil {
			return fmt.Errorf("reset Workflow Agent recovery state: %w", err)
		}
		return nil
	})
}

const maxInt64 = int64(^uint64(0) >> 1)

// lockAgentRecoveryState creates the state row on first use, then locks the
// exact (workspace, Agent) row. INSERT ... ON DUPLICATE KEY UPDATE is both the
// TiDB-compatible lazy migration path and the per-key serialization point for
// concurrent first ticks/resets.
func lockAgentRecoveryState(ctx context.Context, tx coredb.Tx, workspaceID, agentID string, now time.Time) (AgentRecoveryState, error) {
	workspaceID, agentID = strings.TrimSpace(workspaceID), strings.TrimSpace(agentID)
	if workspaceID == "" || agentID == "" {
		return AgentRecoveryState{}, storeerr.BadRequest("workspaceId and agentId are required")
	}
	if _, err := tx.Exec(ctx, `
INSERT INTO workflow_agent_recovery_states (
  workspace_id, agent_id, recovery_count, recovery_epoch, updated_at
) VALUES ($1, $2, 0, 0, $3)
ON DUPLICATE KEY UPDATE workspace_id=VALUES(workspace_id)`, workspaceID, agentID, utc(now)); err != nil {
		return AgentRecoveryState{}, fmt.Errorf("ensure Workflow Agent recovery state: %w", err)
	}
	var state AgentRecoveryState
	if err := tx.QueryRow(ctx, `
SELECT workspace_id, agent_id, recovery_count, recovery_epoch, updated_at
FROM workflow_agent_recovery_states
WHERE workspace_id=$1 AND agent_id=$2
FOR UPDATE`, workspaceID, agentID).Scan(
		&state.WorkspaceID, &state.AgentID, &state.RecoveryCount, &state.RecoveryEpoch, &state.UpdatedAt,
	); err != nil {
		return AgentRecoveryState{}, fmt.Errorf("lock Workflow Agent recovery state: %w", err)
	}
	state.UpdatedAt = utc(state.UpdatedAt)
	return state, nil
}