package backup

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"strings"

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

// Contents is a backup's data: the schema version that wrote it, and each
// table's rows as JSON objects with their keys sorted, in primary-key order.
// An unchanged row is the same bytes in every backup.
type Contents struct {
	Schema int64
	Tables map[string][][]byte
}

// Snapshot reads every table in Tables as of one moment.
func Snapshot(ctx context.Context, pool *pgxpool.Pool) (Contents, error) {
	c := Contents{Tables: make(map[string][][]byte, len(Tables))}
	tx, err := pool.BeginTx(ctx, pgx.TxOptions{IsoLevel: pgx.RepeatableRead, AccessMode: pgx.ReadOnly})
	if err != nil {
		return c, fmt.Errorf("backup: begin: %w", err)
	}
	defer func() { _ = tx.Rollback(ctx) }()

	if c.Schema, err = schemaVersion(ctx, tx); err != nil {
		return c, err
	}
	for _, t := range Tables {
		lines, err := readTable(ctx, tx, t)
		if err != nil {
			return c, err
		}
		c.Tables[t.Name] = lines
	}
	return c, nil
}

func readTable(ctx context.Context, tx pgx.Tx, t Table) ([][]byte, error) {
	pk, err := primaryKey(ctx, tx, t.Name)
	if err != nil {
		return nil, err
	}
	order := make([]string, len(pk))
	for i, col := range pk {
		order[i] = "t." + pgx.Identifier{col}.Sanitize()
	}
	omit := t.Omit
	if omit == nil {
		omit = []string{}
	}
	rows, err := tx.Query(ctx, fmt.Sprintf("SELECT to_jsonb(t) - $1::text[] FROM %s t ORDER BY %s",
		pgx.Identifier{t.Name}.Sanitize(), strings.Join(order, ", ")), omit)
	if err != nil {
		return nil, fmt.Errorf("backup: read %s: %w", t.Name, err)
	}
	lines, err := pgx.CollectRows(rows, func(row pgx.CollectableRow) ([]byte, error) {
		var raw []byte
		if err := row.Scan(&raw); err != nil {
			return nil, err
		}
		return canonical(raw)
	})
	if err != nil {
		return nil, fmt.Errorf("backup: read %s: %w", t.Name, err)
	}
	return lines, nil
}

// canonical re-encodes a row with its keys sorted. Numbers pass through as
// written, so a bigint or a float stays exact.
func canonical(raw []byte) ([]byte, error) {
	dec := json.NewDecoder(bytes.NewReader(raw))
	dec.UseNumber()
	var row map[string]any
	if err := dec.Decode(&row); err != nil {
		return nil, err
	}
	return json.Marshal(row)
}

// primaryKey returns a table's primary-key columns in key order.
func primaryKey(ctx context.Context, tx pgx.Tx, table string) ([]string, error) {
	rows, err := tx.Query(ctx, `
		SELECT a.attname
		  FROM pg_index i
		  JOIN pg_attribute a ON a.attrelid = i.indrelid AND a.attnum = ANY (i.indkey)
		 WHERE i.indrelid = to_regclass($1) AND i.indisprimary
		 ORDER BY array_position(i.indkey, a.attnum)`, table)
	if err != nil {
		return nil, fmt.Errorf("backup: read %s's primary key: %w", table, err)
	}
	cols, err := pgx.CollectRows(rows, pgx.RowTo[string])
	if err != nil {
		return nil, fmt.Errorf("backup: read %s's primary key: %w", table, err)
	}
	if len(cols) == 0 {
		return nil, fmt.Errorf("backup: %s has no primary key to order its rows by", table)
	}
	return cols, nil
}

func schemaVersion(ctx context.Context, tx pgx.Tx) (int64, error) {
	var v int64
	if err := tx.QueryRow(ctx, "SELECT version FROM schema_migrations").Scan(&v); err != nil {
		return 0, fmt.Errorf("backup: read the schema version: %w", err)
	}
	return v, nil
}