package database

import (
	"context"
	"database/sql"
	"fmt"
	"os"
	"sync"
	"time"

	"github.com/tidbcloud/infra-provider/cluster-prober/config"
	"github.com/tidbcloud/infra-provider/cluster-prober/prober"
	"github.com/tidbcloud/infra-provider/cluster-prober/prober/metrics"
	"github.com/tidbcloud/infra-provider/pkg/util"

	"github.com/go-logr/logr"
	_ "github.com/go-sql-driver/mysql"
	"github.com/prometheus/client_golang/prometheus"
)

type Endpoint struct {
	Addr string
	DB   *sql.DB
}

type Context struct {
	context.Context

	Metrics   *metrics.DBProberMetrics
	Logger    logr.Logger
	Endpoints []*Endpoint // endpoints to connect the database
}

func (c *Context) DefaultEndpoint() *Endpoint {
	return c.Endpoints[len(c.Endpoints)-1]
}

type DBProber interface {
	Name() string
	Probe(ctx *Context, ch chan<- prometheus.Metric) error
}

// Prober probes a cluster by accessing the database.
type Prober struct {
	cfg     *config.DatabaseProbeConfig
	metrics *metrics.DBProberMetrics

	user     string
	password *string
	logger   logr.Logger
	probers  []DBProber
}

func New(cfg *config.DatabaseProbeConfig, deps *prober.Dependency) (*Prober, error) {
	user, err := os.ReadFile(cfg.UserFile)
	if err != nil {
		return nil, fmt.Errorf("unable to read user: %w", err)
	}
	var password *string
	if cfg.PasswordFile != nil {
		raw, err := os.ReadFile(*cfg.PasswordFile)
		if err != nil {
			return nil, fmt.Errorf("unable to read password: %w", err)
		}
		password = util.Pointer(string(raw))
	}

	probers := []DBProber{
		&dbWriteProber{},
	}

	return &Prober{
		cfg:      cfg,
		metrics:  metrics.NewDBProberMetrics(cfg.ClusterID),
		user:     string(user),
		password: password,
		logger:   deps.Logger.WithName("DBProber"),
		probers:  probers,
	}, nil
}

// Describe implements prometheus.Collector.
func (p *Prober) Describe(ch chan<- *prometheus.Desc) {
	for _, desc := range p.metrics.Desc() {
		ch <- desc
	}
}

// Collect implements prometheus.Collector.
func (p *Prober) Collect(ch chan<- prometheus.Metric) {
	p.logger.V(4).Info("start probe a cluster over database")

	p.probe(context.Background(), ch)

	p.logger.V(4).Info("probe the cluster over database done")
}

func (p *Prober) probe(ctx context.Context, ch chan<- prometheus.Metric) {
	pctx := &Context{
		Metrics:   p.metrics,
		Logger:    p.logger,
		Endpoints: make([]*Endpoint, 0, len(p.cfg.Addrs)),
	}

	// open db
	endpoints := p.openDBAndPing(ctx, ch)

	// remove unavailable endpoints
	for _, conn := range endpoints {
		if conn != nil {
			pctx.Endpoints = append(pctx.Endpoints, conn)
		}
	}
	if len(pctx.Endpoints) == 0 {
		p.logger.Error(fmt.Errorf("no available database connection"), "skip probe")
		return
	}

	// other probes
	p.logger.V(4).Info("probe the cluster over database with other probers")
	ctx, cancel := context.WithTimeout(context.Background(), p.cfg.ProbeTimeout.Duration)
	defer cancel()
	pctx.Context = ctx
	wg := sync.WaitGroup{}
	for _, prober := range p.probers {
		wg.Add(1)
		go func(prober DBProber) {
			defer wg.Done()
			err := prober.Probe(pctx, ch)
			if err != nil {
				pctx.Logger.Error(err, "probe failed", "prober", prober.Name())
			}
		}(prober)
	}
	wg.Wait()

	// close all db connections
	for _, conn := range pctx.Endpoints {
		if err := conn.DB.Close(); err != nil {
			p.logger.Error(err, "close database failed", "addr", conn.Addr)
		}
	}
}

func (p *Prober) openDBAndPing(ctx context.Context, ch chan<- prometheus.Metric) []*Endpoint {
	wg := sync.WaitGroup{}
	endpoints := make([]*Endpoint, len(p.cfg.Addrs))
	ctx, cancel := context.WithTimeout(ctx, p.cfg.ProbeTimeout.Duration)
	defer cancel()

	// open db and ping concurrently
	for i := range p.cfg.Addrs {
		wg.Add(1)

		go func(index int) {
			defer wg.Done()

			addr := p.cfg.Addrs[index]
			dsn := buildDSN(addr, p.user, p.password)
			db, err := sql.Open("mysql", dsn)
			if err != nil {
				p.logger.Error(err, "open database failed", "addr", addr)
				return
			}
			db.SetMaxOpenConns(1)
			db.SetMaxIdleConns(1)
			db.SetConnMaxLifetime(1 * time.Minute)

			err = db.PingContext(ctx)
			if err != nil {
				p.logger.Error(err, "ping database failed", "addr", addr)
				ch <- p.metrics.DBPing(addr, false)
				return
			}
			ch <- p.metrics.DBPing(addr, true)

			endpoints[index] = &Endpoint{
				Addr: addr,
				DB:   db,
			}
		}(i)
	}
	wg.Wait()

	return endpoints
}

func buildDSN(addr string, user string, password *string) string {
	if password != nil {
		user = fmt.Sprintf("%s:%s", user, *password)
	}
	return fmt.Sprintf("%s@(%s)/?charset=utf8", user, addr)
}
