package harness

import (
	"context"
	"fmt"
	"log"
	"net"
	"net/http"
	"net/url"
	"os"
	"strconv"
	"strings"
	"testing"
	"time"

	"github.com/google/uuid"
	"github.com/jackc/pgx/v5"
	"github.com/testcontainers/testcontainers-go"
	"github.com/testcontainers/testcontainers-go/modules/postgres"
	"github.com/testcontainers/testcontainers-go/modules/rabbitmq"
	"go.uber.org/goleak"

	amqp "github.com/rabbitmq/amqp091-go"

	"github.com/hatchet-dev/hatchet/cmd/hatchet-admin/cli/seed"
	"github.com/hatchet-dev/hatchet/cmd/hatchet-engine/engine"
	"github.com/hatchet-dev/hatchet/cmd/hatchet-migrate/migrate"
	"github.com/hatchet-dev/hatchet/pkg/config/database"
	"github.com/hatchet-dev/hatchet/pkg/config/loader"
	"github.com/hatchet-dev/hatchet/pkg/encryption"
	"github.com/hatchet-dev/hatchet/pkg/random"
)

func getEnvConfig() (string, bool, string, bool, bool) {
	// Get migration strategy: penultimate or latest
	migrateStrategy := os.Getenv("TESTING_MATRIX_MIGRATE")
	if migrateStrategy == "" {
		migrateStrategy = "latest" // Default value
	}

	// Get RabbitMQ enabled status
	rabbitmqEnabled := strings.ToLower(os.Getenv("TESTING_MATRIX_RABBITMQ_ENABLED")) == "true"

	// Get PostgreSQL version
	pgVersion := os.Getenv("TESTING_MATRIX_PG_VERSION")
	if pgVersion == "" {
		pgVersion = "16-alpine" // Default value
	}

	// Get whether optimistic scheduling is enabled
	isOptimistic := strings.ToLower(os.Getenv("TESTING_MATRIX_OPTIMISTIC_SCHEDULING")) == "true"

	// Get whether PgBouncer connection pooling is enabled
	pgBouncerEnabled := strings.ToLower(os.Getenv("TESTING_MATRIX_PGBOUNCER_ENABLED")) == "true"

	return migrateStrategy, rabbitmqEnabled, pgVersion, isOptimistic, pgBouncerEnabled
}

func RunTestWithEngine(m *testing.M) {
	// This runs before all tests
	cleanup := startEngine()

	// Run the tests
	exitCode := m.Run()

	// This runs after all tests
	cleanup()

	// allow a bit of time for the engine to shut down
	time.Sleep(2 * time.Second)

	if exitCode == 0 {
		if err := goleak.Find(
			goleak.IgnoreTopFunction("github.com/hashicorp/golang-lru/v2/expirable.NewLRU[...].func1"),
			goleak.IgnoreTopFunction("github.com/hatchet-dev/hatchet/internal/cache.NewTTL[...].func1"),
			goleak.IgnoreTopFunction("google.golang.org/grpc/internal/resolver/dns.(*dnsResolver).watcher"),
			goleak.IgnoreTopFunction("github.com/testcontainers/testcontainers-go.(*Reaper).connect.func1"),
			goleak.IgnoreTopFunction("go.opencensus.io/stats/view.(*worker).start"),
			goleak.IgnoreTopFunction("google.golang.org/grpc/internal/grpcsync.(*CallbackSerializer).run"),
			goleak.IgnoreTopFunction("internal/poll.runtime_pollWait"),
			goleak.IgnoreTopFunction("google.golang.org/grpc/internal/transport.(*controlBuffer).get"),
			// all engine related packages
			goleak.IgnoreTopFunction("github.com/jackc/pgx/v5/pgxpool.(*Pool).backgroundHealthCheck"),
			goleak.IgnoreTopFunction("github.com/rabbitmq/amqp091-go.(*Connection).heartbeater"),
			goleak.IgnoreTopFunction("github.com/rabbitmq/amqp091-go.(*consumers).buffer"),
			goleak.IgnoreTopFunction("google.golang.org/grpc/internal/transport.(*http2Server).keepalive"),
			goleak.IgnoreTopFunction("github.com/hashicorp/golang-lru/v2/expirable.NewLRU[...].func1"),
			goleak.IgnoreAnyFunction("github.com/hashicorp/golang-lru/v2/expirable.(*LRU[...]).deleteExpired"),
		); err != nil {
			fmt.Fprintf(os.Stderr, "goleak: Errors on successful test run: %v\n", err)
			exitCode = 1
		}
	}

	os.Exit(exitCode)
}

