package client

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

	"github.com/avast/retry-go/v4"
	"github.com/cloudquery/plugin-sdk/v4/message"
	"github.com/cloudquery/plugin-sdk/v4/schema"
	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgconn"
)

func (c *Client) pgTables(ctx context.Context) (map[string]struct{}, error) {
	c.pgTablesToPKConstraintsMu.RLock()
	// This happens when the CLI was invoked with `sync --no-migrate`
	if len(c.pgTablesToPKConstraints) == 0 {
		c.pgTablesToPKConstraintsMu.RUnlock()
		// listTables populates c.pgTablesToPKConstraints
		_, err := c.listTables(ctx)
		if err != nil {
			return nil, err
		}
		c.pgTablesToPKConstraintsMu.RLock()
	}

	list := make(map[string]struct{}, len(c.pgTablesToPKConstraints))
	for k := range c.pgTablesToPKConstraints {
		list[k] = struct{}{}
	}
	c.pgTablesToPKConstraintsMu.RUnlock()
	return list, nil
}

// InsertBatch inserts records into the destination table. It forms part of the writer.MixedBatchWriter interface.
func (c *Client) InsertBatch(ctx context.Context, messages message.WriteInserts) error {
	pgTables, err := c.pgTables(ctx)
	if err != nil {
		return err
	}

	batch := new(pgx.Batch)

	// Accumulate embeddings per table to insert after each successful base flush
	tableToEmbBatch := make(map[string]*embeddingBatch)

	// Queries cache.
	// We may consider LRU cache in the future, but even for 10K records it may be OK to just save.
	queries := make(map[string]string, 100)

	for _, msg := range messages {
		r := msg.Record
		md := r.Schema().Metadata()
		tableName, ok := md.GetValue(schema.MetadataTableName)
		if !ok {
			return errors.New("table name not found in metadata")
		}

		if _, ok := pgTables[tableName]; !ok {
			return fmt.Errorf("table %s not found", tableName)
		}

		sql, ok := queries[tableName]
		if !ok {
			// cache the query
			table := c.normalizeTable(msg.GetTable())
			if len(table.PrimaryKeysIndexes()) > 0 {
				sql = c.upsert(table)
			} else {
				sql = c.insert(table)
			}
			queries[tableName] = sql
		}

		rows := c.transformValues(r)
		for _, rowVals := range rows {
			batch.Queue(sql, rowVals...)
		}

		// If pgvector is configured and this table has a config, accumulate embeddings work
		if c.hasPgVectorConfig() && c.embeddingsRequester != nil && c.spec.GetPgVectorTableConfig(tableName) != nil {
			err := c.addEmbeddingRows(ctx, tableToEmbBatch, tableName, msg.GetTable(), r, rows)
			if err != nil {
				return err
			}
		}
		if int64(batch.Len()) >= c.batchSize {
			if err := c.flushBatch(ctx, batch); err != nil {
				return err
			}
			// After base rows are flushed, insert/upsert embeddings for accumulated tables
			if err := c.insertEmbeddingsBatch(ctx, tableToEmbBatch); err != nil {
				return err
			}
			// Reset for next chunk
			tableToEmbBatch = make(map[string]*embeddingBatch)
			batch = new(pgx.Batch)
		}
	}

	if err := c.flushBatch(ctx, batch); err != nil {
		return err
	}
	return c.insertEmbeddingsBatch(ctx, tableToEmbBatch)
}

func (c *Client) flushBatch(ctx context.Context, batch *pgx.Batch) error {
	err := retry.Do(func() error {
		if batch.Len() == 0 {
			return nil
		}
		err := c.conn.SendBatch(ctx, batch).Close()
		if err != nil {
			return err
		}

		return nil
	}, retry.RetryIf(func(err error) bool {
		var pgErr *pgconn.PgError
		if errors.As(err, &pgErr) {
			return pgErr.Code == "40P01"
		}

		return false
	}),
		retry.Attempts(uint(c.spec.RetryOnDeadlock)+1),
		retry.LastErrorOnly(true),
	)

	if err != nil {
		var pgErr *pgconn.PgError
		if errors.As(err, &pgErr) {
			return fmt.Errorf("failed to execute batch with pgerror: %s: %w", pgErrToStr(pgErr), err)
		}
		return fmt.Errorf("failed to execute batch: %w", err)
	}

	return nil
}

