package database

import (
	"context"
	"errors"
	"fmt"
	"queryorchestration/internal/database/repository"

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

type Pool interface {
	repository.DBTX
	Begin(ctx context.Context) (pgx.Tx, error)
	Ping(ctx context.Context) error
}

func (b *DBConfig) GetDBPool() Pool {
	return b.DBPool
}

func (b *DBConfig) GetDBQueries() *repository.Queries {
	return b.DBQueries
}

func (b *DBConfig) SetDBPoolConfig() error {
	config, err := pgxpool.ParseConfig(b.GetDBURI())
	if err != nil {
		return err
	}

	b.DBEnumTypes = []string{
		"querytype",
	}

	config.AfterConnect = func(ctx context.Context, conn *pgx.Conn) error {
		for _, enumName := range b.DBEnumTypes {
			dt, err := conn.LoadType(ctx, enumName)
			if err != nil {
				return fmt.Errorf("failed to load suitable enum type: %w", err)
			}
			conn.TypeMap().RegisterType(dt)
		}
		return nil
	}

	b.DBPoolConfig = config

	return nil
}

func (b *DBConfig) SetDBPool(ctx context.Context) error {
	err := b.SetDBPoolConfig()
	if err != nil {
		return err
	}

	pool, err := pgxpool.NewWithConfig(ctx, b.DBPoolConfig)
	if err != nil {
		return err
	}

	err = pool.Ping(ctx)
	if err != nil {
		return errors.New("unable to ping database")
	}

	b.DBPool = pool
	b.DBQueries = repository.New(b.DBPool)

	return nil
}