package db

import (
	"context"
	"database/sql"
	"errors"
)

func (db *DB) KVSet(ctx context.Context, key string, value any) error {
	// Upsert into SQLite
	query := `
        INSERT INTO kv_store (key, value, updated_at) 
        VALUES (?, ?, CURRENT_TIMESTAMP)
        ON CONFLICT(key) DO UPDATE SET 
            value = excluded.value, 
            updated_at = CURRENT_TIMESTAMP;`

	_, err := db.write.ExecContext(ctx, query, key, value)
	return err
}

func (db *DB) KVGet(ctx context.Context, key string) (any, error) {
	var rawVal any
	err := db.read.QueryRowContext(ctx, "SELECT value FROM kv_store WHERE key = ?", key).Scan(&rawVal)
	return rawVal, err
}

func (db *DB) KVList(ctx context.Context) (map[string]any, error) {
	rows, err := db.read.QueryContext(ctx, "SELECT key, value FROM kv_store")
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	result := make(map[string]any)
	for rows.Next() {
		var key string
		var value any
		if err := rows.Scan(&key, &value); err != nil {
			return nil, err
		}
		result[key] = value
	}
	return result, rows.Err()
}

func (db *DB) KVInc(ctx context.Context, key string, delta int64) (int64, error) {
	tx, err := db.GetWriteTx(ctx)
	if err != nil {
		return 0, err
	}

	var val any
	err = tx.QueryRowContext(ctx, "SELECT value FROM kv_store WHERE key = ?", key).Scan(&val)
	if err != nil && !errors.Is(err, sql.ErrNoRows) {
		_ = tx.Rollback()
		return 0, err
	}

	if errors.Is(err, sql.ErrNoRows) {
		_, err = tx.ExecContext(ctx, "INSERT INTO kv_store (key, value, updated_at) VALUES (?, ?, CURRENT_TIMESTAMP)", key, delta)
		if err != nil {
			_ = tx.Rollback()
			return 0, err
		}
		return delta, tx.Commit()
	}

	if _, ok := val.(int64); !ok {
		_ = tx.Rollback()
		return 0, errors.New("KVInc: value is not an integer")
	}

	query := `
INSERT INTO kv_store (key, value, updated_at)
VALUES (?, ?, CURRENT_TIMESTAMP)
ON CONFLICT(key) DO UPDATE SET
	value = CAST(kv_store.value AS INTEGER) + excluded.value,
	updated_at = CURRENT_TIMESTAMP
RETURNING value;`

	var newCount int64
	err = tx.QueryRowContext(ctx, query, key, delta).Scan(&newCount)
	if err != nil {
		_ = tx.Rollback()
		return 0, err
	}

	return newCount, tx.Commit()
}

func (db *DB) KVPop(ctx context.Context, key string) (any, error) {
	var rawVal any
	query := `DELETE FROM kv_store WHERE key = ? RETURNING kv_store.value;`
	err := db.write.QueryRowContext(ctx, query, key).Scan(&rawVal)
	return rawVal, err
}