func startEngine() func() {
	setTestingKeysInEnv()

	ctx, cancel := context.WithCancel(context.Background())

	// Get configuration values from environment
	migrateStrategy, rabbitmqEnabled, pgVersion, isOptimistic, pgBouncerEnabled := getEnvConfig()

	log.Printf("Starting engine with migration strategy: %s, RabbitMQ enabled: %t, PostgreSQL version: %s, PgBouncer enabled: %t", migrateStrategy, rabbitmqEnabled, pgVersion, pgBouncerEnabled)

	postgresConnStr, cleanupPostgres := startPostgres(ctx, pgVersion)

	// Start PgBouncer if enabled, but don't use it yet (migrations and seeding go directly to postgres)
	var pgBouncerConnStr string
	var cleanupPgBouncer func() error

	if pgBouncerEnabled {
		pgPort, err := extractPort(postgresConnStr)
		if err != nil {
			log.Fatalf("failed to extract postgres port from connection string: %v", err)
		}

		pgBouncerConnStr, cleanupPgBouncer = startPgBouncer(ctx, pgPort)
	} else {
		cleanupPgBouncer = func() error { return nil }
	}

	grpcPort, err := getFreePort()

	if err != nil {
		log.Fatalf("failed to find available port: %v", err)
	}

	healthPort, err := getFreePort()

	if err != nil {
		log.Fatalf("failed to find available port: %v", err)
	}

	// Use postgres directly for migrations and seeding
	os.Setenv("DATABASE_URL", postgresConnStr)
	os.Setenv("SERVER_GRPC_INSECURE", "true")
	os.Setenv("SERVER_GRPC_PORT", strconv.Itoa(grpcPort))
	os.Setenv("SERVER_GRPC_BROADCAST_ADDRESS", fmt.Sprintf("localhost:%d", grpcPort))
	os.Setenv("SERVER_HEALTHCHECK", "true")
	os.Setenv("SERVER_HEALTHCHECK_PORT", strconv.Itoa(healthPort))
	os.Setenv("HATCHET_CLIENT_TLS_STRATEGY", "none")
	os.Setenv("SERVER_AUTH_COOKIE_DOMAIN", "app.dev.hatchet-tools.com")
	os.Setenv("SERVER_LOGGER_LEVEL", "error")
	os.Setenv("SERVER_LOGGER_FORMAT", "console")
	os.Setenv("DATABASE_LOGGER_LEVEL", "error")
	os.Setenv("DATABASE_LOGGER_FORMAT", "console")
	os.Setenv("SERVER_SECURITY_CHECK_ENABLED", "false")
	os.Setenv("SERVER_ADDITIONAL_LOGGERS_QUEUE_LEVEL", "error")
	os.Setenv("SERVER_ADDITIONAL_LOGGERS_QUEUE_FORMAT", "console")
	os.Setenv("SERVER_ADDITIONAL_LOGGERS_PGXSTATS_LEVEL", "error")
	os.Setenv("SERVER_ADDITIONAL_LOGGERS_PGXSTATS_FORMAT", "console")
	os.Setenv("SERVER_DEFAULT_ENGINE_VERSION", "V1")
	os.Setenv("SERVER_ENABLE_DURABLE_USER_EVENT_LOG", "true")

	var cleanupRabbitMQ func() error
	if rabbitmqEnabled {
		rabbitMQConnStr, rabbitMQCleanup := startRabbitMQ(ctx)
		os.Setenv("SERVER_MSGQUEUE_KIND", "rabbitmq")
		os.Setenv("SERVER_MSGQUEUE_RABBITMQ_URL", rabbitMQConnStr)
		cleanupRabbitMQ = rabbitMQCleanup
	} else {
		os.Setenv("SERVER_MSGQUEUE_KIND", "postgres")
		cleanupRabbitMQ = func() error { return nil }
	}

	if isOptimistic {
		os.Setenv("SERVER_OPTIMISTIC_SCHEDULING_ENABLED", "true")
	} else {
		os.Setenv("SERVER_OPTIMISTIC_SCHEDULING_ENABLED", "false")
	}

	// Run migrations
	if migrateStrategy == "penultimate" {
		migrate.RunMigrations(ctx, migrate.WithUpToPenultimate())
	} else {
		migrate.RunMigrations(ctx)
	}

	// Set higher rate limit for load tests
	os.Setenv("SCHEDULER_CONCURRENCY_RATE_LIMIT", "1000")

	cf := loader.NewConfigLoader("")

	dl, err := cf.InitDataLayer()

	if err != nil {
		log.Fatalf("failed to initialize data layer: %v", err)
	}

	// seed database
	seedDatabase(dl)

	if err := dl.Disconnect(); err != nil {
		log.Fatalf("failed to disconnect data layer: %v", err)
	}

	// set the API token
	tenantUUID, err := uuid.Parse(dl.Seed.DefaultTenantID)
	if err != nil {
		log.Fatalf("failed to parse default tenant ID: %v", err)
	}
	setAPIToken(ctx, cf, tenantUUID)

	// Switch to PgBouncer for the engine if enabled
	if pgBouncerEnabled {
		log.Printf("Switching main pool to PgBouncer: %s", pgBouncerConnStr)
		os.Setenv("DATABASE_PGBOUNCER_URL", pgBouncerConnStr)
	}

	engineCh := make(chan error)

	go func() {
		engineCh <- engine.Run(ctx, cf, "testing")
	}()

	// Return a cleanup function that properly handles shutdown
	return func() {
		cancel()

		err := <-engineCh

		if err != nil {
			log.Fatalf("failed to run engine: %v", err)
		}

		if pgBouncerEnabled {
			err = cleanupPgBouncer()

			if err != nil {
				log.Fatalf("failed to cleanup pgbouncer: %v", err)
			}
		}

		err = cleanupPostgres()

		if err != nil {
			log.Fatalf("failed to cleanup postgres: %v", err)
		}

		if rabbitmqEnabled {
			err = cleanupRabbitMQ()

			if err != nil {
				log.Fatalf("failed to cleanup rabbitmq: %v", err)
			}
		}
	}
}

