package main

import (
	"context"
	"flag"
	"fmt"
	"net/http"
	"os"
	"os/signal"
	"strings"
	"syscall"

	"github.com/prometheus/client_golang/prometheus/promhttp"
	"github.com/tidbcloud/observability/o11ydiagnosis/pkg/keyviz"
	"github.com/tidbcloud/observability/o11ydiagnosis/pkg/topsql"

	gcs "cloud.google.com/go/storage"
	"github.com/Azure/azure-sdk-for-go/sdk/azidentity"
	"github.com/Azure/azure-sdk-for-go/sdk/storage/azblob"
	"github.com/aws/aws-sdk-go-v2/config"
	"github.com/aws/aws-sdk-go-v2/service/s3"
	"github.com/gin-gonic/gin"
	"github.com/pingcap/log"
	"go.uber.org/zap"
	"k8s.io/apimachinery/pkg/runtime"
	ctrl "sigs.k8s.io/controller-runtime"
	"sigs.k8s.io/controller-runtime/pkg/client"

	"net/http/pprof"

	o11yv1 "github.com/tidbcloud/observability/api/regionalserver/resourceworker/v1"
	"github.com/tidbcloud/observability/o11ydiagnosis/pkg/slowquery"
	"github.com/tidbcloud/observability/o11ydiagnosis/pkg/storage"
)

func main() {
	var (
		vendor                 string
		region                 string
		duckDBSwitchTime       int
		duckDBSwitchAllEnabled bool
		duckDBGrayTenantIds    string
	)

	fmt.Printf("os.Args in init(): %v\n", os.Args)
	flag.StringVar(&vendor, "vendor", "local", "vendor")
	flag.StringVar(&region, "region", "us-west-2", "region")
	flag.IntVar(&duckDBSwitchTime, "duckDBSwitchTime", 0, "slow query api switch time in seconds, 0 means no switch")
	flag.BoolVar(&duckDBSwitchAllEnabled, "duckDBSwitchAllEnabled", false, "slow query api switch all tenants to duckdb, default is false")
	flag.StringVar(&duckDBGrayTenantIds, "duckDBGrayTenantIds", "", "slow query api switch gray tenant ids, comma separated, e.g. '111,222'")

	flag.Parse()

	log.Info("starting diagnosis-query",
		zap.String("vendor", vendor),
		zap.String("region", region),
		zap.Int("duckDBSwitchTime", duckDBSwitchTime),
		zap.Bool("duckDBSwitchAllEnabled", duckDBSwitchAllEnabled),
		zap.String("duckDBGrayTenantIds", duckDBGrayTenantIds))

	gin.SetMode(gin.ReleaseMode)

	topsqlOldClusters := map[string]struct{}{}
	topsqlOldClustersEnv := os.Getenv("TOPSQL_OLD_CLUSTERS")
	topsqlOldClustersList := strings.Split(topsqlOldClustersEnv, ",")
	for _, cluster := range topsqlOldClustersList {
		topsqlOldClusters[cluster] = struct{}{}
	}

	store, err := createStore(context.Background(), vendor, region)
	if err != nil {
		log.Panic("failed to create store", zap.Error(err))
	}
	kubecli, err := createKubeClient()
	if err != nil {
		log.Panic("failed to create kubernetes client", zap.Error(err))
	}

	ng := gin.New()
	ng.Use(gin.Logger())
	ng.GET("/metrics", gin.WrapH(promhttp.Handler()))

	// Register pprof routes.
	pprofGroup := ng.Group("/debug/pprof")
	{
		pprofGroup.GET("/", gin.WrapF(pprof.Index))
		pprofGroup.GET("/cmdline", gin.WrapF(pprof.Cmdline))
		pprofGroup.GET("/profile", gin.WrapF(pprof.Profile))
		pprofGroup.POST("/symbol", gin.WrapF(pprof.Symbol))
		pprofGroup.GET("/symbol", gin.WrapF(pprof.Symbol))
		pprofGroup.GET("/trace", gin.WrapF(pprof.Trace))
		pprofGroup.GET("/allocs", gin.WrapH(pprof.Handler("allocs")))
		pprofGroup.GET("/block", gin.WrapH(pprof.Handler("block")))
		pprofGroup.GET("/goroutine", gin.WrapH(pprof.Handler("goroutine")))
		pprofGroup.GET("/heap", gin.WrapH(pprof.Handler("heap")))
		pprofGroup.GET("/mutex", gin.WrapH(pprof.Handler("mutex")))
		pprofGroup.GET("/threadcreate", gin.WrapH(pprof.Handler("threadcreate")))
	}
	apiv1 := ng.Group("/api/v1")

	tenantIds := grayTenantIDStringToMap(duckDBGrayTenantIds)

	duckDBSwitchConfig := slowquery.DuckDBSwitchConfig{
		DuckDBSwitchTime:       duckDBSwitchTime,
		DuckDBSwitchAllEnabled: duckDBSwitchAllEnabled,
		DuckDBGrayTenantIds:    tenantIds,
		Vendor:                 vendor,
		Region:                 region,
	}

	slowquerysvc := slowquery.NewSlowQueryService("/cache", store, kubecli, duckDBSwitchConfig)
	slowquery.InitSlowQueryApi(apiv1, slowquerysvc)

	topsqlsvc := topsql.NewTopSQLService(http.DefaultClient, kubecli, topsqlOldClusters)
	topsql.InitTopSQLApi(apiv1, topsqlsvc)

	keyvizsvc := keyviz.NewService(store)
	keyviz.InitKeyvizApi(apiv1, keyvizsvc)

	ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
	defer stop()
	go func() {
		log.Info("diagnosis-query is serving")
		if err := ng.Run(":8081"); err != nil && err != http.ErrServerClosed {
			log.Panic("start server failed", zap.Error(err))
		}
	}()
	<-ctx.Done()
	stop()
	log.Info("shutting down gracefully, press Ctrl+C again to force")
}

