package persist

import (
	"context"
	"sig-pub/pkg/utils/exit"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"
)

type PGBatchWriter struct {
	pool *pgxpool.Pool
}

func NewPGBatchWriter() *PGBatchWriter {
	return &PGBatchWriter{}
}

func (w *PGBatchWriter) Init(ctx context.Context, connString string) (err error) {
	w.pool, err = pgxpool.New(ctx, connString)
	if err != nil {
		return
	}
	exit.AddHook(w.pool.Close, exit.WithOrderTail())

	return
}

func InsertBatch[T any](w *PGBatchWriter, ctx context.Context, datas []T) (err error) {
	if len(datas) == 0 {
		return
	}

	tableColumns := make(map[string][]string, 3)
	tableValues := make(map[string][][]any, 3)
	for _, data := range datas {
		table, columns, valuesM, e := ReflectGormData(data)
		if e != nil {
			err = e
			return
		}
		if cs, ok := tableColumns[table]; !ok {
			tableColumns[table] = columns
		} else {
			columns = cs
		}
		var values []any
		for _, c := range columns {
			v := valuesM[c]
			values = append(values, v)
		}
		tableValues[table] = append(tableValues[table], values)
	}
	tx, err := w.pool.Begin(ctx)
	if err != nil {
		return
	}
	defer tx.Rollback(ctx)
	for table, values := range tableValues {
		_, err = tx.CopyFrom(ctx, pgx.Identifier{table}, tableColumns[table], pgx.CopyFromRows(values))
		if err != nil {
			return
		}
	}
	err = tx.Commit(ctx)
	return
}