package main

import (
	"context"
	"flag"
	"fmt"
	"net/http"
	"os"
	"os/signal"
	"path/filepath"
	"sync"
	"syscall"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	apiv1 "github.com/tidbcloud/cluster-service-ng/api/cluster-service-ng/v1beta"
	"github.com/tidbcloud/cluster-service-ng/pkg/autoscale/app"
	"github.com/tidbcloud/cluster-service-ng/pkg/autoscale/modules/common"
	"github.com/tidbcloud/cluster-service-ng/pkg/dataplane"
	"github.com/tidbcloud/cluster-service-ng/pkg/dataplanemgr"
	stgdplanemgr "github.com/tidbcloud/cluster-service-ng/pkg/storage/dataplanemanager"
	"gopkg.in/yaml.v2"
	"k8s.io/client-go/rest"
	"k8s.io/client-go/tools/clientcmd"
	"k8s.io/component-base/logs"
	"k8s.io/klog/v2"
)

const discoveryInterval = 60 * time.Second

// engineConfig wraps the autoscale config with engine-specific fields.
type engineConfig struct {
	Autoscale app.AutoscaleConfig `yaml:"autoscale"`
	Engine    engineSpecific      `yaml:"engine"`
}

type engineSpecific struct {
	DataPlaneManagerAddr string `yaml:"dataPlaneManagerAddr"`
	Region               string `yaml:"region"`
	CloudProvider        string `yaml:"cloudProvider"`
	MetricsPort          int    `yaml:"metricsPort"`
}

var autoscaleTiKVActiveScalers = prometheus.NewGaugeVec(
	prometheus.GaugeOpts{
		Name: "autoscale_tikv_active_scalers",
		Help: "Number of currently active TiKV autoscale scalers per shared pool.",
	}, []string{"shared_pool_id"})

func init() {
	prometheus.MustRegister(autoscaleTiKVActiveScalers)
}

func main() {
	var (
		configPath string
		kubeconfig string
	)
	flag.StringVar(&configPath, "config", "", "Path to the configuration file")
	flag.StringVar(&kubeconfig, "kubeconfig", "", "Path to kubeconfig (empty = in-cluster)")
	logs.InitLogs()
	defer logs.FlushLogs()
	flag.Parse()

	if configPath == "" {
		execPath, err := os.Executable()
		if err != nil {
			klog.Fatalf("failed to get executable path: %v", err)
		}
		execDir := filepath.Dir(execPath)
		execName := filepath.Base(execPath)
		configPath = filepath.Join(execDir, execName+".yaml")
	}

	ecfg, err := loadEngineConfig(configPath)
	if err != nil {
		klog.Fatalf("failed to load engine config from %s: %v", configPath, err)
	}

	// Resolve cloudProvider from shared pools at startup so input configs
	// are populated before the app parses them.
	if err := resolveCloudProviderFromPools(ecfg); err != nil {
		klog.Warningf("failed to resolve cloudProvider at startup: %v", err)
	}

	deps, err := buildDependencies(ecfg, kubeconfig)
	if err != nil {
		klog.Fatalf("failed to build dependencies: %v", err)
	}

	metricsPort := ecfg.Engine.MetricsPort
	if metricsPort == 0 {
		metricsPort = 9100
	}
	go func() {
		mux := http.NewServeMux()
		mux.Handle("/metrics", promhttp.Handler())
		klog.Infof("serving metrics on :%d/metrics", metricsPort)
		server := &http.Server{
			Addr:              fmt.Sprintf(":%d", metricsPort),
			Handler:           mux,
			ReadHeaderTimeout: 10 * time.Second,
		}
		if err := server.ListenAndServe(); err != nil {
			klog.Errorf("metrics server error: %v", err)
		}
	}()

	// Create the app and load base config (inputs/outputs/scalers) but NOT clusterScalers.
	// ClusterScalers are managed dynamically by the discovery loop below.
	autoscaleApp, err := app.NewWithDependencies(os.Stderr, deps)
	if err != nil {
		klog.Fatalf("failed to create autoscale app: %v", err)
	}

	scalerRefs, err := autoscaleApp.LoadBaseConfig(&ecfg.Autoscale)
	if err != nil {
		klog.Fatalf("failed to load base config: %v", err)
	}

	scalerUniqName := firstScalerUniqName(&ecfg.Autoscale)
	if scalerUniqName == "" {
		klog.Fatalf("no scalers configured in autoscale config")
	}

	baseURL := findControlPlaneBaseURL(ecfg)
	if baseURL == "" {
		klog.Fatalf("no controlPlaneBaseURL found in input configs")
	}
	cpClient := common.NewControlPlaneClient(baseURL, common.DefaultControlPlaneTimeout)

	pm := &poolManager{
		app:            autoscaleApp,
		scalerRefs:     scalerRefs,
		scalerUniqName: scalerUniqName,
		executorName:   "ScalerInterfaceV2ExecutorSyncLoop",
		cpClient:       cpClient,
		ecfg:           ecfg,
		running:        make(map[string]*scalerHandle),
	}

	// Graceful shutdown.
	ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGTERM, syscall.SIGINT)
	defer stop()

	// Run initial discovery synchronously so the engine starts with scalers.
	pm.discover(ctx)

	// Run periodic discovery loop.
	pm.discoveryLoop(ctx)

	// Context cancelled — stop all running scalers.
	pm.stopAll()
	klog.Info("autoscale engine shut down")
}

