package plan

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

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

type Store struct {
	db coredb.DB
}

func New(db coredb.DB) *Store {
	return &Store{db: db}
}

func (store *Store) CheckOrganizationQuotaInTx(ctx context.Context, tx coredb.Tx, organizationID string, metric MetricKey, delta int64) error {
	trimmedOrganizationID := strings.TrimSpace(organizationID)
	if trimmedOrganizationID == "" {
		return storeerr.BadRequest("organizationId is required")
	}
	if !ValidMetric(metric) {
		return storeerr.BadRequest("metric is invalid")
	}
	if delta <= 0 {
		return storeerr.BadRequest("quota delta must be positive")
	}
	if metric != MetricSeats && metric != MetricAgents {
		return nil
	}
	if err := lockOrganizationForQuotaCheck(ctx, tx, trimmedOrganizationID); err != nil {
		return err
	}
	return store.checkOrganizationSeatEntitlementQuotaInTx(ctx, tx, trimmedOrganizationID, metric, delta)
}

func (store *Store) checkOrganizationSeatEntitlementQuotaInTx(ctx context.Context, tx coredb.Tx, organizationID string, metric MetricKey, delta int64) error {
	now := time.Now().UTC()
	entitlement, err := seatbilling.LoadEffectiveEntitlement(ctx, tx, organizationID, now, true)
	if err != nil {
		return err
	}
	limit := entitlementLimit(entitlement, metric)
	if metric == MetricAgents {
		grants, _, err := loadActiveResourceGrantSummary(ctx, tx, organizationID, metric, QuotaScopeAgentOwner, now)
		if err != nil {
			return err
		}
		limit += grants
	}
	if limit < 0 {
		return nil
	}
	current, err := loadOrganizationMetricUsageForEntitlement(ctx, tx, organizationID, metric, entitlement)
	if err != nil {
		return err
	}
	if current+delta <= limit {
		return nil
	}
	planKey := PlanKey(strings.TrimSpace(entitlement.PlanKey))
	planName := planNameForSeatEntitlement(entitlement)
	return QuotaExceededError{
		OrganizationID: organizationID,
		PlanKey:        planKey,
		PlanName:       planName,
		Metric:         metric,
		Limit:          limit,
		Current:        current,
		Delta:          delta,
		Scope:          quotaScopeForMetric(metric),
	}
}

func (store *Store) CheckWorkspaceAgentOwnerQuotaInTx(ctx context.Context, tx coredb.Tx, workspaceID string, ownerUserID string, delta int64) error {
	trimmedWorkspaceID := strings.TrimSpace(workspaceID)
	if trimmedWorkspaceID == "" {
		return storeerr.BadRequest("workspaceId is required")
	}
	trimmedOwnerUserID := strings.TrimSpace(ownerUserID)
	if trimmedOwnerUserID == "" {
		return storeerr.BadRequest("ownerUserId is required")
	}
	if delta <= 0 {
		return storeerr.BadRequest("quota delta must be positive")
	}
	organizationID, err := workspaceOrganizationID(ctx, tx, trimmedWorkspaceID)
	if err != nil {
		return err
	}
	if err := lockOrganizationForQuotaCheck(ctx, tx, organizationID); err != nil {
		return err
	}
	now := time.Now().UTC()
	entitlement, err := seatbilling.LoadEffectiveEntitlement(ctx, tx, organizationID, now, true)
	if err != nil {
		return err
	}
	limit := entitlementLimit(entitlement, MetricAgents)
	if grants, _, err := loadActiveResourceGrantSummary(ctx, tx, organizationID, MetricAgents, QuotaScopeAgentOwner, now); err != nil {
		return err
	} else {
		limit += grants
	}
	if limit < 0 {
		return nil
	}
	current, err := loadAgentQuotaUsageForOwner(ctx, tx, organizationID, trimmedOwnerUserID)
	if err != nil {
		return err
	}
	if current+delta <= limit {
		return nil
	}
	planKey := PlanKey(strings.TrimSpace(entitlement.PlanKey))
	planName := planNameForSeatEntitlement(entitlement)
	return QuotaExceededError{
		OrganizationID: organizationID,
		PlanKey:        planKey,
		PlanName:       planName,
		Metric:         MetricAgents,
		Limit:          limit,
		Current:        current,
		Delta:          delta,
		Scope:          quotaScopeForMetric(MetricAgents),
	}
}

func quotaScopeForMetric(metric MetricKey) QuotaScope {
	switch metric {
	case MetricAgents:
		return QuotaScopeAgentOwner
	case MetricManagedMemories, MetricSkillRepositories:
		return QuotaScopeWorkspaceMember
	default:
		return QuotaScopeOrganization
	}
}

func entitlementLimit(entitlement seatbilling.SeatEntitlement, metric MetricKey) int64 {
	if metric == MetricAgents {
		return entitlement.OrganizationAgentLimit
	}
	return entitlement.HumanSeatLimit
}