func (*Client) insert(table *schema.Table) string {
	var sb strings.Builder
	tableName := table.Name
	sb.WriteString("insert into ")
	sb.WriteString(pgx.Identifier{tableName}.Sanitize())
	sb.WriteString(" (")
	columns := table.Columns
	columnsLen := len(columns)
	for i, c := range columns {
		sb.WriteString(pgx.Identifier{c.Name}.Sanitize())
		if i < columnsLen-1 {
			sb.WriteString(",")
		} else {
			sb.WriteString(") values (")
		}
	}
	for i := range columns {
		sb.WriteString(fmt.Sprintf("$%d", i+1))
		if i < columnsLen-1 {
			sb.WriteString(",")
		} else {
			sb.WriteString(")")
		}
	}
	return sb.String()
}

func (c *Client) upsert(table *schema.Table) string {
	if c.pgType == pgTypeCrateDB {
		return c.upsertCrateDB(table)
	}
	var sb strings.Builder
	sb.WriteString(c.insert(table))
	columns := table.Columns
	columnsLen := len(columns)
	constraintName := table.PkConstraintName
	sb.WriteString(" on conflict on constraint ")
	sb.WriteString(pgx.Identifier{constraintName}.Sanitize())
	sb.WriteString(" do update set ")
	for i, column := range columns {
		sb.WriteString(pgx.Identifier{column.Name}.Sanitize())
		sb.WriteString("=excluded.") // excluded references the new values
		sb.WriteString(pgx.Identifier{column.Name}.Sanitize())
		if i < columnsLen-1 {
			sb.WriteString(",")
		} else {
			sb.WriteString("")
		}
	}

	return sb.String()
}

// CrateDB differs from normal Postgres in that it does not support constraint-name
// based upserts, and errors out if a primary key is written to during upsert.
func (c *Client) upsertCrateDB(table *schema.Table) string {
	var sb strings.Builder
	sb.WriteString(c.insert(table))
	columns := table.Columns
	if len(table.PrimaryKeysIndexes()) == len(table.Columns) {
		sb.WriteString(" on conflict do nothing")
		return sb.String()
	}
	pks := table.PrimaryKeys()
	sb.WriteString(fmt.Sprintf(" on conflict (%s) ", strings.Join(pks, ",")))

	sb.WriteString(" do update set ")
	cols := make([]string, 0, len(columns)-len(pks))
	for _, column := range columns {
		if column.PrimaryKey {
			continue
		}
		cols = append(cols, pgx.Identifier{column.Name}.Sanitize()+"=excluded."+pgx.Identifier{column.Name}.Sanitize())
	}
	sb.WriteString(strings.Join(cols, ","))

	return sb.String()
}

func pgErrToStr(err *pgconn.PgError) string {
	var sb strings.Builder
	sb.WriteString("severity: ")
	sb.WriteString(err.Severity)
	sb.WriteString(", code: ")
	sb.WriteString(err.Code)
	sb.WriteString(", message: ")
	sb.WriteString(err.Message)
	sb.WriteString(", detail :")
	sb.WriteString(err.Detail)
	sb.WriteString(", hint: ")
	sb.WriteString(err.Hint)
	sb.WriteString(", position: ")
	sb.WriteString(strconv.FormatInt(int64(err.Position), 10))
	sb.WriteString(", internal_position: ")
	sb.WriteString(strconv.FormatInt(int64(err.InternalPosition), 10))
	sb.WriteString(", internal_query: ")
	sb.WriteString(err.InternalQuery)
	sb.WriteString(", where: ")
	sb.WriteString(err.Where)
	sb.WriteString(", schema_name: ")
	sb.WriteString(err.SchemaName)
	sb.WriteString(", table_name: ")
	sb.WriteString(err.TableName)
	sb.WriteString(", column_name: ")
	sb.WriteString(err.ColumnName)
	sb.WriteString(", data_type_name: ")
	sb.WriteString(err.DataTypeName)
	sb.WriteString(", constraint_name: ")
	sb.WriteString(err.ConstraintName)
	sb.WriteString(", file: ")
	sb.WriteString(err.File)
	sb.WriteString(", line: ")
	sb.WriteString(strconv.FormatUint(uint64(err.Line), 10))
	sb.WriteString(", routine: ")
	sb.WriteString(err.Routine)
	return sb.String()
}
