package workspace

import (
	"context"
	"database/sql"
	"errors"
	"regexp"
	"testing"
	"time"

	"github.com/DATA-DOG/go-sqlmock"
	"github.com/go-sql-driver/mysql"
	"github.com/stretchr/testify/require"

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

func TestEnsureWorkspacePlacementInTxWritesActivePlacement(t *testing.T) {
	t.Parallel()

	now := time.Date(2026, 5, 28, 1, 2, 3, 0, time.UTC)
	tx := &placementFakeTx{}

	err := EnsureWorkspacePlacementInTx(context.Background(), tx, " wsp_1 ", " workspace-default ", now)

	require.NoError(t, err)
	require.Contains(t, tx.query, "INSERT INTO workspace_placements")
	require.Contains(t, tx.query, "ON DUPLICATE KEY UPDATE")
	require.Equal(t, []any{"wsp_1", "workspace-default", PlacementStateActive, now}, tx.args)
}

func TestEnsureWorkspacePlacementInTxValidatesInputs(t *testing.T) {
	t.Parallel()

	err := EnsureWorkspacePlacementInTx(context.Background(), &placementFakeTx{}, " ", "workspace-default", time.Time{})
	require.EqualError(t, err, "workspace placement workspace id is required")

	err = EnsureWorkspacePlacementInTx(context.Background(), &placementFakeTx{}, "wsp_1", " ", time.Time{})
	require.EqualError(t, err, "workspace placement group id is required for workspace wsp_1")

	err = EnsureWorkspacePlacementInTx(context.Background(), nil, "wsp_1", "workspace-default", time.Time{})
	require.EqualError(t, err, "workspace placement tx is not configured for workspace wsp_1")
}

func TestBackfillWorkspacePlacements(t *testing.T) {
	t.Parallel()

	sqlDB, mock, err := sqlmock.New()
	require.NoError(t, err)
	defer sqlDB.Close()

	store := New(coredb.WrapSQLDB(sqlDB), Options{})
	mock.ExpectExec(regexp.QuoteMeta("INSERT INTO workspace_placements")).
		WithArgs("workspace-default", PlacementStateActive).
		WillReturnResult(sqlmock.NewResult(0, 2))

	inserted, err := store.BackfillWorkspacePlacements(context.Background(), " workspace-default ")

	require.NoError(t, err)
	require.Equal(t, int64(2), inserted)
	require.NoError(t, mock.ExpectationsWereMet())
}

func TestBackfillWorkspacePlacementsRequiresGroup(t *testing.T) {
	t.Parallel()

	inserted, err := (&Store{}).BackfillWorkspacePlacements(context.Background(), " ")

	require.Zero(t, inserted)
	require.EqualError(t, err, "workspace group id is required")
}

func TestRepairMissingWorkspacePlacement(t *testing.T) {
	t.Parallel()

	sqlDB, mock, err := sqlmock.New()
	require.NoError(t, err)
	defer sqlDB.Close()

	store := New(coredb.WrapSQLDB(sqlDB), Options{})
	mock.ExpectBegin()
	mock.ExpectQuery("SELECT id[[:space:]]+FROM workspaces").
		WithArgs("wsp_1").
		WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow("wsp_1"))
	mock.ExpectQuery("SELECT state[[:space:]]+FROM workspace_groups").
		WithArgs("workspace-large-a").
		WillReturnRows(sqlmock.NewRows([]string{"state"}).AddRow(WorkspaceGroupStateActive))
	mock.ExpectQuery("SELECT COUNT\\(\\*\\)[[:space:]]+FROM workspace_placements").
		WithArgs("wsp_1").
		WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(0))
	mock.ExpectExec("INSERT INTO workspace_placements").
		WithArgs("wsp_1", "workspace-large-a", PlacementStateActive).
		WillReturnResult(sqlmock.NewResult(0, 1))
	mock.ExpectCommit()

	placement, err := store.RepairMissingWorkspacePlacement(context.Background(), " wsp_1 ", " workspace-large-a ")

	require.NoError(t, err)
	require.Equal(t, WorkspacePlacement{
		WorkspaceID:      "wsp_1",
		WorkspaceGroupID: "workspace-large-a",
		State:            PlacementStateActive,
		PlacementEpoch:   1,
	}, placement)
	require.NoError(t, mock.ExpectationsWereMet())
}