func planNameForSeatEntitlement(entitlement seatbilling.SeatEntitlement) string {
	switch strings.TrimSpace(entitlement.PlanKey) {
	case string(PlanFree):
		return "Free"
	case seatbilling.PlanKeyTeam:
		return "Team"
	case seatbilling.PlanKeyEnterprise:
		return "Enterprise"
	case seatbilling.PlanKeyMax:
		return "Max"
	default:
		return strings.TrimSpace(entitlement.PlanKey)
	}
}

func (store *Store) CheckWorkspaceQuotaInTx(ctx context.Context, tx coredb.Tx, workspaceID string, metric MetricKey, delta int64) error {
	organizationID, err := workspaceOrganizationID(ctx, tx, workspaceID)
	if err != nil {
		return err
	}
	return store.CheckOrganizationQuotaInTx(ctx, tx, organizationID, metric, delta)
}

func (store *Store) LoadOrganizationPlanView(ctx context.Context, organizationID string) (OrganizationPlanView, error) {
	if store == nil || store.db == nil {
		return OrganizationPlanView{}, fmt.Errorf("plan store is not configured")
	}
	trimmedOrganizationID := strings.TrimSpace(organizationID)
	if trimmedOrganizationID == "" {
		return OrganizationPlanView{}, storeerr.BadRequest("organizationId is required")
	}
	return loadOrganizationPlanView(ctx, store.db, trimmedOrganizationID, time.Now().UTC())
}

func (store *Store) LoadWorkspaceUsageView(ctx context.Context, workspaceID string, ownerUserID string) (WorkspaceUsageView, error) {
	if store == nil || store.db == nil {
		return WorkspaceUsageView{}, fmt.Errorf("plan store is not configured")
	}
	trimmedWorkspaceID := strings.TrimSpace(workspaceID)
	if trimmedWorkspaceID == "" {
		return WorkspaceUsageView{}, storeerr.BadRequest("workspaceId is required")
	}
	trimmedOwnerUserID := strings.TrimSpace(ownerUserID)
	if trimmedOwnerUserID == "" {
		return WorkspaceUsageView{}, storeerr.BadRequest("ownerUserId is required")
	}
	organizationID, err := workspaceOrganizationID(ctx, store.db, trimmedWorkspaceID)
	if err != nil {
		return WorkspaceUsageView{}, err
	}
	workspaceUsage, err := loadWorkspaceUsage(ctx, store.db, trimmedWorkspaceID)
	if err != nil {
		return WorkspaceUsageView{}, err
	}
	organizationView, err := loadOrganizationPlanView(ctx, store.db, organizationID, time.Now().UTC())
	if err != nil {
		return WorkspaceUsageView{}, err
	}
	ownerAgentUsage, err := loadOrganizationOwnedAgentUsage(ctx, store.db, organizationID, trimmedOwnerUserID)
	if err != nil {
		return WorkspaceUsageView{}, err
	}
	organizationView.Usage[MetricAgents] = ownerAgentUsage
	return WorkspaceUsageView{
		MachineCount: workspaceUsage.MachineCount,
		AgentCount:   workspaceUsage.AgentCount,
		ChannelCount: workspaceUsage.ChannelCount,
		Organization: organizationView,
	}, nil
}

func (store *Store) LoadWorkspacePlanKey(ctx context.Context, workspaceID string) (PlanKey, error) {
	if store == nil || store.db == nil {
		return "", fmt.Errorf("plan store is not configured")
	}
	organizationID, err := workspaceOrganizationID(ctx, store.db, workspaceID)
	if err != nil {
		return "", err
	}
	entitlement, err := seatbilling.LoadEffectiveEntitlement(ctx, store.db, organizationID, time.Now().UTC(), false)
	if err != nil {
		return "", err
	}
	return PlanKey(strings.TrimSpace(entitlement.PlanKey)), nil
}

func (store *Store) LoadWorkspaceResourceLimit(ctx context.Context, workspaceID string, ownerUserID string, metric MetricKey) (WorkspaceResourceLimit, error) {
	if store == nil || store.db == nil {
		return WorkspaceResourceLimit{}, fmt.Errorf("plan store is not configured")
	}
	trimmedWorkspaceID := strings.TrimSpace(workspaceID)
	if trimmedWorkspaceID == "" {
		return WorkspaceResourceLimit{}, storeerr.BadRequest("workspaceId is required")
	}
	trimmedOwnerUserID := strings.TrimSpace(ownerUserID)
	if trimmedOwnerUserID == "" {
		return WorkspaceResourceLimit{}, storeerr.BadRequest("ownerUserId is required")
	}
	if metric != MetricManagedMemories && metric != MetricSkillRepositories {
		return WorkspaceResourceLimit{}, storeerr.BadRequest("metric is invalid")
	}
	now := time.Now().UTC()
	organizationID, err := workspaceOrganizationID(ctx, store.db, trimmedWorkspaceID)
	if err != nil {
		return WorkspaceResourceLimit{}, err
	}
	if err := requireOrganizationExists(ctx, store.db, organizationID); err != nil {
		return WorkspaceResourceLimit{}, err
	}
	entitlement, err := seatbilling.LoadEffectiveEntitlement(ctx, store.db, organizationID, now, false)
	if err != nil {
		return WorkspaceResourceLimit{}, err
	}
	scope := quotaScopeForMetric(metric)
	return WorkspaceResourceLimit{
		OrganizationID: organizationID,
		PlanKey:        PlanKey(strings.TrimSpace(entitlement.PlanKey)),
		PlanName:       planNameForSeatEntitlement(entitlement),
		Metric:         metric,
		Scope:          scope,
		Limit:          -1,
	}, nil
}

