package plan

import (
	"context"
	"fmt"
	"testing"
	"time"

	"github.com/stretchr/testify/require"

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

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

	ctx := context.Background()
	db := startPlanRollbackDB(t, ctx)
	now := time.Date(2026, 6, 8, 10, 0, 0, 0, time.UTC)
	seedPlanOrganization(t, ctx, db, "org_plan_free_seats", now)
	for index := 0; index < int(seatbilling.FreeHumanSeatLimit); index++ {
		seedPlanHumanMember(t, ctx, db, "org_plan_free_seats", index, now)
	}

	store := New(db)
	err := coredb.WithTx(ctx, db, func(tx coredb.Tx) error {
		return store.CheckOrganizationQuotaInTx(ctx, tx, "org_plan_free_seats", MetricSeats, 1)
	})

	var quotaErr QuotaExceededError
	require.ErrorAs(t, err, &quotaErr)
	require.Equal(t, PlanFree, quotaErr.PlanKey)
	require.Equal(t, "Free", quotaErr.PlanName)
	require.Equal(t, MetricSeats, quotaErr.Metric)
	require.EqualValues(t, seatbilling.FreeHumanSeatLimit, quotaErr.Limit)
	require.EqualValues(t, seatbilling.FreeHumanSeatLimit, quotaErr.Current)
	require.EqualValues(t, 1, quotaErr.Delta)
}

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

	ctx := context.Background()
	db := startPlanRollbackDB(t, ctx)
	now := activePlanFixtureTime()
	seedPlanOrganization(t, ctx, db, "org_plan_timed_grant", now)
	for index := 0; index < int(seatbilling.FreeHumanSeatLimit); index++ {
		seedPlanHumanMember(t, ctx, db, "org_plan_timed_grant", index, now)
	}
	seedPlanSeatGrant(t, ctx, db, "grant_plan_active", "org_plan_timed_grant", 2, now.Add(30*24*time.Hour), nil, now.Add(-time.Hour))
	seedPlanSeatGrant(t, ctx, db, "grant_plan_expired", "org_plan_timed_grant", 9, now.Add(-24*time.Hour), nil, now.Add(-48*time.Hour))
	seedPlanSeatGrant(t, ctx, db, "grant_plan_revoked", "org_plan_timed_grant", 9, now.Add(30*24*time.Hour), ptrPlanTime(now.Add(-30*time.Minute)), now.Add(-time.Hour))

	store := New(db)
	view, err := store.LoadOrganizationPlanView(ctx, "org_plan_timed_grant")
	require.NoError(t, err)
	require.EqualValues(t, seatbilling.FreeHumanSeatLimit+2, view.Limits[MetricSeats])
	require.EqualValues(t, seatbilling.FreeHumanSeatLimit, view.Usage[MetricSeats])

	err = coredb.WithTx(ctx, db, func(tx coredb.Tx) error {
		return store.CheckOrganizationQuotaInTx(ctx, tx, "org_plan_timed_grant", MetricSeats, 2)
	})
	require.NoError(t, err)

	err = coredb.WithTx(ctx, db, func(tx coredb.Tx) error {
		return store.CheckOrganizationQuotaInTx(ctx, tx, "org_plan_timed_grant", MetricSeats, 3)
	})

	var quotaErr QuotaExceededError
	require.ErrorAs(t, err, &quotaErr)
	require.Equal(t, MetricSeats, quotaErr.Metric)
	require.EqualValues(t, seatbilling.FreeHumanSeatLimit+2, quotaErr.Limit)
	require.EqualValues(t, seatbilling.FreeHumanSeatLimit, quotaErr.Current)
	require.EqualValues(t, 3, quotaErr.Delta)
}

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

	ctx := context.Background()
	db := startPlanRollbackDB(t, ctx)
	now := time.Date(2026, 6, 8, 10, 0, 0, 0, time.UTC)
	seedPlanOrganization(t, ctx, db, "org_plan_free_agents", now)
	seedPlanWorkspace(t, ctx, db, "org_plan_free_agents", "wsp_plan_free_agents", now)
	seedPlanHumanMember(t, ctx, db, "org_plan_free_agents", 0, now)
	seedPlanMachine(t, ctx, db, "wsp_plan_free_agents", "mac_plan_free_agents_owner", "usr_org_plan_free_agents", now)
	seedPlanMachine(t, ctx, db, "wsp_plan_free_agents", "mac_plan_free_agents_peer", "usr_org_plan_free_agents_0", now)
	for index := 0; index < int(seatbilling.FreeAgentsPerHumanOwnerLimit); index++ {
		seedPlanAgent(t, ctx, db, "wsp_plan_free_agents", "mac_plan_free_agents_owner", index, now)
	}

	store := New(db)
	peerUsage, err := store.LoadWorkspaceUsageView(ctx, "wsp_plan_free_agents", "usr_org_plan_free_agents_0")
	require.NoError(t, err)
	require.EqualValues(t, 0, peerUsage.Organization.Usage[MetricAgents])

	ownerUsage, err := store.LoadWorkspaceUsageView(ctx, "wsp_plan_free_agents", "usr_org_plan_free_agents")
	require.NoError(t, err)
	require.EqualValues(t, seatbilling.FreeAgentsPerHumanOwnerLimit, ownerUsage.Organization.Usage[MetricAgents])

	err = coredb.WithTx(ctx, db, func(tx coredb.Tx) error {
		return store.CheckWorkspaceAgentOwnerQuotaInTx(ctx, tx, "wsp_plan_free_agents", "usr_org_plan_free_agents_0", 1)
	})
	require.NoError(t, err)

	err = coredb.WithTx(ctx, db, func(tx coredb.Tx) error {
		return store.CheckWorkspaceAgentOwnerQuotaInTx(ctx, tx, "wsp_plan_free_agents", "usr_org_plan_free_agents", 1)
	})

	var quotaErr QuotaExceededError
	require.ErrorAs(t, err, &quotaErr)
	require.Equal(t, PlanFree, quotaErr.PlanKey)
	require.Equal(t, "Free", quotaErr.PlanName)
	require.Equal(t, MetricAgents, quotaErr.Metric)
	require.EqualValues(t, seatbilling.FreeAgentsPerHumanOwnerLimit, quotaErr.Limit)
	require.EqualValues(t, seatbilling.FreeAgentsPerHumanOwnerLimit, quotaErr.Current)
	require.EqualValues(t, 1, quotaErr.Delta)
	require.Equal(t, QuotaScopeAgentOwner, quotaErr.Scope)
}

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

	ctx := context.Background()
	db := startPlanRollbackDB(t, ctx)
	now := activePlanFixtureTime()
	seedPlanOrganization(t, ctx, db, "org_plan_paid_agents", now)
	seedPlanWorkspace(t, ctx, db, "org_plan_paid_agents", "wsp_plan_paid_agents", now)
	seedPlanMachine(t, ctx, db, "wsp_plan_paid_agents", "mac_plan_paid_agents", "usr_org_plan_paid_agents", now)
	seedPlanHumanMember(t, ctx, db, "org_plan_paid_agents", 0, now)
	seedPlanMachine(t, ctx, db, "wsp_plan_paid_agents", "mac_plan_paid_agents_peer", "usr_org_plan_paid_agents_0", now)
	seedPlanSeatSubscription(t, ctx, db, "org_plan_paid_agents", "active", "enterprise", "month", 1, 20, 20, 9900, now.Add(-time.Hour), now.Add(30*24*time.Hour), now)
	for index := 0; index < 20; index++ {
		seedPlanAgent(t, ctx, db, "wsp_plan_paid_agents", "mac_plan_paid_agents", index, now)
	}

	store := New(db)
	err := coredb.WithTx(ctx, db, func(tx coredb.Tx) error {
		return store.CheckWorkspaceAgentOwnerQuotaInTx(ctx, tx, "wsp_plan_paid_agents", "usr_org_plan_paid_agents_0", 1)
	})
	require.NoError(t, err)

	err = coredb.WithTx(ctx, db, func(tx coredb.Tx) error {
		return store.CheckWorkspaceAgentOwnerQuotaInTx(ctx, tx, "wsp_plan_paid_agents", "usr_org_plan_paid_agents", 1)
	})

	var quotaErr QuotaExceededError
	require.ErrorAs(t, err, &quotaErr)
	require.Equal(t, PlanKey("enterprise"), quotaErr.PlanKey)
	require.Equal(t, "Enterprise", quotaErr.PlanName)
	require.Equal(t, MetricAgents, quotaErr.Metric)
	require.EqualValues(t, 20, quotaErr.Limit)
	require.EqualValues(t, 20, quotaErr.Current)
	require.EqualValues(t, 1, quotaErr.Delta)
	require.Equal(t, QuotaScopeAgentOwner, quotaErr.Scope)
}

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

	ctx := context.Background()
	db := startPlanRollbackDB(t, ctx)
	now := activePlanFixtureTime()
	seedPlanOrganization(t, ctx, db, "org_plan_resource_grants", now)
	seedPlanWorkspace(t, ctx, db, "org_plan_resource_grants", "wsp_plan_resource_grants", now)
	seedPlanMachine(t, ctx, db, "wsp_plan_resource_grants", "mac_plan_resource_grants", "usr_org_plan_resource_grants", now)
	for index := 0; index < int(seatbilling.FreeHumanSeatLimit); index++ {
		seedPlanHumanMember(t, ctx, db, "org_plan_resource_grants", index, now)
	}
	for index := 0; index < int(seatbilling.FreeAgentsPerHumanOwnerLimit); index++ {
		seedPlanAgent(t, ctx, db, "wsp_plan_resource_grants", "mac_plan_resource_grants", index, now)
	}
	seedPlanSeatGrant(t, ctx, db, "grant_plan_seats_extra", "org_plan_resource_grants", 2, now.Add(30*24*time.Hour), nil, now.Add(-time.Hour))
	seedPlanResourceGrant(t, ctx, db, "grant_plan_agent_extra", "org_plan_resource_grants", MetricAgents, QuotaScopeAgentOwner, 1, nil, now.Add(-time.Hour))

	store := New(db)
	view, err := store.LoadOrganizationPlanView(ctx, "org_plan_resource_grants")
	require.NoError(t, err)
	require.EqualValues(t, seatbilling.FreeHumanSeatLimit+2, view.Limits[MetricSeats])
	require.EqualValues(t, seatbilling.FreeAgentsPerHumanOwnerLimit+1, view.Limits[MetricAgents])

	err = coredb.WithTx(ctx, db, func(tx coredb.Tx) error {
		return store.CheckOrganizationQuotaInTx(ctx, tx, "org_plan_resource_grants", MetricSeats, 2)
	})
	require.NoError(t, err)

	err = coredb.WithTx(ctx, db, func(tx coredb.Tx) error {
		return store.CheckWorkspaceAgentOwnerQuotaInTx(ctx, tx, "wsp_plan_resource_grants", "usr_org_plan_resource_grants", 1)
	})
	require.NoError(t, err)

	seedPlanHumanMember(t, ctx, db, "org_plan_resource_grants", int(seatbilling.FreeHumanSeatLimit), now)
	seedPlanAgent(t, ctx, db, "wsp_plan_resource_grants", "mac_plan_resource_grants", int(seatbilling.FreeAgentsPerHumanOwnerLimit), now)

	err = coredb.WithTx(ctx, db, func(tx coredb.Tx) error {
		return store.CheckOrganizationQuotaInTx(ctx, tx, "org_plan_resource_grants", MetricSeats, 2)
	})
	var seatsQuotaErr QuotaExceededError
	require.ErrorAs(t, err, &seatsQuotaErr)
	require.Equal(t, MetricSeats, seatsQuotaErr.Metric)
	require.EqualValues(t, seatbilling.FreeHumanSeatLimit+2, seatsQuotaErr.Limit)
	require.EqualValues(t, seatbilling.FreeHumanSeatLimit+1, seatsQuotaErr.Current)

	err = coredb.WithTx(ctx, db, func(tx coredb.Tx) error {
		return store.CheckWorkspaceAgentOwnerQuotaInTx(ctx, tx, "wsp_plan_resource_grants", "usr_org_plan_resource_grants", 1)
	})
	var agentsQuotaErr QuotaExceededError
	require.ErrorAs(t, err, &agentsQuotaErr)
	require.Equal(t, MetricAgents, agentsQuotaErr.Metric)
	require.EqualValues(t, seatbilling.FreeAgentsPerHumanOwnerLimit+1, agentsQuotaErr.Limit)
	require.EqualValues(t, seatbilling.FreeAgentsPerHumanOwnerLimit+1, agentsQuotaErr.Current)
	require.Equal(t, QuotaScopeAgentOwner, agentsQuotaErr.Scope)
}

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

	ctx := context.Background()
	db := startPlanRollbackDB(t, ctx)
	store := New(db)

	for _, metric := range []MetricKey{MetricWorkspaces, MetricMachines, MetricChannels} {
		err := coredb.WithTx(ctx, db, func(tx coredb.Tx) error {
			return store.CheckOrganizationQuotaInTx(ctx, tx, "missing_org_is_ignored_for_legacy_metrics", metric, 1)
		})
		require.NoError(t, err)
	}
}

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

	ctx := context.Background()
	db := startPlanRollbackDB(t, ctx)
	now := time.Date(2026, 6, 8, 10, 0, 0, 0, time.UTC)
	seedPlanOrganization(t, ctx, db, "org_plan_resource_limits", now)
	seedPlanWorkspace(t, ctx, db, "org_plan_resource_limits", "wsp_plan_resource_limits", now)
	seedPlanResourceGrant(t, ctx, db, "grant_plan_memory_extra", "org_plan_resource_limits", MetricManagedMemories, QuotaScopeWorkspaceMember, 3, nil, now.Add(-time.Hour))

	store := New(db)
	limit, err := store.LoadWorkspaceResourceLimit(ctx, "wsp_plan_resource_limits", "usr_org_plan_resource_limits", MetricManagedMemories)
	require.NoError(t, err)
	require.Equal(t, "org_plan_resource_limits", limit.OrganizationID)
	require.Equal(t, PlanFree, limit.PlanKey)
	require.Equal(t, "Free", limit.PlanName)
	require.Equal(t, MetricManagedMemories, limit.Metric)
	require.Equal(t, QuotaScopeWorkspaceMember, limit.Scope)
	require.EqualValues(t, -1, limit.Limit)
}

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

	ctx := context.Background()
	db := startPlanRollbackDB(t, ctx)
	now := activePlanFixtureTime()
	seedPlanOrganization(t, ctx, db, "org_plan_resource_paid", now)
	seedPlanWorkspace(t, ctx, db, "org_plan_resource_paid", "wsp_plan_resource_paid", now)
	seedPlanSeatSubscription(t, ctx, db, "org_plan_resource_paid", "active", "enterprise", "year", 1, 20, 20, 9900, now.Add(-time.Hour), now.Add(30*24*time.Hour), now)

	store := New(db)
	limit, err := store.LoadWorkspaceResourceLimit(ctx, "wsp_plan_resource_paid", "usr_org_plan_resource_paid", MetricSkillRepositories)
	require.NoError(t, err)
	require.Equal(t, PlanKey("enterprise"), limit.PlanKey)
	require.Equal(t, "Enterprise", limit.PlanName)
	require.Equal(t, MetricSkillRepositories, limit.Metric)
	require.Equal(t, QuotaScopeWorkspaceMember, limit.Scope)
	require.EqualValues(t, -1, limit.Limit)
}

