package observability

import (
	"context"
	"errors"
	"fmt"
	"log/slog"
	"net"
	"net/http"
	"strings"
	"sync"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"go.opentelemetry.io/otel/propagation"
	"go.opentelemetry.io/otel/trace"
	"go.opentelemetry.io/otel/trace/noop"

	"github.com/tidbcloud/AgentTeam/internal/config"
	"github.com/tidbcloud/AgentTeam/internal/observability/logging"
	"github.com/tidbcloud/AgentTeam/internal/observability/metrics"
	"github.com/tidbcloud/AgentTeam/internal/observability/tracing"
)

const metricsServerShutdownTimeout = 5 * time.Second

type Option func(*Provider)

type Provider struct {
	cfg             config.ObservabilityConfig
	logger          *slog.Logger
	reg             *prometheus.Registry
	server          *http.Server
	http            *metrics.HTTPMiddleware
	daemonTransport *metrics.DaemonTransport
	httpTracing     *tracing.HTTPServer
	tracer          trace.Tracer
	tracerProvider  trace.TracerProvider
	traceShutdown   func(context.Context) error
	propagator      propagation.TextMapPropagator

	mu          sync.RWMutex
	address     string
	env         string
	shutdownMu  sync.Mutex
	metricsDone bool
	traceDone   bool
}

func WithEnvironment(environment string) Option {
	return func(provider *Provider) {
		provider.env = strings.TrimSpace(environment)
	}
}

func WithTracerProvider(tracerProvider trace.TracerProvider) Option {
	return func(provider *Provider) {
		if tracerProvider != nil {
			provider.tracerProvider = tracerProvider
		}
	}
}

func WithPropagator(propagator propagation.TextMapPropagator) Option {
	return func(provider *Provider) {
		if propagator != nil {
			provider.propagator = propagator
		}
	}
}

func New(cfg config.ObservabilityConfig, logger *slog.Logger, dbPool metrics.DBPoolSnapshotProvider, options ...Option) (*Provider, error) {
	provider, err := NewBase(cfg, logger, options...)
	if err != nil {
		return nil, err
	}
	if dbPool == nil {
		provider.logger.Warn("database pool metrics disabled: db pool is not configured")
	} else if err := provider.RegisterDBPool(dbPool); err != nil {
		provider.shutdownAfterInitializationFailure()
		return nil, err
	}
	return provider, nil
}

func NewBase(cfg config.ObservabilityConfig, logger *slog.Logger, options ...Option) (*Provider, error) {
	if logger == nil {
		logger = slog.Default()
	}

	provider := &Provider{
		cfg:        cfg,
		logger:     logger,
		reg:        prometheus.NewRegistry(),
		propagator: tracing.NewPropagator(),
	}
	for _, option := range options {
		if option != nil {
			option(provider)
		}
	}
	if provider.tracerProvider == nil {
		tracerProvider, shutdown, err := newConfiguredTracerProvider(context.Background(), cfg, provider.env)
		if err != nil {
			return nil, err
		}
		provider.tracerProvider = tracerProvider
		provider.traceShutdown = shutdown
	}
	provider.tracer = provider.tracerProvider.Tracer(tracing.InstrumentationName)
	provider.httpTracing = tracing.NewHTTPServer(provider.tracer, provider.propagator)
	if cfg.DBStatementEnabled && strings.EqualFold(provider.env, "production") {
		provider.logger.Warn("database SQL statement tracing is enabled; tracing records caller-supplied SQL text, audit query construction before enabling in production")
	}

	if err := metrics.RegisterRuntime(provider.reg); err != nil {
		provider.shutdownAfterInitializationFailure()
		return nil, fmt.Errorf("register runtime metrics: %w", err)
	}
	httpMetrics, err := metrics.RegisterHTTP(provider.reg)
	if err != nil {
		provider.shutdownAfterInitializationFailure()
		return nil, fmt.Errorf("register http metrics: %w", err)
	}
	provider.http = httpMetrics
	daemonTransportMetrics, err := metrics.RegisterDaemonTransport(provider.reg)
	if err != nil {
		provider.shutdownAfterInitializationFailure()
		return nil, fmt.Errorf("register daemon transport metrics: %w", err)
	}
	provider.daemonTransport = daemonTransportMetrics

	metricsAddress := strings.TrimSpace(cfg.MetricsAddress)
	if metricsAddress != "" {
		if strings.EqualFold(provider.env, "production") && isBindAllInterfaces(metricsAddress) {
			logger.Warn(
				"observability metrics listener binds all interfaces; enforce network isolation",
				slog.String("address", metricsAddress),
			)
		}
		provider.server = newMetricsServer(metricsAddress, provider.reg, provider.logger)
	}

	return provider, nil
}

func (provider *Provider) RegisterDBPool(dbPool metrics.DBPoolSnapshotProvider) error {
	if provider == nil {
		return nil
	}
	if dbPool == nil {
		return errors.New("database pool metrics provider is nil")
	}
	if err := metrics.RegisterDBPool(provider.reg, dbPool); err != nil {
		return fmt.Errorf("register database pool metrics: %w", err)
	}
	return nil
}