func TestRepairMissingWorkspacePlacementRejectsExistingPlacement(t *testing.T) {
	t.Parallel()

	sqlDB, mock, err := sqlmock.New()
	require.NoError(t, err)
	defer sqlDB.Close()

	store := New(coredb.WrapSQLDB(sqlDB), Options{})
	mock.ExpectBegin()
	mock.ExpectQuery("SELECT id[[:space:]]+FROM workspaces").
		WithArgs("wsp_1").
		WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow("wsp_1"))
	mock.ExpectQuery("SELECT state[[:space:]]+FROM workspace_groups").
		WithArgs("workspace-large-a").
		WillReturnRows(sqlmock.NewRows([]string{"state"}).AddRow(WorkspaceGroupStateActive))
	mock.ExpectQuery("SELECT COUNT\\(\\*\\)[[:space:]]+FROM workspace_placements").
		WithArgs("wsp_1").
		WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(1))
	mock.ExpectRollback()

	_, err = store.RepairMissingWorkspacePlacement(context.Background(), "wsp_1", "workspace-large-a")

	require.ErrorIs(t, err, storeerr.ErrConflict)
	require.EqualError(t, err, "workspace wsp_1 already has a placement")
	require.NoError(t, mock.ExpectationsWereMet())
}

func TestRepairMissingWorkspacePlacementMapsConcurrentInsertToConflict(t *testing.T) {
	t.Parallel()

	sqlDB, mock, err := sqlmock.New()
	require.NoError(t, err)
	defer sqlDB.Close()

	store := New(coredb.WrapSQLDB(sqlDB), Options{})
	mock.ExpectBegin()
	mock.ExpectQuery("SELECT id[[:space:]]+FROM workspaces").
		WithArgs("wsp_1").
		WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow("wsp_1"))
	mock.ExpectQuery("SELECT state[[:space:]]+FROM workspace_groups").
		WithArgs("workspace-large-a").
		WillReturnRows(sqlmock.NewRows([]string{"state"}).AddRow(WorkspaceGroupStateActive))
	mock.ExpectQuery("SELECT COUNT\\(\\*\\)[[:space:]]+FROM workspace_placements").
		WithArgs("wsp_1").
		WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(0))
	mock.ExpectExec("INSERT INTO workspace_placements").
		WithArgs("wsp_1", "workspace-large-a", PlacementStateActive).
		WillReturnError(&mysql.MySQLError{Number: 1062, Message: "duplicate placement"})
	mock.ExpectRollback()

	_, err = store.RepairMissingWorkspacePlacement(context.Background(), "wsp_1", "workspace-large-a")

	require.ErrorIs(t, err, storeerr.ErrConflict)
	require.EqualError(t, err, "workspace wsp_1 already has a placement")
	require.NoError(t, mock.ExpectationsWereMet())
}

func TestRepairMissingWorkspacePlacementRejectsUnknownWorkspace(t *testing.T) {
	t.Parallel()

	sqlDB, mock, err := sqlmock.New()
	require.NoError(t, err)
	defer sqlDB.Close()

	store := New(coredb.WrapSQLDB(sqlDB), Options{})
	mock.ExpectBegin()
	mock.ExpectQuery("SELECT id[[:space:]]+FROM workspaces").
		WithArgs("wsp_missing").
		WillReturnError(sql.ErrNoRows)
	mock.ExpectRollback()

	_, err = store.RepairMissingWorkspacePlacement(context.Background(), "wsp_missing", "workspace-large-a")

	require.ErrorIs(t, err, storeerr.ErrNotFound)
	require.EqualError(t, err, "workspace wsp_missing not found")
	require.NoError(t, mock.ExpectationsWereMet())
}