func seedPlanOrganization(t *testing.T, ctx context.Context, db coredb.DB, organizationID string, now time.Time) {
	t.Helper()
	userID := "usr_" + organizationID
	mustPlanExec(t, ctx, db, `
INSERT INTO users (id, email, name, display_name, password_hash, is_admin, created_at, updated_at)
VALUES ($1, $2, $3, $4, 'hash', FALSE, $5, $5)
`, userID, organizationID+"@example.com", userID, userID, now)
	mustPlanExec(t, ctx, db, `
INSERT INTO organizations (id, slug, name, created_by_user_id, created_at, updated_at)
VALUES ($1, $2, $3, $4, $5, $5)
`, organizationID, organizationID, organizationID, userID, now)
}

func activePlanFixtureTime() time.Time {
	// The quota store evaluates paid subscriptions and seat grants against the
	// current wall clock. Anchor active-period fixtures to test time so the
	// integration tests do not expire as calendar dates pass.
	return time.Now().UTC().Truncate(time.Second)
}

func seedPlanHumanMember(t *testing.T, ctx context.Context, db coredb.DB, organizationID string, index int, now time.Time) {
	t.Helper()
	userID := fmt.Sprintf("usr_%s_%d", organizationID, index)
	memberID := fmt.Sprintf("orgm_%s_%d", organizationID, index)
	mustPlanExec(t, ctx, db, `
INSERT INTO users (id, email, name, display_name, password_hash, is_admin, created_at, updated_at)
VALUES ($1, $2, $3, $4, 'hash', FALSE, $5, $5)
`, userID, userID+"@example.com", userID, userID, now)
	mustPlanExec(t, ctx, db, `
INSERT INTO organization_members (id, organization_id, user_id, org_role, status, source, joined_at, created_at, updated_at)
VALUES ($1, $2, $3, 'org_member', 'active', 'local', $4, $4, $4)
`, memberID, organizationID, userID, now)
}

