package database

import (
	"database/sql"
	"fmt"

	"github.com/tidbcloud/infra-provider/pkg/util"

	"github.com/prometheus/client_golang/prometheus"
)

const (
	probedDB          = "test"
	probedTable       = "db_write_prober"
	probedDataSize    = 1024
	probedRecordCount = 3
)

// dbWriteProber writes a table to probe whether the database is normal.
// Because the connection may be up but the database can't be read or written
type dbWriteProber struct{}

func (p *dbWriteProber) Name() string {
	return "db_write_prober"
}

func (p *dbWriteProber) Probe(ctx *Context, ch chan<- prometheus.Metric) error {
	table := fmt.Sprintf("%s.%s", probedDB, probedTable)
	err := p.probe(ctx.DefaultEndpoint().DB, table, probedRecordCount, probedDataSize)
	if err != nil {
		ch <- ctx.Metrics.DBWrite(false)
		return err
	}
	ch <- ctx.Metrics.DBWrite(true)
	return nil
}

func (p *dbWriteProber) probe(conn *sql.DB, table string, recordCnt, dataSize int) error {
	// create the table if not exists
	createExpr := `
CREATE TABLE IF NOT EXISTS %s (
    id BIGINT NOT NULL AUTO_INCREMENT,
    raw_bytes BLOB NOT NULL,
    PRIMARY KEY (id)
);
`
	_, err := conn.Exec(fmt.Sprintf(createExpr, table))
	if err != nil {
		return fmt.Errorf("unable to create table: %w", err)
	}

	// insert records
	insertExpr := `
INSERT INTO %s (raw_bytes) VALUES ('%s');
`
	for i := 0; i < recordCnt; i++ {
		value := util.RandStringRunes(dataSize)
		_, err := conn.Exec(fmt.Sprintf(insertExpr, table, value))
		if err != nil {
			return fmt.Errorf("unable to insert record: %w", err)
		}
	}

	// delete records
	deleteExpr := `
DELETE FROM %s;
`
	_, err = conn.Exec(fmt.Sprintf(deleteExpr, table))
	if err != nil {
		return fmt.Errorf("unable to delete record: %w", err)
	}

	return nil
}