// poolManager dynamically manages one scaler goroutine per shared pool.
//
// In the TiKV capacity scaling flow, the framework's ClusterId / TargetClusterId
// holds a **SharedPool ID** (not a tenant cluster ID). This is because TiKV nodes
// are shared across all tenant clusters within a SharedPool — scaling decisions
// are made at the SharedPool level. See ClusterGatewayDataProvider and
// ScalerTiKVCapacityV1 for where sharedPoolID is used as clusterId.
type scalerHandle struct {
	cancel context.CancelFunc
	done   <-chan struct{}
}

type poolManager struct {
	app            *app.App
	scalerRefs     map[string]app.ScalerRefInfo
	scalerUniqName string
	executorName   string
	cpClient       *common.ControlPlaneClient
	ecfg           *engineConfig // retained for re-resolving cloudProvider

	mu      sync.Mutex
	running map[string]*scalerHandle // sharedPoolID → handle
}

func (pm *poolManager) discoveryLoop(ctx context.Context) {
	ticker := time.NewTicker(discoveryInterval)
	defer ticker.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			pm.discover(ctx)
		}
	}
}

// discover lists shared pools and reconciles running scalers.
//
// for K8s client init in stub mode. Engine uses gateway mode so that path is never reached.
// Fixing ParseConfig's signature would require changing 10+ call sites across all modules.
//
//nolint:contextcheck // contextcheck traces into ParseConfig which uses context.Background()
func (pm *poolManager) discover(ctx context.Context) {
	resp, err := pm.cpClient.ListSharedPools(ctx)
	if err != nil {
		klog.Errorf("discovery: ListSharedPools failed: %v", err)
		return
	}

	pools := resp.GetSharedPools()

	// Re-resolve cloudProvider if it was empty at startup (e.g. no pools existed yet).
	if pm.ecfg.Engine.CloudProvider == "" {
		if cp := resolveCloudProvider(pools); cp != "" {
			injectCloudProvider(pm.ecfg, cp)
			klog.Infof("discovery: resolved cloudProvider=%s from pool provider info", cp)
		}
	}

	desired := make(map[string]bool)
	for _, sp := range pools {
		switch sp.GetPhase() {
		case apiv1.SharedPool_AVAILABLE, apiv1.SharedPool_MODIFYING:
			desired[sp.GetId()] = true
		}
	}

	// Lock covers both stop and start phases. RunOneClusterScaler does in-memory
	// config parsing + goroutine launch, so holding the lock here is fine.
	// Revisit if executor creation becomes heavier (e.g. network calls).
	pm.mu.Lock()
	defer pm.mu.Unlock()

	// Stop scalers for removed / deleting pools.
	for id, h := range pm.running {
		if !desired[id] {
			klog.Infof("discovery: stopping scaler for removed pool %s", id)
			h.cancel()
			delete(pm.running, id)
			autoscaleTiKVActiveScalers.DeleteLabelValues(id)
		}
	}

	// Start scalers for newly discovered pools.
	for id := range desired {
		if _, exists := pm.running[id]; exists {
			continue
		}
		cfg := &app.ClusterScalerConfig{
			ScalerUniqName:  pm.scalerUniqName,
			TargetClusterId: id,
			ExecutorName:    pm.executorName,
		}
		cancel, done, err := pm.app.RunOneClusterScaler(cfg, pm.scalerRefs)
		if err != nil {
			klog.Errorf("discovery: failed to start scaler for pool %s: %v", id, err)
			continue
		}
		handle := &scalerHandle{cancel: cancel, done: done}
		pm.running[id] = handle
		autoscaleTiKVActiveScalers.WithLabelValues(id).Set(1)
		klog.Infof("discovery: started scaler for pool %s", id)

		go pm.watchScaler(id, handle)
	}

	klog.Infof("discovery: %d pools active, %d scalers running", len(desired), len(pm.running))
}

func (pm *poolManager) watchScaler(poolID string, watched *scalerHandle) {
	<-watched.done

	pm.mu.Lock()
	defer pm.mu.Unlock()
	if current := pm.running[poolID]; current == watched {
		klog.Warningf("scaler for pool %s exited unexpectedly, removing from running set", poolID)
		delete(pm.running, poolID)
		autoscaleTiKVActiveScalers.DeleteLabelValues(poolID)
	}
}

func (pm *poolManager) stopAll() {
	pm.mu.Lock()
	defer pm.mu.Unlock()
	for id, h := range pm.running {
		klog.Infof("shutdown: stopping scaler for pool %s", id)
		h.cancel()
		autoscaleTiKVActiveScalers.DeleteLabelValues(id)
	}
	pm.running = nil
}