func TestRepairMissingWorkspacePlacementRejectsUnknownGroup(t *testing.T) {
	t.Parallel()

	sqlDB, mock, err := sqlmock.New()
	require.NoError(t, err)
	defer sqlDB.Close()

	store := New(coredb.WrapSQLDB(sqlDB), Options{})
	mock.ExpectBegin()
	mock.ExpectQuery("SELECT id[[:space:]]+FROM workspaces").
		WithArgs("wsp_1").
		WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow("wsp_1"))
	mock.ExpectQuery("SELECT state[[:space:]]+FROM workspace_groups").
		WithArgs("workspace-missing").
		WillReturnError(sql.ErrNoRows)
	mock.ExpectRollback()

	_, err = store.RepairMissingWorkspacePlacement(context.Background(), "wsp_1", "workspace-missing")

	require.ErrorIs(t, err, storeerr.ErrConflict)
	require.EqualError(t, err, "workspace group workspace-missing does not exist")
	require.NoError(t, mock.ExpectationsWereMet())
}

func TestRepairMissingWorkspacePlacementRejectsNonActiveGroup(t *testing.T) {
	t.Parallel()

	for _, state := range []string{WorkspaceGroupStateDraining, WorkspaceGroupStateDisabled, "paused"} {
		state := state
		t.Run(state, func(t *testing.T) {
			t.Parallel()

			sqlDB, mock, err := sqlmock.New()
			require.NoError(t, err)
			defer sqlDB.Close()

			store := New(coredb.WrapSQLDB(sqlDB), Options{})
			mock.ExpectBegin()
			mock.ExpectQuery("SELECT id[[:space:]]+FROM workspaces").
				WithArgs("wsp_1").
				WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow("wsp_1"))
			mock.ExpectQuery("SELECT state[[:space:]]+FROM workspace_groups").
				WithArgs("workspace-large-a").
				WillReturnRows(sqlmock.NewRows([]string{"state"}).AddRow(state))
			mock.ExpectRollback()

			_, err = store.RepairMissingWorkspacePlacement(context.Background(), "wsp_1", "workspace-large-a")

			require.ErrorIs(t, err, storeerr.ErrConflict)
			require.EqualError(t, err, "workspace group workspace-large-a is not active: state="+state)
			require.NoError(t, mock.ExpectationsWereMet())
		})
	}
}

func TestRepairMissingWorkspacePlacementValidatesInputs(t *testing.T) {
	t.Parallel()

	_, err := (&Store{}).RepairMissingWorkspacePlacement(context.Background(), " ", "workspace-large-a")
	require.ErrorIs(t, err, storeerr.ErrBadRequest)
	require.EqualError(t, err, "workspaceId is required")

	_, err = (&Store{}).RepairMissingWorkspacePlacement(context.Background(), "wsp_1", " ")
	require.ErrorIs(t, err, storeerr.ErrBadRequest)
	require.EqualError(t, err, "workspaceGroupId is required")
}

func TestCountWorkspacesMissingPlacement(t *testing.T) {
	t.Parallel()

	sqlDB, mock, err := sqlmock.New()
	require.NoError(t, err)
	defer sqlDB.Close()

	store := New(coredb.WrapSQLDB(sqlDB), Options{})
	mock.ExpectQuery("SELECT COUNT").WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(3))

	count, err := store.CountWorkspacesMissingPlacement(context.Background())

	require.NoError(t, err)
	require.Equal(t, int64(3), count)
	require.NoError(t, mock.ExpectationsWereMet())
}