func seedPlanSeatGrant(t *testing.T, ctx context.Context, db coredb.DB, grantID string, organizationID string, seatDelta int64, expiresAt time.Time, revokedAt *time.Time, createdAt time.Time) {
	t.Helper()
	mustPlanExec(t, ctx, db, `
INSERT INTO organization_resource_quota_grants (
  id,
  organization_id,
  metric,
  quota_scope,
  limit_delta,
  expires_at,
  source_type,
  source_ref,
  request_id,
  reason,
  created_by_operator,
  revoked_at,
  revoke_reason,
  created_at,
  updated_at
) VALUES ($1, $2, $3, $4, $5, $6, 'operator_manual', '', $7, 'launch gift', $8, $9, '', $10, $10)
`, grantID, organizationID, MetricSeats, QuotaScopeOrganization, seatDelta, expiresAt, grantID+"_request", "usr_"+organizationID, revokedAt, createdAt)
}

func seedPlanResourceGrant(t *testing.T, ctx context.Context, db coredb.DB, grantID string, organizationID string, metric MetricKey, scope QuotaScope, limitDelta int64, expiresAt *time.Time, createdAt time.Time) {
	t.Helper()
	mustPlanExec(t, ctx, db, `
INSERT INTO organization_resource_quota_grants (
  id,
  organization_id,
  metric,
  quota_scope,
  limit_delta,
  expires_at,
  source_type,
  source_ref,
  request_id,
  reason,
  created_by_operator,
  revoke_reason,
  created_at,
  updated_at
) VALUES ($1, $2, $3, $4, $5, $6, 'operator_manual', '', $7, 'test grant', $8, '', $9, $9)
`, grantID, organizationID, metric, scope, limitDelta, expiresAt, grantID+"_request", "usr_"+organizationID, createdAt)
}

