package main

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

	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/aws"
	"github.com/aws/aws-sdk-go-v2/config"
	awsconfig "github.com/aws/aws-sdk-go-v2/config"
	"github.com/aws/aws-sdk-go-v2/service/s3"
	"github.com/gin-contrib/cors"
	"github.com/gin-gonic/gin"
	"github.com/pingcap/log"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	swaggerFiles "github.com/swaggo/files"
	ginSwagger "github.com/swaggo/gin-swagger"
	"github.com/tidbcloud/observability/o11ydiagnosis/pkg/apimetrics"
	"github.com/tidbcloud/observability/o11ydiagnosis/pkg/conprof"
	"github.com/tidbcloud/observability/o11ydiagnosis/pkg/keyviz"
	"github.com/tidbcloud/observability/o11ydiagnosis/pkg/query"
	"github.com/tidbcloud/observability/o11ydiagnosis/pkg/topru"
	"github.com/tidbcloud/observability/o11ydiagnosis/pkg/topsql"
	"github.com/tidbcloud/observability/o11ydiagnosis/pkg/topsqlv2"
	"go.uber.org/zap"
	"k8s.io/apimachinery/pkg/runtime"
	ctrl "sigs.k8s.io/controller-runtime"
	"sigs.k8s.io/controller-runtime/pkg/client"

	docs "github.com/tidbcloud/observability/api/diagnosisqueryswagger"
	o11yv1 "github.com/tidbcloud/observability/api/regionalserver/resourceworker/v1"
	"github.com/tidbcloud/observability/o11ydiagnosis/pkg/slowquery"
	"github.com/tidbcloud/observability/o11ydiagnosis/pkg/sqls"
	"github.com/tidbcloud/observability/o11ydiagnosis/pkg/storage"
	"github.com/tidbcloud/observability/o11ydiagnosis/pkg/utils/metacache"
)