func TestVerifyWorkspacePlacements(t *testing.T) {
	t.Parallel()

	sqlDB, mock, err := sqlmock.New()
	require.NoError(t, err)
	defer sqlDB.Close()

	store := New(coredb.WrapSQLDB(sqlDB), Options{})
	mock.ExpectQuery("SELECT").
		WithArgs(
			"workspace-default",
			PlacementStateDraining,
			PlacementStateMoving,
			PlacementStateDisabled,
			PlacementStateActive,
			PlacementStateDraining,
			PlacementStateMoving,
			PlacementStateDisabled,
		).
		WillReturnRows(sqlmock.NewRows([]string{
			"total_workspace_count",
			"placement_count",
			"missing_placement_count",
			"non_default_placement_count",
			"inactive_placement_count",
			"malformed_placement_count",
		}).AddRow(4, 3, 1, 2, 1, 0))

	result, err := store.VerifyWorkspacePlacements(context.Background(), " workspace-default ")

	require.NoError(t, err)
	require.Equal(t, WorkspacePlacementVerification{
		WorkspaceGroupID:         "workspace-default",
		TotalWorkspaceCount:      4,
		PlacementCount:           3,
		MissingPlacementCount:    1,
		NonDefaultPlacementCount: 2,
		InactivePlacementCount:   1,
		MalformedPlacementCount:  0,
	}, result)
	require.NoError(t, mock.ExpectationsWereMet())
}

func TestVerifyWorkspacePlacementsRequiresGroup(t *testing.T) {
	t.Parallel()

	result, err := (&Store{}).VerifyWorkspacePlacements(context.Background(), " ")

	require.Zero(t, result)
	require.EqualError(t, err, "workspace group id is required")
}

func TestIsWorkspaceRoutableToGroup(t *testing.T) {
	t.Parallel()

	sqlDB, mock, err := sqlmock.New()
	require.NoError(t, err)
	defer sqlDB.Close()

	store := New(coredb.WrapSQLDB(sqlDB), Options{})
	mock.ExpectQuery("SELECT 1").
		WithArgs("wsp_1", "workspace-default", PlacementStateActive, PlacementStateDraining, PlacementStateMoving).
		WillReturnRows(sqlmock.NewRows([]string{"exists"}).AddRow(1))
	mock.ExpectQuery("SELECT 1").
		WithArgs("wsp_2", "workspace-default", PlacementStateActive, PlacementStateDraining, PlacementStateMoving).
		WillReturnRows(sqlmock.NewRows([]string{"exists"}))

	owned, err := store.IsWorkspaceRoutableToGroup(context.Background(), " wsp_1 ", " workspace-default ")
	require.NoError(t, err)
	require.True(t, owned)

	owned, err = store.IsWorkspaceRoutableToGroup(context.Background(), "wsp_2", "workspace-default")
	require.NoError(t, err)
	require.False(t, owned)
	require.NoError(t, mock.ExpectationsWereMet())
}

func TestIsWorkspaceRoutableToGroupRequiresInputs(t *testing.T) {
	t.Parallel()

	owned, err := (&Store{}).IsWorkspaceRoutableToGroup(context.Background(), " ", "workspace-default")
	require.False(t, owned)
	require.EqualError(t, err, "workspace id is required")

	owned, err = (&Store{}).IsWorkspaceRoutableToGroup(context.Background(), "wsp_1", " ")
	require.False(t, owned)
	require.EqualError(t, err, "workspace group id is required")
}

func TestListActiveWorkspaceIDsForGroup(t *testing.T) {
	t.Parallel()

	sqlDB, mock, err := sqlmock.New()
	require.NoError(t, err)
	defer sqlDB.Close()

	store := New(coredb.WrapSQLDB(sqlDB), Options{})
	mock.ExpectQuery("SELECT workspace_id").
		WithArgs("workspace-default", PlacementStateActive).
		WillReturnRows(sqlmock.NewRows([]string{"workspace_id"}).AddRow("wsp_1").AddRow(" wsp_2 "))

	workspaceIDs, err := store.ListActiveWorkspaceIDsForGroup(context.Background(), " workspace-default ")

	require.NoError(t, err)
	require.Equal(t, []string{"wsp_1", "wsp_2"}, workspaceIDs)
	require.NoError(t, mock.ExpectationsWereMet())
}

func TestListRoutableWorkspaceIDsForGroup(t *testing.T) {
	t.Parallel()

	sqlDB, mock, err := sqlmock.New()
	require.NoError(t, err)
	defer sqlDB.Close()

	store := New(coredb.WrapSQLDB(sqlDB), Options{})
	mock.ExpectQuery("SELECT workspace_id").
		WithArgs("workspace-default", PlacementStateActive, PlacementStateDraining, PlacementStateMoving).
		WillReturnRows(sqlmock.NewRows([]string{"workspace_id"}).AddRow("wsp_active").AddRow(" wsp_moving "))

	workspaceIDs, err := store.ListRoutableWorkspaceIDsForGroup(context.Background(), " workspace-default ")

	require.NoError(t, err)
	require.Equal(t, []string{"wsp_active", "wsp_moving"}, workspaceIDs)
	require.NoError(t, mock.ExpectationsWereMet())
}

