package capacitycontroller

import (
	"context"
	"encoding/json"
	"errors"
	"log/slog"
	"net/http"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/easyai/easyai-ai-gateway/apps/api/internal/config"
	"github.com/easyai/easyai-ai-gateway/apps/api/internal/executionpool"
	"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
	"github.com/jackc/pgx/v5"
)

const (
	controllerReconcileInterval = 10 * time.Second
	controllerLeadershipRetry   = 5 * time.Second
	controllerDeletionCost      = -1000
)

type capacityStore interface {
	TryAcquireCapacityControllerLeadership(context.Context) (store.Leadership, bool, error)
	WorkerQueueRuntime(context.Context) (store.WorkerQueueRuntime, error)
	ListPoolQueueRuntime(context.Context) ([]store.PoolQueueRuntime, error)
	ListWorkerInstanceRuntime(context.Context) ([]store.WorkerInstanceRuntime, error)
	CapacityDatabaseHealth(context.Context) (store.CapacityDatabaseHealth, error)
	MarkWorkerDraining(context.Context, string) error
	ReactivateWorkerInstance(context.Context, string) error
	PublishDesiredCapacity(context.Context, executionpool.DesiredCapacity) error
}

type OrchestratorAdapter interface {
	PoolState(context.Context, string, string) (PoolInfrastructureState, error)
	ScalePool(context.Context, string, string, int) error
	SetInstanceTerminationPriority(context.Context, string, int) error
}

type Status struct {
	Leader       bool                     `json:"leader"`
	LastRunAt    time.Time                `json:"lastRunAt,omitempty"`
	LastError    string                   `json:"lastError,omitempty"`
	Queue        store.WorkerQueueRuntime `json:"queue"`
	Plan         Plan                     `json:"plan"`
	ScaleActions uint64                   `json:"scaleActions"`
}

type Controller struct {
	cfg              config.Config
	store            capacityStore
	orchestrator     OrchestratorAdapter
	pools            []config.ExecutionPoolCapacityConfig
	logger           *slog.Logger
	expectedRevision string
	now              func() time.Time
	highSince        time.Time
	lowSince         time.Time
	statusMu         sync.RWMutex
	status           Status
}

func New(
	cfg config.Config,
	db capacityStore,
	orchestrator OrchestratorAdapter,
	logger *slog.Logger,
) *Controller {
	return &Controller{
		cfg: cfg, store: db, orchestrator: orchestrator, logger: logger,
		expectedRevision: strings.TrimSpace(os.Getenv("AI_GATEWAY_REVISION")),
		now:              time.Now,
		pools:            mustCapacityPools(cfg),
	}
}

func mustCapacityPools(cfg config.Config) []config.ExecutionPoolCapacityConfig {
	pools, _ := cfg.CapacityPools()
	return pools
}

func (controller *Controller) Run(ctx context.Context) {
	for ctx.Err() == nil {
		leadership, acquired, err := controller.store.TryAcquireCapacityControllerLeadership(ctx)
		if err != nil {
			controller.setError(err)
			if !waitContext(ctx, controllerLeadershipRetry) {
				return
			}
			continue
		}
		if !acquired {
			controller.setLeader(false)
			controller.clearError()
			if !waitContext(ctx, controllerLeadershipRetry) {
				return
			}
			continue
		}
		controller.setLeader(true)
		controller.clearError()
		controller.runLeader(ctx, leadership)
		leadership.Release()
		controller.setLeader(false)
	}
}

func (controller *Controller) runLeader(ctx context.Context, leadership store.Leadership) {
	ticker := time.NewTicker(controllerReconcileInterval)
	defer ticker.Stop()
	if err := controller.reconcile(ctx); err != nil {
		controller.logError("initial capacity reconciliation failed", err)
	}
	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			keepAliveCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
			err := leadership.KeepAlive(keepAliveCtx)
			cancel()
			if err != nil {
				controller.logError("capacity controller leadership lost", err)
				return
			}
			if err := controller.reconcile(ctx); err != nil {
				controller.logError("capacity reconciliation failed", err)
			}
		}
	}
}