func ptrPlanTime(value time.Time) *time.Time {
	return &value
}

func seedPlanWorkspace(t *testing.T, ctx context.Context, db coredb.DB, organizationID string, workspaceID string, now time.Time) {
	t.Helper()
	mustPlanExec(t, ctx, db, `
INSERT INTO workspaces (id, organization_id, slug, name, latest_message_seq, created_at, updated_at)
VALUES ($1, $2, $3, $3, 0, $4, $4)
`, workspaceID, organizationID, workspaceID, now)
}

func seedPlanMachine(t *testing.T, ctx context.Context, db coredb.DB, workspaceID string, machineID string, createdByUserID string, now time.Time) {
	t.Helper()
	mustPlanExec(t, ctx, db, `
INSERT INTO machines (id, workspace_id, created_by_user_id, name, api_key, status, created_at, updated_at)
VALUES ($1, $2, $3, $1, $1, 'active', $4, $4)
`, machineID, workspaceID, createdByUserID, now)
}

func seedPlanAgent(t *testing.T, ctx context.Context, db coredb.DB, workspaceID string, machineID string, index int, now time.Time) {
	t.Helper()
	memberID := fmt.Sprintf("mem_agent_%d", index)
	agentID := fmt.Sprintf("agt_plan_%d", index)
	mustPlanExec(t, ctx, db, `
INSERT INTO members (id, workspace_id, type, display_name, status, created_at, updated_at)
VALUES ($1, $2, 'agent', $1, 'active', $3, $3)
`, memberID, workspaceID, now)
	mustPlanExec(t, ctx, db, `
INSERT INTO agents (id, workspace_id, machine_id, member_id, name, runtime, model, description, auth_token_hash, status, activity_detail, created_at, updated_at)
VALUES ($1, $2, $3, $4, $1, 'codex', 'gpt-5', '', 'hash', 'inactive', '', $5, $5)
`, agentID, workspaceID, machineID, memberID, now)
}

func seedPlanSeatSubscription(t *testing.T, ctx context.Context, db coredb.DB, organizationID string, status string, planKey string, billingPeriod string, seatQuantity int, managedAgentsPerSeat int64, maxHumanSeats int, unitPriceCents int64, startsAt time.Time, endsAt time.Time, now time.Time) {
	t.Helper()
	mustPlanExec(t, ctx, db, `
INSERT INTO organization_seat_subscriptions (
  organization_id,
  status,
  plan_key,
  billing_period,
  seat_quantity,
  managed_agents_per_seat_snapshot,
  max_human_seats_snapshot,
  unit_price_cents_snapshot,
  currency,
  current_period_start,
  current_period_end,
  created_at,
  updated_at
) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, 'CNY', $9, $10, $11, $11)
`, organizationID, status, planKey, billingPeriod, seatQuantity, managedAgentsPerSeat, maxHumanSeats, unitPriceCents, startsAt, endsAt, now)
}

func mustPlanExec(t *testing.T, ctx context.Context, db coredb.DB, query string, args ...any) {
	t.Helper()
	_, err := db.Exec(ctx, query, args...)
	require.NoError(t, err)
}