package firefly

import (
	"context"
	"flag"
	"github.com/jackc/pgx/v4"
	"github.com/jackc/pgx/v4/pgxpool"
	"github.com/sirupsen/logrus"
)

var (
	pgsqlConn = flag.String("pgsql", "", "Connection spec for postgresql database")
	pool      *pgxpool.Pool
)

func Initialize() {
	var err error
	config, err := pgxpool.ParseConfig(*pgsqlConn)
	if err != nil {
		panic(err)
	}
	if config.MaxConns == 0 {
		config.MaxConns = 8
	}
	if pool, err = pgxpool.Connect(context.Background(), *pgsqlConn); err != nil {
		panic(err)
	}
	if conn, err := pool.Acquire(context.Background()); err != nil {
		panic(err)
	} else {
		conn.Release()
	}
}

func AddSpeedTestResult(ip string, speed map[string]float64) {
	conn, err := pool.Acquire(context.Background())
	if err != nil {
		logrus.WithError(err).Error("Unable to get database connection")
		return
	}
	defer conn.Release()
	tx, err := conn.Begin(context.Background())
	if err != nil {
		logrus.WithError(err).Error("Unable to start database transaction")
		return
	}
	row := tx.QueryRow(context.Background(),
		"INSERT INTO speedtest_entry (id, ip, timestamp) VALUES (DEFAULT, $1, now()) RETURNING id", ip)
	var id int64
	if err = row.Scan(&id); err != nil {
		logrus.WithError(err).Error("Unable to insert new speedtest result entry")
		_ = tx.Rollback(context.Background())
		return
	}
	batch := &pgx.Batch{}
	for name, value := range speed {
		batch.Queue("INSERT INTO speedtest_result (entry_id, node, speed) VALUES ($1, $2, $3)",
			id, name, value)
	}
	results := tx.SendBatch(context.Background(), batch)
	if err = results.Close(); err != nil {
		logrus.WithError(err).Error("Unable to insert speedtest results")
		_ = tx.Rollback(context.Background())
		return
	}
	if err = tx.Commit(context.Background()); err != nil {
		logrus.WithError(err).Error("Unable to commit the database transaction")
		_ = tx.Rollback(context.Background())
	}
}