func startPostgres(ctx context.Context, pgVersion string) (string, func() error) {
	// Find an available port for PostgreSQL
	pgPort, err := findAvailablePort(5432)
	if err != nil {
		log.Fatalf("failed to find available port for postgres: %v", err)
	}

	postgresContainer, err := postgres.Run(
		ctx,
		fmt.Sprintf("postgres:%s", pgVersion),
		postgres.WithDatabase("test"),
		postgres.WithUsername("user"),
		postgres.WithPassword("password"),
		testcontainers.WithHostPortAccess(pgPort),
	)

	if err != nil {
		log.Fatalf("failed to start postgres container: %v", err)
	}

	connStr, err := postgresContainer.ConnectionString(ctx, "sslmode=disable")
	if err != nil {
		log.Fatalf("failed to get connection string: %v", err)
	}

	// loop until the database is ready
	for i := 0; i < 10; i++ {
		var db *pgx.Conn
		db, err = pgx.Connect(ctx, connStr)

		if err != nil {
			time.Sleep(time.Second * 2)
			continue
		}

		// make sure we can ping the database
		err = db.Ping(ctx)

		if err != nil {
			time.Sleep(time.Second * 2)
			continue
		}

		db.Close(ctx)

		return connStr, func() error {
			ctx, cancel := context.WithTimeout(context.Background(), time.Second*10)
			defer cancel()
			if err := postgresContainer.Terminate(ctx); err != nil {
				return fmt.Errorf("failed to terminate postgres container: %w", err)
			}
			return nil
		}
	}

	log.Fatalf("failed to connect to postgres container after 10 attempts: %v", err)

	// this should never be reached
	return "", func() error {
		return nil
	}
}