func (controller *Controller) reconcile(ctx context.Context) error {
	queue, err := controller.store.WorkerQueueRuntime(ctx)
	if err != nil {
		return err
	}
	poolQueues, err := controller.store.ListPoolQueueRuntime(ctx)
	if err != nil {
		return err
	}
	demandByPool := make(map[string]int, len(poolQueues))
	for _, poolQueue := range poolQueues {
		demandByPool[poolQueue.PoolID] = poolQueue.Queued + poolQueue.Running
	}
	instances, err := controller.store.ListWorkerInstanceRuntime(ctx)
	if err != nil {
		return err
	}
	database, err := controller.store.CapacityDatabaseHealth(ctx)
	if err != nil {
		return err
	}
	now := controller.now()
	poolResources := make([]PoolResources, 0, len(controller.pools))
	infrastructureStates := make(map[string]PoolInfrastructureState, len(controller.pools))
	for _, pool := range controller.pools {
		if pool.MaxReplicas == 0 {
			continue
		}
		state, stateErr := controller.orchestrator.PoolState(ctx, pool.ID, pool.AdapterRef)
		if stateErr != nil {
			return stateErr
		}
		infrastructureStates[pool.ID] = state
		poolResources = append(poolResources, PoolResources{
			PoolID: pool.ID, CurrentReplicas: state.CurrentReplicas,
			MinReplicas: pool.MinReplicas, MaxReplicas: pool.MaxReplicas,
			Demand:                   demandByPool[pool.ID],
			AllocatableMemoryBytes:   state.AllocatableMemoryBytes,
			UsedMemoryBytes:          state.UsedMemoryBytes,
			WorkerRequestMemoryBytes: state.WorkerRequestMemoryBytes,
			AllocatableMilliCPU:      state.AllocatableMilliCPU,
			UsedMilliCPU:             state.UsedMilliCPU,
			WorkerRequestMilliCPU:    state.WorkerRequestMilliCPU,
			MemoryPressure:           state.MemoryPressure,
			Nodes:                    nodeResources(state.Nodes),
		})
	}
	currentTotal := 0
	for _, pool := range poolResources {
		currentTotal += pool.CurrentReplicas
	}
	target := controller.cfg.WorkerTargetOutstandingPerReplica
	if target < 1 {
		target = 2 * controller.cfg.AsyncWorkerInstanceHardLimit
	}
	rawDesired := (max(queue.Queued+queue.Running, 0) + target - 1) / target
	if rawDesired > currentTotal {
		if controller.highSince.IsZero() {
			controller.highSince = now
		}
		controller.lowSince = time.Time{}
	} else if rawDesired < currentTotal && queue.Queued == 0 {
		if controller.lowSince.IsZero() {
			controller.lowSince = now
		}
		controller.highSince = time.Time{}
	} else {
		controller.highSince = time.Time{}
		controller.lowSince = time.Time{}
	}
	revisionHealthy := controller.revisionsMatch(instances, infrastructureStates)
	scaleUpEligible := !controller.highSince.IsZero() &&
		now.Sub(controller.highSince) >= time.Duration(controller.cfg.WorkerScaleUpWindowSeconds)*time.Second &&
		revisionHealthy
	scaleDownEligible := !controller.lowSince.IsZero() &&
		now.Sub(controller.lowSince) >= time.Duration(controller.cfg.WorkerScaleDownStabilizationSeconds)*time.Second &&
		revisionHealthy
	plan := CalculatePlan(PlanInput{
		Queued: queue.Queued, Running: queue.Running,
		InstanceSlots:               controller.cfg.AsyncWorkerInstanceHardLimit,
		TargetOutstandingPerReplica: target,
		MemoryTargetPercent:         controller.cfg.NodeMemoryTargetPercent,
		MemoryHardPercent:           controller.cfg.NodeMemoryHardPercent,
		CPUTargetPercent:            controller.cfg.NodeCPUTargetPercent,
		DatabaseConnections:         database.Connections,
		DatabaseConnectionBudget:    controller.cfg.PostgresConnectionBudget,
		NonWorkerConnectionBudget:   controller.cfg.PostgresNonWorkerConnectionBudget,
		WorkerDatabasePoolMax:       controller.cfg.WorkerDatabaseMaxConns,
		SynchronousDatabasePeers:    database.SynchronousPeers,
		ScaleUpEligible:             scaleUpEligible,
		ScaleDownEligible:           scaleDownEligible,
		Pools:                       poolResources,
	})
	if !revisionHealthy && plan.FrozenReason == "" {
		plan.FrozenReason = "release_revision_mismatch"
	}
	for _, poolPlan := range plan.Pools {
		reason := "capacity_plan"
		if plan.FrozenReason != "" {
			reason = "frozen:" + plan.FrozenReason
		}
		if err := controller.store.PublishDesiredCapacity(ctx, executionpool.DesiredCapacity{
			PoolID: poolPlan.PoolID, Desired: poolPlan.DesiredReplicas,
			Reason: reason, ValidUntil: now.Add(2 * controllerReconcileInterval),
		}); err != nil {
			return err
		}
	}
	if controller.cfg.WorkerAutoscalingEnabled {
		for _, poolPlan := range plan.Pools {
			switch {
			case poolPlan.DesiredReplicas > poolPlan.CurrentReplicas:
				if err := controller.orchestrator.ScalePool(ctx, poolPlan.PoolID, controller.adapterRef(poolPlan.PoolID), poolPlan.DesiredReplicas); err != nil {
					return err
				}
				controller.observeScale(poolPlan.PoolID, poolPlan.CurrentReplicas, poolPlan.DesiredReplicas, "scale_up")
			case poolPlan.DesiredReplicas < poolPlan.CurrentReplicas:
				if err := controller.reconcileScaleDown(ctx, now, poolPlan, instances); err != nil {
					return err
				}
			}
		}
	}
	controller.statusMu.Lock()
	controller.status.LastRunAt = now
	controller.status.LastError = ""
	controller.status.Queue = queue
	controller.status.Plan = plan
	controller.statusMu.Unlock()
	return nil
}

