package run

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

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

const (
	placementActive   = "active"
	placementDraining = "draining"
	placementMoving   = "moving"
)

func placementRoutable(state string) bool {
	switch strings.TrimSpace(state) {
	case placementActive, placementDraining, placementMoving:
		return true
	default:
		return false
	}
}

// ClaimRunnableRuns returns at most Limit independently fenced Run turns. A
// race may yield fewer results; callers simply scan again after the batch.
func (store *Store) ClaimRunnableRuns(ctx context.Context, params ClaimRunnableParams) ([]ClaimedRun, error) {
	if err := store.Validate(); err != nil {
		return nil, err
	}
	params.WorkerID = strings.TrimSpace(params.WorkerID)
	if params.WorkerID == "" || len(params.WorkerID) > 128 || params.LeaseTTL <= 0 {
		return nil, storeerr.BadRequest("a bounded workerId and positive leaseTTL are required")
	}
	limit, err := boundedLimit(params.Limit)
	if err != nil {
		return nil, err
	}
	if params.Now.IsZero() {
		params.Now = time.Now()
	}
	params.Now = utc(params.Now)
	candidates, err := store.listRunnableCandidates(ctx, params.Now, limit)
	if err != nil {
		return nil, err
	}
	claimed := make([]ClaimedRun, 0, len(candidates))
	for _, candidate := range candidates {
		turn, ok, err := store.claimRunnableCandidate(ctx, candidate, params.WorkerID, params.Now, params.LeaseTTL)
		if err != nil {
			return nil, err
		}
		if ok {
			claimed = append(claimed, turn)
		}
	}
	return claimed, nil
}

func (store *Store) listRunnableCandidates(ctx context.Context, now time.Time, limit int) ([]RunnableCandidate, error) {
	rows, err := store.db.Query(ctx, `
SELECT run.id, run.workspace_id,
       placement.workspace_group_id, placement.state, placement.placement_epoch
FROM workflow_run_instances run
JOIN workspace_placements placement ON placement.workspace_id=run.workspace_id
WHERE run.status IN ('queued','running','waiting')
  AND run.runnable_at <= $1
  AND (run.lease_token IS NULL OR run.lease_expires_at <= $1)
  AND placement.workspace_group_id=COALESCE(NULLIF($2,''),placement.workspace_group_id)
  AND placement.state IN ('active','draining','moving')
ORDER BY run.runnable_at ASC, run.id ASC
LIMIT $3`, now, store.workspaceGroupID, limit)
	if err != nil {
		return nil, fmt.Errorf("list runnable workflow Runs: %w", err)
	}
	defer rows.Close()
	values := make([]RunnableCandidate, 0, limit)
	for rows.Next() {
		var value RunnableCandidate
		if err := rows.Scan(&value.runID, &value.workspaceID, &value.workspaceGroup, &value.placementState, &value.placementEpoch); err != nil {
			return nil, fmt.Errorf("scan runnable workflow Run: %w", err)
		}
		values = append(values, value)
	}
	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("iterate runnable workflow Runs: %w", err)
	}
	return values, nil
}

func (store *Store) claimRunnableCandidate(ctx context.Context, candidate RunnableCandidate, workerID string, now time.Time, ttl time.Duration) (ClaimedRun, bool, error) {
	if candidate.runID == "" || candidate.workspaceID == "" || candidate.workspaceGroup == "" || !placementRoutable(candidate.placementState) || candidate.placementEpoch <= 0 {
		return ClaimedRun{}, false, storeerr.BadRequest("workflow Run candidate was not issued by this store")
	}
	leaseToken, err := identifiers.NewEntityID("wrl")
	if err != nil {
		return ClaimedRun{}, false, fmt.Errorf("generate workflow Run turn lease: %w", err)
	}
	leaseExpires := utc(now.Add(ttl))
	var claimed ClaimedRun
	var ok bool
	err = coredb.WithTx(ctx, store.db, func(tx coredb.Tx) error {
		current, err := store.lockCandidatePlacement(ctx, tx, candidate.workspaceID)
		if err != nil {
			return err
		}
		if current.workspaceGroup != candidate.workspaceGroup || current.state != candidate.placementState || current.epoch != candidate.placementEpoch ||
			(store.workspaceGroupID != "" && store.workspaceGroupID != current.workspaceGroup) {
			return nil
		}
		result, err := tx.Exec(ctx, `
UPDATE workflow_run_instances
SET status='running',
    started_at=COALESCE(started_at,$4),
    lease_owner=$2, lease_token=$3, lease_expires_at=$5,
    revision=revision+1, updated_at=$4
WHERE workspace_id=$1 AND id=$6
  AND status IN ('queued','running','waiting')
  AND runnable_at <= $4
  AND (lease_token IS NULL OR lease_expires_at <= $4)`, candidate.workspaceID, workerID, leaseToken, now, leaseExpires, candidate.runID)
		if err != nil {
			return fmt.Errorf("claim workflow Run turn: %w", err)
		}
		if result.RowsAffected() == 0 {
			return nil
		}
		run, err := scanRun(tx.QueryRow(ctx, runSelectSQL()+` WHERE r.workspace_id=$1 AND r.id=$2`, candidate.workspaceID, candidate.runID))
		if err != nil {
			return fmt.Errorf("load claimed workflow Run: %w", err)
		}
		claimed = ClaimedRun{Run: run, Fence: TurnFence{WorkspaceID: run.WorkspaceID, RunID: run.ID, LeaseToken: leaseToken}}
		ok = true
		return nil
	})
	return claimed, ok, err
}

