package main

import (
	"context"
	"database/sql/driver"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"os"

	embeddedpostgres "github.com/fergusstrange/embedded-postgres"
	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"
)

type TestData struct {
	Id    string     `db:"id"`
	Value *TestValue `db:"value"`
}

func (e *TestData) RewriteQuery(ctx context.Context, conn *pgx.Conn, sql string, args []any) (newSQL string, newArgs []any, err error) {
	return `
		INSERT INTO
			pgx_test (
				id,
				value
			)
		VALUES
			($1, $2)
	`, []any{e.Id, e.Value}, nil
}

type TestValue struct {
	Data string `json:"data"`
}

func (t *TestValue) Scan(src interface{}) error {
	if src == nil {
		return nil
	}

	var source []byte
	switch val := src.(type) {
	case string:
		source = []byte(val)
	case []byte:
		source = val
	default:
		return errors.New("incompatible source type")
	}

	if err := json.Unmarshal(source, &t); err != nil {
		return fmt.Errorf("unable to unmarshal source data: %w", err)
	}

	return nil
}

func (t *TestValue) Value() (driver.Value, error) {
	if t == nil {
		return nil, nil
	}

	testValueBytes, err := json.Marshal(t)
	if err != nil {
		return nil, fmt.Errorf("unable to marshal test value data: %w", err)
	}

	return testValueBytes, nil
}

func main() {
	ctx := context.TODO()

	logger := slog.New(slog.NewJSONHandler(os.Stdout, nil))
	slog.SetDefault(logger)

	postgres := embeddedpostgres.NewDatabase(embeddedpostgres.DefaultConfig().Port(5555).Logger(nil))
	if err := postgres.Start(); err != nil {
		slog.Error("unable to start embedded postgres", "error", err.Error())
		os.Exit(1)
	}
	defer func() {
		if err := postgres.Stop(); err != nil {
			slog.Error("unable to stop embedded postgres", "error", err.Error())
			os.Exit(1)
		}
	}()

	test(ctx)

}

func test(ctx context.Context) {

	connectionUrl := "postgres://postgres:postgres@127.0.0.1:5555/postgres"

	pgxConfig, err := pgxpool.ParseConfig(connectionUrl)
	if err != nil {
		slog.Error("cannot parse postgres config", "error", err.Error())
		return
	}

	pool, err := pgxpool.NewWithConfig(ctx, pgxConfig)
	if err != nil {
		slog.Error("unable to create pgx connection pool", "error", err.Error())
		return
	}

	if _, err := pool.Exec(ctx, `CREATE TABLE "public"."pgx_test" ("id" text NOT NULL PRIMARY KEY, "value" JSONB);`); err != nil {
		slog.Error("unable to create table 'pgx_test'", "error", err.Error())
		return
	}

	first := &TestData{
		Id: "first",
		Value: &TestValue{
			Data: "test",
		},
	}

	if _, err := pool.Exec(ctx, "", first); err != nil {
		slog.Error("unable to insert first value into table 'pgx_test'", "error", err.Error())
		return
	}

	second := &TestData{
		Id:    "second",
		Value: nil,
	}

	if _, err := pool.Exec(ctx, "", second); err != nil {
		slog.Error("unable to insert second value into table 'pgx_test'", "error", err.Error())
		return
	}

	rows, err := pool.Query(ctx, "SELECT id, value from pgx_test")
	if err != nil {
		slog.Error("unable to query rows from table 'pgx_test'", "error", err.Error())
		return
	}

	testData, err := pgx.CollectRows(rows, pgx.RowToAddrOfStructByName[TestData])
	if err != nil {
		slog.Error("unable to collect rows from table 'pgx_test'", "error", err.Error())
		return
	}

	if testData[0].Value == nil || testData[0].Value.Data != "test" {
		slog.Error("expected testData[0].Value.Data to be 'test'")
		return
	}
	slog.Info("testData[0].Value.Data was 'test'")

	if testData[1].Value != nil {
		slog.Error("expected testData[1].Value to be nil struct")
		return
	}

	slog.Info("testData[1].Value was nil struct")
}