func nodeResources(states []KubernetesNodeState) []NodeResources {
	nodes := make([]NodeResources, 0, len(states))
	for _, state := range states {
		nodes = append(nodes, NodeResources{
			NodeName:               state.NodeName,
			CurrentReplicas:        state.CurrentReplicas,
			AllocatableMemoryBytes: state.AllocatableMemoryBytes,
			UsedMemoryBytes:        state.UsedMemoryBytes,
			WorkerUsedMemoryBytes:  state.WorkerUsedMemoryBytes,
			AllocatableMilliCPU:    state.AllocatableMilliCPU,
			UsedMilliCPU:           state.UsedMilliCPU,
			WorkerUsedMilliCPU:     state.WorkerUsedMilliCPU,
			MemoryPressure:         state.MemoryPressure,
		})
	}
	return nodes
}

func (controller *Controller) reconcileScaleDown(
	ctx context.Context,
	now time.Time,
	poolPlan PoolPlan,
	instances []store.WorkerInstanceRuntime,
) error {
	for _, instance := range instances {
		if workerPoolID(instance) != poolPlan.PoolID || instance.Status != "draining" {
			continue
		}
		if instance.RunningTasks == 0 && instance.ActiveLeases == 0 {
			if strings.TrimSpace(instance.OrchestratorInstanceRef) == "" {
				return errors.New("worker orchestrator instance reference is required for scale down")
			}
			if err := controller.orchestrator.SetInstanceTerminationPriority(ctx, instance.OrchestratorInstanceRef, controllerDeletionCost); err != nil {
				return err
			}
			if err := controller.orchestrator.ScalePool(ctx, poolPlan.PoolID, controller.adapterRef(poolPlan.PoolID), poolPlan.CurrentReplicas-1); err != nil {
				return err
			}
			controller.observeScale(poolPlan.PoolID, poolPlan.CurrentReplicas, poolPlan.CurrentReplicas-1, "drained_scale_down")
			return nil
		}
		if instance.DrainingAt != nil &&
			now.Sub(*instance.DrainingAt) >= time.Duration(controller.cfg.WorkerDrainTimeoutSeconds)*time.Second {
			if err := controller.store.ReactivateWorkerInstance(ctx, instance.InstanceID); err != nil && !errors.Is(err, pgx.ErrNoRows) {
				return err
			}
			controller.observeScale(poolPlan.PoolID, poolPlan.CurrentReplicas, poolPlan.CurrentReplicas, "drain_timeout")
		}
		return nil
	}
	var candidate *store.WorkerInstanceRuntime
	for index := range instances {
		instance := &instances[index]
		if workerPoolID(*instance) != poolPlan.PoolID || instance.Status != "active" {
			continue
		}
		if candidate == nil ||
			instance.RunningTasks+instance.ActiveLeases < candidate.RunningTasks+candidate.ActiveLeases {
			candidate = instance
		}
	}
	if candidate == nil {
		return nil
	}
	if err := controller.store.MarkWorkerDraining(ctx, candidate.InstanceID); err != nil && !errors.Is(err, pgx.ErrNoRows) {
		return err
	}
	controller.observeScale(poolPlan.PoolID, poolPlan.CurrentReplicas, poolPlan.CurrentReplicas, "drain_started")
	return nil
}

