package workspace

import (
	"context"
	"database/sql"
	"fmt"
	"strings"

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

func loadWorkspacePlacementMigrationForUpdate(ctx context.Context, tx coredb.Tx, migrationID string) (WorkspacePlacementMigration, error) {
	return scanWorkspacePlacementMigration(tx.QueryRow(ctx, `
SELECT migration_id, workspace_id, source_group_id, target_group_id, state, placement_epoch_start, placement_epoch_current, requested_by, reason, idempotency_key, last_error, created_at, updated_at, completed_at
FROM workspace_placement_migrations
WHERE migration_id = $1
FOR UPDATE
`, migrationID), migrationID)
}

func loadWorkspacePlacementMigration(ctx context.Context, queryer coredb.QueryRower, migrationID string) (WorkspacePlacementMigration, error) {
	return scanWorkspacePlacementMigration(queryer.QueryRow(ctx, `
SELECT migration_id, workspace_id, source_group_id, target_group_id, state, placement_epoch_start, placement_epoch_current, requested_by, reason, idempotency_key, last_error, created_at, updated_at, completed_at
FROM workspace_placement_migrations
WHERE migration_id = $1
`, migrationID), migrationID)
}

func loadActiveWorkspacePlacementMigration(ctx context.Context, queryer coredb.QueryRower, workspaceID string) (WorkspacePlacementMigration, error) {
	return scanWorkspacePlacementMigration(queryer.QueryRow(ctx, `
SELECT migration_id, workspace_id, source_group_id, target_group_id, state, placement_epoch_start, placement_epoch_current, requested_by, reason, idempotency_key, last_error, created_at, updated_at, completed_at
FROM workspace_placement_migrations
WHERE workspace_id = $1
  AND state IN ($2, $3, $4, $5, $6)
ORDER BY updated_at DESC, created_at DESC
LIMIT 1
`, workspaceID, PlacementMigrationStateRequested, PlacementMigrationStatePreflight, PlacementMigrationStateDraining, PlacementMigrationStateMoving, PlacementMigrationStateVerifying), "active workspace placement migration for workspace "+workspaceID)
}

func loadWorkspacePlacementMigrationByIdempotency(ctx context.Context, queryer coredb.QueryRower, idempotencyKey string) (WorkspacePlacementMigration, error) {
	return scanWorkspacePlacementMigration(queryer.QueryRow(ctx, `
SELECT migration_id, workspace_id, source_group_id, target_group_id, state, placement_epoch_start, placement_epoch_current, requested_by, reason, idempotency_key, last_error, created_at, updated_at, completed_at
FROM workspace_placement_migrations
WHERE idempotency_key = $1
`, idempotencyKey), idempotencyKey)
}

func scanWorkspacePlacementMigration(row coredb.Row, migrationID string) (WorkspacePlacementMigration, error) {
	var migration WorkspacePlacementMigration
	var lastError sql.NullString
	var completedAt sql.NullTime
	if err := row.Scan(
		&migration.MigrationID,
		&migration.WorkspaceID,
		&migration.SourceGroupID,
		&migration.TargetGroupID,
		&migration.State,
		&migration.PlacementEpochStart,
		&migration.PlacementEpochCurrent,
		&migration.RequestedBy,
		&migration.Reason,
		&migration.IdempotencyKey,
		&lastError,
		&migration.CreatedAt,
		&migration.UpdatedAt,
		&completedAt,
	); err != nil {
		if coredb.IsNotFound(err) {
			return WorkspacePlacementMigration{}, storeerr.NotFound(fmt.Sprintf("workspace placement migration %s not found", strings.TrimSpace(migrationID)))
		}
		return WorkspacePlacementMigration{}, fmt.Errorf("scan workspace placement migration %s: %w", strings.TrimSpace(migrationID), err)
	}
	migration.LastError = strings.TrimSpace(lastError.String)
	if completedAt.Valid {
		completed := completedAt.Time
		migration.CompletedAt = &completed
	}
	return migration, nil
}