// @title Diagnosis Query API
// @version 1.0
// @description API for TiDB Cloud Observability diagnosis: TopSQL, TopRU, Slow Query, Key Visualizer, Statements, Continuous Profiling. For AI and human developers.
// @license.name Apache-2.0
// @license.url https://www.apache.org/licenses/LICENSE-2.0.html
// @host localhost:8081
// @BasePath /api/v1
// @schemes http https
func main() {
	var (
		vendor                 string
		region                 string
		duckDBSwitchTime       int
		duckDBSwitchAllEnabled bool
		duckDBGrayTenantIds    string
		athenaWorkGroup        string
		athenaDatabase         string
		athenaCatalog          string
		athenaOutputLocation   string
	)

	fmt.Printf("os.Args in init(): %v\n", os.Args)
	flag.StringVar(&vendor, "vendor", "aws", "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.StringVar(&athenaWorkGroup, "athena-workgroup", "", "Athena workgroup for /api/v1/olap (AWS only)")
	flag.StringVar(&athenaDatabase, "athena-database", "", "Athena database for /api/v1/olap (AWS only)")
	flag.StringVar(&athenaCatalog, "athena-catalog", "", "Athena catalog for /api/v1/olap (AWS only)")
	flag.StringVar(&athenaOutputLocation, "athena-output-location", "", "Athena query result S3 location for /api/v1/olap (AWS only)")

	flag.Parse()

	o11yJwtAuthEnabled := strings.ToLower(strings.TrimSpace(os.Getenv("O11Y_JWT_AUTH_ENABLED"))) == "true"
	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),
		zap.Bool("o11yJwtAuthEnabled", o11yJwtAuthEnabled))
	// When O11Y_JWT_AUTH_ENABLED is enabled, token verification is performed locally with JWT_SECRET.

	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()))

	// Swagger UI: /swagger/index.html
	docs.SwaggerInfo.BasePath = "/api/v1"
	ng.GET("/swagger/*any", ginSwagger.WrapHandler(swaggerFiles.Handler))

	// Register pprof handlers
	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")
	apiv1.Use(apimetrics.Middleware())
	apiv1.Use(o11yJwtAuthMiddleware())

	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)

	topsqlsvcv2 := topsqlv2.NewSQLsService("/cache", store, kubecli, topsqlv2.DuckDBConfig{Vendor: vendor, Region: region}, topsqlsvc)
	topsqlv2.InitSQLsApi(apiv1, topsqlsvcv2)

	toprusvc := topru.NewService("/cache", store, kubecli, topsqlv2.DuckDBConfig{Vendor: vendor, Region: region})
	topru.InitTopRUApi(apiv1, toprusvc)

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

	sqlsvc := sqls.NewSQLsService("/cache", store, kubecli, sqls.DuckDBConfig{Vendor: vendor, Region: region})
	sqls.InitSQLsApi(apiv1, sqlsvc)

	conprofsvc := conprof.NewConProfService(vendor, store, "/cache")
	conprof.InitConProfApi(apiv1, conprofsvc)

	queryCfg := query.DefaultConfigFromEnv(vendor, region)
	queryCfg.ApplyOverrides(athenaWorkGroup, athenaDatabase, athenaCatalog, athenaOutputLocation)
	queryCfg.ApplyLogsSettings(os.Getenv("BUCKET_NAME"), os.Getenv("O11Y_LOGS_TABLE_SUFFIX"))
	querysvc, err := query.NewService(queryCfg, metacache.NewMetaCache(kubecli, store))
	if err != nil {
		log.Panic("failed to create query service", zap.Error(err))
	}
	query.InitOlapApi(apiv1, querysvc)

	corsConfig := cors.DefaultConfig()
	corsConfig.AllowAllOrigins = true
	corsConfig.AllowHeaders = append(corsConfig.AllowHeaders,
		"authorization",
		"x-provider",
		"x-region",
		"x-env",
		"x-biz-type",
		"x-org-id",
		"x-cluster-id",
		"x-deploy-type",
		"x-region-id",
		"x-application-id",
	)
	ng.Group("/external/slowquery").
		Use(cors.New(corsConfig)).
		Any("/:bizType/tenant/:tenantID/application/:clusterID/api/v1/*path", func(c *gin.Context) {
			bizType := c.Param("bizType")
			tenantID := c.Param("tenantID")
			clusterID := c.Param("clusterID")
			path := c.Param("path")
			c.Request.Header.Set("x-provider", vendor)
			c.Request.Header.Set("x-region", region)
			c.Request.Header.Set("x-biz-type", bizType)
			c.Request.Header.Set("x-org-id", tenantID)
			c.Request.Header.Set("x-cluster-id", clusterID)
			c.Request.Header.Set("x-deploy-type", "premium")
			if c.Request.Header.Get("x-env") == "" {
				switch bizType {
				case "dev-tidbcloud-nextgen":
					c.Request.Header.Set("x-env", "dev")
				case "staging-tidbcloud-nextgen":
					c.Request.Header.Set("x-env", "staging")
				case "prod-tidbcloud-nextgen":
					c.Request.Header.Set("x-env", "prod")
				}
			}
			c.Request.URL.Path = "/api/v1" + path
			ng.ServeHTTP(c.Writer, c.Request)
		})

	// Route for <o11y_id>/diagnosis/api/v1/ prefix
	// Parameters are extracted from headers:
	// - X-Region-Id: format <vendor>-<region>, e.g. aws-ap-southeast-1
	// - X-Biz-Type: biz type, e.g. byoc
	// - X-Org-Id: tidbcloud tenant id
	// - X-Application-Id: tidbcloud cluster id
	ng.Group("/:o11yId/diagnosis/api/v1").
		Use(cors.New(corsConfig)).
		Any("/*path", func(c *gin.Context) {
			path := c.Param("path")

			// Get parameters from headers
			regionId := c.GetHeader("X-Region-Id")
			bizType := c.GetHeader("X-Biz-Type")
			orgId := c.GetHeader("X-Org-Id")
			applicationId := c.GetHeader("X-Application-Id")

			// Parse X-Region-Id to extract vendor and region
			// Format: <vendor>-<region>, e.g. aws-ap-southeast-1
			var parsedVendor, parsedRegion string
			if regionId != "" {
				parts := strings.SplitN(regionId, "-", 2)
				if len(parts) == 2 {
					parsedVendor = parts[0]
					parsedRegion = parts[1]
				} else {
					log.Warn("Invalid X-Region-Id format", zap.String("region-id", regionId))
					c.JSON(http.StatusBadRequest, gin.H{
						"error": "Invalid X-Region-Id format, expected <vendor>-<region>",
					})
					return
				}
			} else {
				// Fallback to default vendor and region if not provided
				parsedVendor = vendor
				parsedRegion = region
			}

			// Set headers for downstream processing
			c.Request.Header.Set("x-provider", parsedVendor)
			c.Request.Header.Set("x-region", parsedRegion)

			if bizType != "" {
				c.Request.Header.Set("x-biz-type", bizType)
			}

			if orgId != "" {
				c.Request.Header.Set("x-org-id", orgId)
			}

			if applicationId != "" {
				c.Request.Header.Set("x-cluster-id", applicationId)
			}
			c.Request.Header.Set("x-deploy-type", "premium")
			// Set env based on biz type if not already set
			if c.Request.Header.Get("x-env") == "" && bizType != "" {
				switch bizType {
				case "dev-tidbcloud-nextgen":
					c.Request.Header.Set("x-env", "dev")
				case "staging-tidbcloud-nextgen":
					c.Request.Header.Set("x-env", "staging")
				case "prod-tidbcloud-nextgen":
					c.Request.Header.Set("x-env", "prod")
				}
			}
			// Rewrite path to /api/v1/...
			c.Request.URL.Path = "/api/v1" + path
			ng.ServeHTTP(c.Writer, c.Request)
		})

	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 "alicloud":
		//customEndpoint := fmt.Sprintf("https://o11y-prod-shared-%s.oss-%s-internal.aliyuncs.com", region, region)
		customEndpoint := fmt.Sprintf("http://oss-%s-internal.aliyuncs.com", region)
		var cfg aws.Config
		var err error
		log.Info("AWS_ENDPOINT_S3 detected: Using S3-Compatible (OSS) mode.", zap.String("endpoint", customEndpoint))
		customResolver := aws.EndpointResolverWithOptionsFunc(func(service, r string, options ...interface{}) (aws.Endpoint, error) {
			return aws.Endpoint{
				PartitionID:   "oss",
				URL:           customEndpoint,
				SigningRegion: region,
			}, nil
		})
		cfg, err = awsconfig.LoadDefaultConfig(context.TODO(),
			awsconfig.WithRegion(region),
			awsconfig.WithEndpointResolverWithOptions(customResolver))
		cfg.BaseEndpoint = aws.String(customEndpoint)
		if err != nil {
			panic(err)
		}
		return storage.NewAwsStorage(s3.NewFromConfig(cfg)), nil
	case "local":
		return storage.NewLocalStorage(), 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})
}