func grayTenantIDStringToMap(s string) map[string]bool {
	result := make(map[string]bool)
	if s == "" {
		return result
	}

	parts := strings.Split(s, ",")
	for _, part := range parts {
		trimmed := strings.TrimSpace(part)
		if trimmed != "" {
			result[trimmed] = true
		}
	}
	return result
}

func createStore(ctx context.Context, vendor, region string) (storage.Storage, error) {
	switch vendor {
	case "aws":
		cfg, err := config.LoadDefaultConfig(ctx, config.WithRegion(region))
		if err != nil {
			return nil, err
		}
		return storage.NewAwsStorage(s3.NewFromConfig(cfg)), nil
	case "gcp":
		client, err := gcs.NewClient(ctx)
		if err != nil {
			return nil, err
		}
		return storage.NewGcpStorage(client), nil
	case "azure":
		cred, err := azidentity.NewDefaultAzureCredential(nil)
		if err != nil {
			return nil, fmt.Errorf("failed to create azure credential: %w", err)
		}
		client, err := azblob.NewClient(fmt.Sprintf("https://o11yprod%s.blob.core.windows.net", region), cred, nil)
		if err != nil {
			return nil, fmt.Errorf("failed to create azure blob client: %w", err)
		}
		return storage.NewAzureStorage(client), nil
	case "local":
		cfg, err := config.LoadDefaultConfig(ctx, config.WithRegion(region))
		if err != nil {
			return nil, err
		}
		return storage.NewAwsStorage(s3.NewFromConfig(cfg)), nil
	default:
		return nil, fmt.Errorf("invalid vendor: %s", vendor)
	}
}

func createKubeClient() (client.Client, error) {
	scheme := runtime.NewScheme()
	if err := o11yv1.AddToScheme(scheme); err != nil {
		return nil, err
	}
	kubecfg := ctrl.GetConfigOrDie()
	return client.New(kubecfg, client.Options{Scheme: scheme})
}