// Patial Tech.
// Author, Ankit Patial

package pgm

import (
	"context"
	"errors"
	"strconv"
	"strings"

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

type (
	Update interface {
		Set(field Field, val any) UpdateClause
		SetMap(fields map[Field]any) UpdateClause
	}

	UpdateClause interface {
		Update
		Where(cond ...Conditioner) WhereOrExec
	}

	WhereOrExec interface {
		Where(cond ...Conditioner) WhereOrExec
		Execute
	}

	updateQry struct {
		table     string
		cols      []string
		condition []Conditioner
		args      []any
		debug     bool
	}
)

func (q *updateQry) Set(field Field, val any) UpdateClause {
	col := field.Name()
	q.cols = append(q.cols, col+"=$"+strconv.Itoa(len(q.args)+1))
	q.args = append(q.args, val)
	return q
}

func (q *updateQry) SetMap(cols map[Field]any) UpdateClause {
	for k, v := range cols {
		q.Set(k, v)
	}
	return q
}

func (q *updateQry) Where(cond ...Conditioner) WhereOrExec {
	q.condition = append(q.condition, cond...)
	return q
}

func (q *updateQry) Exec(ctx context.Context) error {
	if len(q.cols) == 0 {
		return errors.New("update query has no columns to update: call Set() before Exec()")
	}
	_, err := poolPGX.Load().Exec(ctx, q.String(), q.args...)
	if err != nil {
		return err
	}

	return nil
}

func (q *updateQry) ExecTx(ctx context.Context, tx pgx.Tx) error {
	if len(q.cols) == 0 {
		return errors.New("update query has no columns to update: call Set() before ExecTx()")
	}
	_, err := tx.Exec(ctx, q.String(), q.args...)
	if err != nil {
		return err
	}

	return nil
}

func (q *updateQry) String() string {
	sb := getSB()
	defer putSB(sb)

	n := 7 + len(q.table) + 5 // "UPDATE q.table SET
	for _, col := range q.cols {
		n += len(col) + 5
	}
	if len(q.condition) > 0 {
		n += 7 + len(q.condition)*5 // WHERE with condition, 5 is just avg small min val
	}
	sb.Grow(n)

	// UPDATE
	sb.WriteString("UPDATE " + q.table + " SET ")
	sb.WriteString(strings.Join(q.cols, ", "))

	// WHERE
	if len(q.condition) > 0 {
		sb.WriteString(" WHERE ")
		var argIdx int
		for i, c := range q.condition {
			argIdx = len(q.args)
			if i > 0 {
				sb.WriteString(" AND ")
			}
			sb.WriteString(c.Condition(&q.args, argIdx))
		}
	}

	return sb.String()
}