func startRabbitMQ(ctx context.Context) (string, func() error) {
	rabbitContainer, err := rabbitmq.Run(
		ctx,
		"rabbitmq:3-management-alpine",
	)

	if err != nil {
		log.Fatalf("failed to start rabbitmq container: %v", err)
	}

	// Get the connection URL for RabbitMQ
	amqpURI, err := rabbitContainer.AmqpURL(ctx)
	if err != nil {
		log.Fatalf("failed to get AMQP URL: %v", err)
	}

	// loop until RabbitMQ is ready
	for i := 0; i < 10; i++ {
		var conn *amqp.Connection
		conn, err = amqp.Dial(amqpURI)

		if err != nil {
			time.Sleep(time.Second * 2)
			continue
		}

		// make sure we can create a channel
		var ch *amqp.Channel
		ch, err = conn.Channel()

		if err != nil {
			conn.Close()
			time.Sleep(time.Second * 2)
			continue
		}

		ch.Close()
		conn.Close()

		return amqpURI, func() error {
			ctx, cancel := context.WithTimeout(context.Background(), time.Second*10)
			defer cancel()
			if err := rabbitContainer.Terminate(ctx); err != nil {
				return fmt.Errorf("failed to terminate rabbitmq container: %w", err)
			}
			return nil
		}
	}

	log.Fatalf("failed to connect to rabbitmq container after 10 attempts: %v", err)

	// this should never be reached
	return "", func() error {
		return nil
	}
}

func extractPort(connStr string) (int, error) {
	u, err := url.Parse(connStr)
	if err != nil {
		return 0, fmt.Errorf("failed to parse connection string: %w", err)
	}

	port, err := strconv.Atoi(u.Port())
	if err != nil {
		return 0, fmt.Errorf("failed to parse port from connection string: %w", err)
	}

	return port, nil
}

func startPgBouncer(ctx context.Context, pgPort int) (string, func() error) {
	pgBouncerContainer, err := testcontainers.GenericContainer(ctx, testcontainers.GenericContainerRequest{
		ContainerRequest: testcontainers.ContainerRequest{
			Image:           "edoburu/pgbouncer:v1.25.1-p0",
			ExposedPorts:    []string{"5432/tcp"},
			HostAccessPorts: []int{pgPort},
			Env: map[string]string{
				"DATABASE_URL":            fmt.Sprintf("postgres://user:password@host.testcontainers.internal:%d/test", pgPort),
				"POOL_MODE":               "transaction",
				"MAX_CLIENT_CONN":         "500",
				"DEFAULT_POOL_SIZE":       "50",
				"AUTH_TYPE":               "scram-sha-256",
				"MAX_PREPARED_STATEMENTS": "256",
			},
		},
		Started: true,
	})

	if err != nil {
		log.Fatalf("failed to start pgbouncer container: %v", err)
	}

	host, err := pgBouncerContainer.Host(ctx)
	if err != nil {
		log.Fatalf("failed to get pgbouncer host: %v", err)
	}

	mappedPort, err := pgBouncerContainer.MappedPort(ctx, "5432/tcp")
	if err != nil {
		log.Fatalf("failed to get pgbouncer mapped port: %v", err)
	}

	connStr := fmt.Sprintf(
		"postgresql://user:password@%s:%s/test?sslmode=disable",
		host, mappedPort.Port(),
	)

	// loop until pgbouncer is ready
	for i := 0; i < 10; i++ {
		var db *pgx.Conn
		db, err = pgx.Connect(ctx, connStr)

		if err != nil {
			time.Sleep(time.Second * 2)
			continue
		}

		err = db.Ping(ctx)

		if err != nil {
			db.Close(ctx)
			time.Sleep(time.Second * 2)
			continue
		}

		db.Close(ctx)

		return connStr, func() error {
			ctx, cancel := context.WithTimeout(context.Background(), time.Second*30)
			defer cancel()

			err = pgBouncerContainer.Terminate(ctx)
			if err != nil {
				return fmt.Errorf("failed to terminate pgbouncer container: %w", err)
			}

			return nil
		}
	}

	log.Fatalf("failed to connect to pgbouncer container after 10 attempts: %v", err)

	// this should never be reached
	return "", func() error {
		return nil
	}
}