func workerPoolID(instance store.WorkerInstanceRuntime) string {
	if strings.TrimSpace(instance.PoolID) != "" {
		return instance.PoolID
	}
	return instance.Site
}

func (controller *Controller) adapterRef(poolID string) string {
	for _, pool := range controller.pools {
		if pool.ID == poolID {
			return pool.AdapterRef
		}
	}
	return poolID
}

func (controller *Controller) revisionsMatch(
	instances []store.WorkerInstanceRuntime,
	states map[string]PoolInfrastructureState,
) bool {
	if controller.expectedRevision == "" {
		return true
	}
	for _, state := range states {
		if state.Revision != "" && state.Revision != controller.expectedRevision {
			return false
		}
	}
	for _, instance := range instances {
		if instance.Revision != "" && instance.Revision != controller.expectedRevision {
			return false
		}
	}
	return true
}

func (controller *Controller) Status() Status {
	controller.statusMu.RLock()
	defer controller.statusMu.RUnlock()
	return controller.status
}

func (controller *Controller) Handler() http.Handler {
	mux := http.NewServeMux()
	mux.HandleFunc("GET /healthz", func(w http.ResponseWriter, _ *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"ok":true,"service":"easyai-capacity-controller"}`))
	})
	mux.HandleFunc("GET /readyz", func(w http.ResponseWriter, _ *http.Request) {
		status := controller.Status()
		// Followers are intentionally idle while the database advisory lock is
		// held by the elected leader. They are still ready to take over and must
		// not make a two-replica Deployment permanently fail its rollout.
		if status.LastError != "" {
			http.Error(w, `{"ok":false}`, http.StatusServiceUnavailable)
			return
		}
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"ok":true}`))
	})
	mux.HandleFunc("GET /status", func(w http.ResponseWriter, _ *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		_ = json.NewEncoder(w).Encode(controller.Status())
	})
	return mux
}

func (controller *Controller) setLeader(leader bool) {
	controller.statusMu.Lock()
	controller.status.Leader = leader
	controller.statusMu.Unlock()
}

func (controller *Controller) setError(err error) {
	controller.statusMu.Lock()
	controller.status.LastError = err.Error()
	controller.statusMu.Unlock()
}

func (controller *Controller) clearError() {
	controller.statusMu.Lock()
	controller.status.LastError = ""
	controller.statusMu.Unlock()
}

func (controller *Controller) logError(message string, err error) {
	controller.setError(err)
	if controller.logger != nil {
		controller.logger.Error(message, "error", err)
	}
}

func (controller *Controller) observeScale(site string, from int, to int, reason string) {
	controller.statusMu.Lock()
	controller.status.ScaleActions++
	controller.statusMu.Unlock()
	if controller.logger != nil {
		controller.logger.Info("worker capacity action",
			"site", site,
			"fromReplicas", from,
			"toReplicas", to,
			"reason", reason,
		)
	}
}

func waitContext(ctx context.Context, duration time.Duration) bool {
	timer := time.NewTimer(duration)
	defer timer.Stop()
	select {
	case <-ctx.Done():
		return false
	case <-timer.C:
		return true
	}
}