func loadOrganizationPlanView(ctx context.Context, queryer coredb.QueryExecer, organizationID string, now time.Time) (OrganizationPlanView, error) {
	if err := requireOrganizationExists(ctx, queryer, organizationID); err != nil {
		return OrganizationPlanView{}, err
	}
	entitlement, err := seatbilling.LoadEffectiveEntitlement(ctx, queryer, organizationID, now, false)
	if err != nil {
		return OrganizationPlanView{}, err
	}
	limits, err := loadEntitlementLimits(ctx, queryer, organizationID, entitlement, now)
	if err != nil {
		return OrganizationPlanView{}, err
	}
	usage, err := loadOrganizationUsage(ctx, queryer, organizationID, entitlement)
	if err != nil {
		return OrganizationPlanView{}, err
	}
	overLimit := false
	for _, metric := range quotaMetrics {
		limit, ok := limits[metric]
		if !ok {
			return OrganizationPlanView{}, fmt.Errorf("entitlement %s missing required metric %s", entitlement.PlanKey, metric)
		}
		if limit >= 0 && usage[metric] > limit {
			overLimit = true
		}
	}
	return OrganizationPlanView{
		OrganizationID: organizationID,
		PlanKey:        PlanKey(strings.TrimSpace(entitlement.PlanKey)),
		Status:         "active",
		Usage:          usage,
		Limits:         limits,
		OverLimit:      overLimit,
	}, nil
}

func requireOrganizationExists(ctx context.Context, queryer coredb.QueryRower, organizationID string) error {
	row := queryer.QueryRow(ctx, `
SELECT 1
FROM organizations
WHERE id = $1
`, organizationID)
	var marker int
	if err := row.Scan(&marker); err != nil {
		if coredb.IsNotFound(err) {
			return storeerr.NotFound("organization not found")
		}
		return fmt.Errorf("load organization: %w", err)
	}
	return nil
}

func lockOrganizationForQuotaCheck(ctx context.Context, tx coredb.Tx, organizationID string) error {
	row := tx.QueryRow(ctx, `
SELECT id
FROM organizations
WHERE id = $1
FOR UPDATE
`, organizationID)
	var marker string
	if err := row.Scan(&marker); err != nil {
		if coredb.IsNotFound(err) {
			return storeerr.NotFound("organization not found")
		}
		return fmt.Errorf("lock organization for quota check: %w", err)
	}
	return nil
}

func loadEntitlementLimits(ctx context.Context, queryer coredb.QueryRower, organizationID string, entitlement seatbilling.SeatEntitlement, now time.Time) (map[MetricKey]int64, error) {
	limits := map[MetricKey]int64{
		MetricSeats:  entitlement.HumanSeatLimit,
		MetricAgents: entitlement.OrganizationAgentLimit,
	}
	agentGrants, _, err := loadActiveResourceGrantSummary(ctx, queryer, organizationID, MetricAgents, QuotaScopeAgentOwner, now)
	if err != nil {
		return nil, err
	}
	limits[MetricAgents] += agentGrants
	return limits, nil
}

func loadActiveResourceGrantSummary(ctx context.Context, queryer coredb.QueryRower, organizationID string, metric MetricKey, scope QuotaScope, now time.Time) (int64, *time.Time, error) {
	var grants sql.NullInt64
	var nextExpiresAt sql.NullTime
	if err := queryer.QueryRow(ctx, `
SELECT COALESCE(SUM(limit_delta), 0), MIN(expires_at)
FROM organization_resource_quota_grants
WHERE organization_id = $1
  AND metric = $2
  AND quota_scope = $3
  AND revoked_at IS NULL
  AND (expires_at IS NULL OR $4 < expires_at)
`, organizationID, metric, scope, now.UTC()).Scan(&grants, &nextExpiresAt); err != nil {
		return 0, nil, fmt.Errorf("load active resource grants: %w", err)
	}
	var expiresAt *time.Time
	if nextExpiresAt.Valid {
		copy := nextExpiresAt.Time.UTC()
		expiresAt = &copy
	}
	return grants.Int64, expiresAt, nil
}

func workspaceOrganizationID(ctx context.Context, queryer coredb.QueryRower, workspaceID string) (string, error) {
	trimmedWorkspaceID := strings.TrimSpace(workspaceID)
	if trimmedWorkspaceID == "" {
		return "", storeerr.BadRequest("workspaceId is required")
	}
	row := queryer.QueryRow(ctx, `
SELECT organization_id
FROM workspaces
WHERE id = $1
`, trimmedWorkspaceID)
	var organizationID string
	if err := row.Scan(&organizationID); err != nil {
		if coredb.IsNotFound(err) {
			return "", storeerr.ErrNotFound
		}
		return "", fmt.Errorf("load workspace organization: %w", err)
	}
	return organizationID, nil
}