// resolveCloudProviderFromPools does a one-time ListSharedPools to resolve
// cloudProvider from the shared pool's Provider info and inject it into
// input configs that don't already have it set.
func resolveCloudProviderFromPools(ecfg *engineConfig) error {
	baseURL := findControlPlaneBaseURL(ecfg)
	if baseURL == "" {
		return fmt.Errorf("no controlPlaneBaseURL found in input configs")
	}

	cli := common.NewControlPlaneClient(baseURL, common.DefaultControlPlaneTimeout)
	resp, err := cli.ListSharedPools(context.Background())
	if err != nil {
		return fmt.Errorf("ListSharedPools: %w", err)
	}

	cloudProvider := resolveCloudProvider(resp.GetSharedPools())
	if cloudProvider == "" {
		return nil
	}

	injectCloudProvider(ecfg, cloudProvider)
	klog.Infof("resolved cloudProvider=%s from shared pool provider info", cloudProvider)
	return nil
}

func findControlPlaneBaseURL(ecfg *engineConfig) string {
	for _, input := range ecfg.Autoscale.Inputs {
		if input.Config != nil {
			if v, ok := input.Config["controlPlaneBaseURL"]; ok {
				if s, ok := v.(string); ok && s != "" {
					return s
				}
			}
		}
	}
	return ""
}

func firstScalerUniqName(cfg *app.AutoscaleConfig) string {
	if len(cfg.Scalers) > 1 {
		klog.Warningf("multiple scalers configured (%d), only the first will be used", len(cfg.Scalers))
	}
	for _, s := range cfg.Scalers {
		return s.UniqName
	}
	return ""
}

// injectCloudProvider sets Engine.CloudProvider and populates the "cloudProvider"
// key in each input config that doesn't already have it.
func injectCloudProvider(ecfg *engineConfig, cp string) {
	if ecfg.Engine.CloudProvider == "" {
		ecfg.Engine.CloudProvider = cp
	}
	for i := range ecfg.Autoscale.Inputs {
		if ecfg.Autoscale.Inputs[i].Config == nil {
			ecfg.Autoscale.Inputs[i].Config = map[string]any{}
		}
		if _, exists := ecfg.Autoscale.Inputs[i].Config["cloudProvider"]; !exists {
			ecfg.Autoscale.Inputs[i].Config["cloudProvider"] = cp
		}
	}
}

// resolveCloudProvider maps ProviderType to the lowercase string used in autoscale code.
func resolveCloudProvider(pools []*apiv1.SharedPool) string {
	for _, sp := range pools {
		if sp.GetProvider() == nil {
			continue
		}
		switch sp.GetProvider().GetType() {
		case apiv1.ProviderType_AWS:
			return "aws"
		case apiv1.ProviderType_GCP:
			return "gcp"
		case apiv1.ProviderType_AZURE:
			return "azure"
		case apiv1.ProviderType_ALICLOUD:
			return "alicloud"
		}
	}
	return ""
}

func loadEngineConfig(path string) (*engineConfig, error) {
	data, err := os.ReadFile(path)
	if err != nil {
		return nil, fmt.Errorf("read config file: %w", err)
	}
	var cfg engineConfig
	if err := yaml.Unmarshal(data, &cfg); err != nil {
		return nil, fmt.Errorf("unmarshal config: %w", err)
	}
	return &cfg, nil
}

func buildDependencies(cfg *engineConfig, kubeconfigPath string) (app.Dependencies, error) {
	deps := app.Dependencies{}

	const defaultDataPlaneManagerAddr = "dataplane-manager.dataplane-manager:10000"
	if cfg.Engine.DataPlaneManagerAddr == "" {
		cfg.Engine.DataPlaneManagerAddr = defaultDataPlaneManagerAddr
		klog.Infof("engine.dataPlaneManagerAddr not configured, using default: %s", defaultDataPlaneManagerAddr)
	}

	// K8s rest config (in-cluster or from kubeconfig)
	var restCfg *rest.Config
	var err error
	if kubeconfigPath != "" {
		restCfg, err = clientcmd.BuildConfigFromFlags("", kubeconfigPath)
	} else {
		restCfg, err = rest.InClusterConfig()
	}
	if err != nil {
		return deps, fmt.Errorf("build k8s config: %w", err)
	}

	// Dataplane manager gRPC client
	dpmClient, err := dataplanemgr.NewClient(cfg.Engine.DataPlaneManagerAddr)
	if err != nil {
		return deps, fmt.Errorf("create dataplane manager client: %w", err)
	}

	// Dataplane repo (loads dataplane connection info from dataplane-manager)
	dataplaneRepo, err := stgdplanemgr.NewDataplaneRepo(
		cfg.Engine.Region,
		cfg.Engine.CloudProvider,
		dpmClient,
	)
	if err != nil {
		return deps, fmt.Errorf("create dataplane repo: %w", err)
	}

	deps.DataplaneManager = dataplane.NewManager(restCfg, dataplaneRepo)
	return deps, nil
}