type placementFence struct {
	workspaceGroup string
	state          string
	epoch          int64
}

func (store *Store) lockCandidatePlacement(ctx context.Context, tx coredb.Tx, workspaceID string) (placementFence, error) {
	var value placementFence
	if err := tx.QueryRow(ctx, `
SELECT workspace_group_id,state,placement_epoch
FROM workspace_placements WHERE workspace_id=$1 FOR UPDATE`, workspaceID).Scan(&value.workspaceGroup, &value.state, &value.epoch); err != nil {
		if coredb.IsNotFound(err) {
			return placementFence{}, fmt.Errorf("workspace placement for workflow Run %s is missing: %w", workspaceID, err)
		}
		return placementFence{}, fmt.Errorf("lock workflow Run workspace placement: %w", err)
	}
	if !placementRoutable(value.state) {
		return placementFence{}, storeerr.Conflict("workflow Run workspace placement is not routable")
	}
	return value, nil
}

// LoadRunForUpdate locks and validates a live turn lease. The placement lock
// must already have been acquired when a caller also performs placement work.
func (store *TxStore) LoadRunForUpdate(ctx context.Context, fence TurnFence, now time.Time) (Run, error) {
	if err := store.Validate(); err != nil {
		return Run{}, err
	}
	fence.WorkspaceID, fence.RunID, fence.LeaseToken = strings.TrimSpace(fence.WorkspaceID), strings.TrimSpace(fence.RunID), strings.TrimSpace(fence.LeaseToken)
	if fence.WorkspaceID == "" || fence.RunID == "" || fence.LeaseToken == "" {
		return Run{}, storeerr.BadRequest("complete workflow Run turn fence is required")
	}
	if now.IsZero() {
		now = time.Now()
	}
	now = utc(now)
	value, err := scanRun(store.tx.QueryRow(ctx, runSelectSQL()+`
WHERE r.workspace_id=$1 AND r.id=$2 FOR UPDATE`, fence.WorkspaceID, fence.RunID))
	if err != nil {
		return Run{}, mapReadError("lock workflow Run turn", err)
	}
	if value.Terminal() || value.LeaseToken != fence.LeaseToken || value.LeaseExpires == nil || !value.LeaseExpires.After(now) {
		return Run{}, storeerr.Conflict("workflow Run turn lease changed or expired")
	}
	return value, nil
}

func (store *Store) WithClaimedRun(ctx context.Context, fence TurnFence, now time.Time, fn func(*TxStore, Run) error) error {
	if err := store.Validate(); err != nil {
		return err
	}
	if fn == nil {
		return storeerr.BadRequest("workflow Run transaction callback is required")
	}
	return coredb.WithTx(ctx, store.db, func(tx coredb.Tx) error {
		txStore := store.InTx(tx)
		run, err := txStore.LoadRunForUpdate(ctx, fence, now)
		if err != nil {
			return err
		}
		return fn(txStore, run)
	})
}

func (store *Store) ReleaseRunTurn(ctx context.Context, fence TurnFence, runnableAt, now time.Time) error {
	if err := store.Validate(); err != nil {
		return err
	}
	if runnableAt.IsZero() {
		return storeerr.BadRequest("workflow Run runnableAt is required")
	}
	if now.IsZero() {
		now = time.Now()
	}
	runnableAt, now = utc(runnableAt), utc(now)
	return coredb.WithTx(ctx, store.db, func(tx coredb.Tx) error {
		txStore := store.InTx(tx)
		run, err := txStore.LoadRunForUpdate(ctx, fence, now)
		if err != nil {
			return err
		}
		result, err := tx.Exec(ctx, `
UPDATE workflow_run_instances
SET runnable_at=$4, lease_owner=NULL, lease_token=NULL, lease_expires_at=NULL,
    revision=revision+1, updated_at=$5
WHERE workspace_id=$1 AND id=$2 AND lease_token=$3 AND revision=$6`, fence.WorkspaceID, fence.RunID, fence.LeaseToken, runnableAt, now, run.Revision)
		if err != nil {
			return fmt.Errorf("release workflow Run turn: %w", err)
		}
		if result.RowsAffected() == 0 {
			return storeerr.Conflict("workflow Run turn changed")
		}
		return nil
	})
}