func TestGetActiveWorkspacePlacementMigration(t *testing.T) {
	t.Parallel()

	now := time.Date(2026, 5, 31, 4, 5, 6, 0, time.UTC)
	sqlDB, mock, err := sqlmock.New()
	require.NoError(t, err)
	defer sqlDB.Close()

	store := New(coredb.WrapSQLDB(sqlDB), Options{})
	mock.ExpectQuery("SELECT migration_id").
		WithArgs(
			"wsp_1",
			PlacementMigrationStateRequested,
			PlacementMigrationStatePreflight,
			PlacementMigrationStateDraining,
			PlacementMigrationStateMoving,
			PlacementMigrationStateVerifying,
		).
		WillReturnRows(sqlmock.NewRows([]string{
			"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",
		}).AddRow(
			"wpm_1",
			"wsp_1",
			"workspace-default",
			"workspace-large-a",
			PlacementMigrationStateMoving,
			int64(2),
			int64(3),
			"opsusr_1",
			"move large workspace",
			"idem-1",
			sql.NullString{},
			now,
			now,
			sql.NullTime{},
		))

	migration, err := store.GetActiveWorkspacePlacementMigration(context.Background(), " wsp_1 ")

	require.NoError(t, err)
	require.NotNil(t, migration)
	require.Equal(t, "wpm_1", migration.MigrationID)
	require.Equal(t, "workspace-large-a", migration.TargetGroupID)
	require.Equal(t, PlacementMigrationStateMoving, migration.State)
	require.NoError(t, mock.ExpectationsWereMet())
}

func TestGetActiveWorkspacePlacementMigrationReturnsNilWhenIdle(t *testing.T) {
	t.Parallel()

	sqlDB, mock, err := sqlmock.New()
	require.NoError(t, err)
	defer sqlDB.Close()

	store := New(coredb.WrapSQLDB(sqlDB), Options{})
	mock.ExpectQuery("SELECT migration_id").
		WithArgs(
			"wsp_1",
			PlacementMigrationStateRequested,
			PlacementMigrationStatePreflight,
			PlacementMigrationStateDraining,
			PlacementMigrationStateMoving,
			PlacementMigrationStateVerifying,
		).
		WillReturnError(sql.ErrNoRows)

	migration, err := store.GetActiveWorkspacePlacementMigration(context.Background(), "wsp_1")

	require.NoError(t, err)
	require.Nil(t, migration)
	require.NoError(t, mock.ExpectationsWereMet())
}

func TestGetActiveWorkspacePlacementMigrationRequiresWorkspaceID(t *testing.T) {
	t.Parallel()

	migration, err := (&Store{}).GetActiveWorkspacePlacementMigration(context.Background(), " ")

	require.Nil(t, migration)
	require.EqualError(t, err, "workspaceId is required")
}

type placementFakeTx struct {
	query string
	args  []any
	err   error
}

func (tx *placementFakeTx) QueryRow(context.Context, string, ...any) coredb.Row {
	return bootstrapErrorRow{err: errors.New("unexpected QueryRow")}
}

func (tx *placementFakeTx) Query(context.Context, string, ...any) (coredb.Rows, error) {
	return nil, errors.New("unexpected Query")
}

func (tx *placementFakeTx) Exec(_ context.Context, query string, args ...any) (coredb.Result, error) {
	tx.query = query
	tx.args = append([]any(nil), args...)
	if tx.err != nil {
		return nil, tx.err
	}
	return bootstrapRowsAffectedResult(1), nil
}

func (tx *placementFakeTx) Commit(context.Context) error   { return nil }
func (tx *placementFakeTx) Rollback(context.Context) error { return nil }