func (provider *Provider) Registry() prometheus.Registerer {
	if provider == nil {
		return nil
	}
	return provider.reg
}

func (provider *Provider) HTTPMetrics() *metrics.HTTPMiddleware {
	if provider == nil {
		return nil
	}
	return provider.http
}

func (provider *Provider) HTTPTracing() *tracing.HTTPServer {
	if provider == nil {
		return nil
	}
	return provider.httpTracing
}

func (provider *Provider) DaemonTransportMetrics() *metrics.DaemonTransport {
	if provider == nil {
		return nil
	}
	return provider.daemonTransport
}

func (provider *Provider) WrapHTTPClient(client *http.Client, name string) *http.Client {
	if provider == nil {
		return tracing.WrapHTTPClient(client, nil, nil, name)
	}
	return tracing.WrapHTTPClient(client, provider.tracer, provider.propagator, name)
}

func (provider *Provider) SQLOpener() tracing.SQLOpener {
	if provider == nil {
		return nil
	}
	return tracing.NewSQLOpener(provider.tracerProvider, tracing.DBOptions{
		StatementEnabled: provider.cfg.DBStatementEnabled,
	})
}

func (provider *Provider) Tracer() trace.Tracer {
	if provider == nil {
		return noop.NewTracerProvider().Tracer(tracing.InstrumentationName)
	}
	if provider.tracer == nil {
		return noop.NewTracerProvider().Tracer(tracing.InstrumentationName)
	}
	return provider.tracer
}

func (provider *Provider) Address() string {
	if provider == nil {
		return ""
	}
	provider.mu.RLock()
	defer provider.mu.RUnlock()
	return provider.address
}

func (provider *Provider) Run(ctx context.Context) error {
	if provider == nil || provider.server == nil {
		return nil
	}

	listener, err := net.Listen("tcp", provider.server.Addr)
	if err != nil {
		return err
	}
	provider.setAddress(listener.Addr().String())

	errCh := make(chan error, 1)
	logging.Go(provider.logger, func() {
		provider.logger.Info("metrics server listening", slog.String("address", provider.Address()))
		errCh <- provider.server.Serve(listener)
	})

	select {
	case err := <-errCh:
		return err
	case <-ctx.Done():
		shutdownCtx, cancel := context.WithTimeout(context.Background(), metricsServerShutdownTimeout)
		defer cancel()
		if err := provider.Shutdown(shutdownCtx); err != nil && !errors.Is(err, http.ErrServerClosed) {
			return err
		}
		return <-errCh
	}
}

func (provider *Provider) Shutdown(ctx context.Context) error {
	if provider == nil {
		return nil
	}
	provider.shutdownMu.Lock()
	defer provider.shutdownMu.Unlock()

	var shutdownErrors []error
	if !provider.metricsDone && provider.server != nil {
		if err := provider.server.Shutdown(ctx); err != nil {
			if errors.Is(err, http.ErrServerClosed) {
				provider.metricsDone = true
			} else {
				shutdownErrors = append(shutdownErrors, err)
			}
		} else {
			provider.metricsDone = true
		}
	}
	if !provider.traceDone && provider.traceShutdown != nil {
		if err := provider.traceShutdown(ctx); err != nil {
			shutdownErrors = append(shutdownErrors, err)
		} else {
			provider.traceDone = true
		}
	}
	return errors.Join(shutdownErrors...)
}

func (provider *Provider) shutdownAfterInitializationFailure() {
	if provider == nil {
		return
	}
	shutdownCtx, cancel := context.WithTimeout(context.Background(), metricsServerShutdownTimeout)
	defer cancel()
	if shutdownErr := provider.Shutdown(shutdownCtx); shutdownErr != nil && !errors.Is(shutdownErr, http.ErrServerClosed) {
		provider.logger.Warn("observability shutdown after initialization failure", slog.Any("error", shutdownErr))
	}
}

func (provider *Provider) setAddress(address string) {
	provider.mu.Lock()
	defer provider.mu.Unlock()
	provider.address = address
}

func newMetricsServer(address string, reg *prometheus.Registry, logger *slog.Logger) *http.Server {
	mux := http.NewServeMux()
	mux.Handle("/metrics", promhttp.HandlerFor(reg, promhttp.HandlerOpts{}))
	return &http.Server{
		Addr:              address,
		Handler:           mux,
		ErrorLog:          logging.HTTPErrorLog(logger),
		ReadHeaderTimeout: 5 * time.Second,
	}
}

func isBindAllInterfaces(address string) bool {
	host, _, err := net.SplitHostPort(strings.TrimSpace(address))
	if err != nil {
		return false
	}
	switch host {
	case "", "0.0.0.0", "::", "[::]":
		return true
	default:
		return false
	}
}