func seedDatabase(dc *database.Layer) {
	log.Printf("Seeding database")

	err := seed.SeedDatabase(dc)

	if err != nil {
		log.Fatalf("could not seed database: %v", err)
	}

	log.Printf("Seeding database complete")
}

func setAPIToken(ctx context.Context, cf *loader.ConfigLoader, tenantID uuid.UUID) {
	log.Printf("Generating API token for Hatchet server")

	cleanup, server, err := cf.CreateServerFromConfig("testing")

	if err != nil {
		log.Fatalf("could not create server config: %v", err)
	}

	expiresAt := time.Now().Add(time.Hour * 24 * 30)

	defaultTok, err := server.Auth.JWTManager.GenerateTenantToken(
		ctx,
		tenantID,
		"testing",
		false,
		&expiresAt,
	)

	if err != nil {
		log.Fatalf("could not generate token: %v", err)
	}

	err = cleanup()

	if err != nil {
		log.Fatalf("could not cleanup server: %v", err)
	}

	err = server.Disconnect()

	if err != nil {
		log.Fatalf("could not disconnect server: %v", err)
	}

	os.Setenv("HATCHET_CLIENT_TOKEN", defaultTok.Token)

	log.Printf("Generated API token for tenant %s", tenantID)
}

func setTestingKeysInEnv() {
	log.Println("Generating encryption keys for Hatchet server")

	cookieHashKey, err := random.Generate(16)

	if err != nil {
		log.Fatalf("could not generate hash key for instance: %v", err)
	}

	cookieBlockKey, err := random.Generate(16)

	if err != nil {
		log.Fatalf("could not generate block key for instance: %v", err)
	}

	_ = os.Setenv("SERVER_AUTH_COOKIE_SECRETS", fmt.Sprintf("%s %s", cookieHashKey, cookieBlockKey))

	masterKeyBytes, privateEc256, publicEc256, _, err := encryption.GenerateLocalKeys()

	if err != nil {
		log.Fatalf("could not generate local keys: %v", err)
	}

	_ = os.Setenv("SERVER_ENCRYPTION_MASTER_KEYSET", string(masterKeyBytes))
	_ = os.Setenv("SERVER_ENCRYPTION_JWT_PRIVATE_KEYSET", string(privateEc256))
	_ = os.Setenv("SERVER_ENCRYPTION_JWT_PUBLIC_KEYSET", string(publicEc256))
}

// findAvailablePort returns an available port starting from the given port
// It will only attempt a maximum of 100 ports before giving up
func findAvailablePort(startPort int) (int, error) {
	port := startPort
	maxAttempts := 100
	// Use min to limit the search to either port+100 or 65535, whichever is smaller
	maxPort := min(startPort+maxAttempts-1, 65535)

	for port <= maxPort {
		addr := ":" + strconv.Itoa(port)
		listener, err := net.Listen("tcp", addr)

		if err == nil {
			// Port is available, close the listener and return the port
			listener.Close()
			return port, nil
		}

		// Try the next port
		port++
	}

	return 0, fmt.Errorf("no available port found in range %d-%d", startPort, maxPort)
}

func WaitEngineReady(ctx context.Context, timeout time.Duration) error {
	ctx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()

	healthPort := os.Getenv("SERVER_HEALTHCHECK_PORT")
	if healthPort == "" {
		return fmt.Errorf("SERVER_HEALTHCHECK_PORT not set")
	}

	addr := fmt.Sprintf("http://localhost:%s/ready", healthPort)

	ticker := time.NewTicker(time.Second)
	defer ticker.Stop()

	for {
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-ticker.C:
			req, _ := http.NewRequestWithContext(ctx, http.MethodGet, addr, nil)
			resp, err := http.DefaultClient.Do(req) //nolint:gosec
			if err != nil {
				continue
			}
			resp.Body.Close()
			if resp.StatusCode == http.StatusOK {
				return nil
			}
		}
	}
}

func getFreePort() (int, error) {
	listener, err := net.Listen("tcp", ":0") //nolint:gosec
	if err != nil {
		return 0, err
	}
	port := listener.Addr().(*net.TCPAddr).Port
	_ = listener.Close()
	return port, nil
}
