banger/internal/store/store.go

package store

import (
	"context"
	"database/sql"
	"encoding/json"
	"errors"
	"fmt"
	"net/url"
	"path/filepath"
	"sync"
	"time"

	_ "modernc.org/sqlite"

	"banger/internal/model"
)

type Store struct {
	db      *sql.DB
	writeMu sync.Mutex
}

func Open(path string) (*Store, error) {
	dsn, err := sqliteDSN(path)
	if err != nil {
		return nil, err
	}
	db, err := sql.Open("sqlite", dsn)
	if err != nil {
		return nil, err
	}
	store := &Store{db: db}
	if err := store.migrate(); err != nil {
		_ = db.Close()
		return nil, err
	}
	return store, nil
}

func (s *Store) Close() error {
	return s.db.Close()
}

func sqliteDSN(path string) (string, error) {
	absPath, err := filepath.Abs(path)
	if err != nil {
		return "", fmt.Errorf("resolve sqlite path: %w", err)
	}
	query := url.Values{}
	for _, pragma := range []string{
		"journal_mode(WAL)",
		"synchronous(NORMAL)",
		"foreign_keys(1)",
		"busy_timeout(5000)",
		"temp_store(MEMORY)",
		"wal_autocheckpoint(1000)",
		"journal_size_limit(67108864)",
	} {
		query.Add("_pragma", pragma)
	}
	return (&url.URL{
		Scheme:   "file",
		Path:     filepath.ToSlash(absPath),
		RawQuery: query.Encode(),
	}).String(), nil
}

func (s *Store) migrate() error {
	stmts := []string{
		`CREATE TABLE IF NOT EXISTS images (
			id TEXT PRIMARY KEY,
			name TEXT NOT NULL UNIQUE,
			managed INTEGER NOT NULL DEFAULT 0,
			artifact_dir TEXT,
			rootfs_path TEXT NOT NULL,
			work_seed_path TEXT,
			kernel_path TEXT NOT NULL,
			initrd_path TEXT,
			modules_dir TEXT,
			packages_path TEXT,
			build_size TEXT,
			seeded_ssh_public_key_fingerprint TEXT,
			docker INTEGER NOT NULL DEFAULT 0,
			created_at TEXT NOT NULL,
			updated_at TEXT NOT NULL
		);`,
		`CREATE TABLE IF NOT EXISTS vms (
			id TEXT PRIMARY KEY,
			name TEXT NOT NULL UNIQUE,
			image_id TEXT NOT NULL,
			guest_ip TEXT NOT NULL UNIQUE,
			state TEXT NOT NULL,
			created_at TEXT NOT NULL,
			updated_at TEXT NOT NULL,
			last_touched_at TEXT NOT NULL,
			spec_json TEXT NOT NULL,
			runtime_json TEXT NOT NULL,
			stats_json TEXT NOT NULL DEFAULT '{}',
			FOREIGN KEY(image_id) REFERENCES images(id) ON DELETE RESTRICT
		);`,
	}
	for _, stmt := range stmts {
		if _, err := s.db.Exec(stmt); err != nil {
			return err
		}
	}
	if err := ensureColumnExists(s.db, "images", "work_seed_path", "TEXT"); err != nil {
		return err
	}
	if err := ensureColumnExists(s.db, "images", "seeded_ssh_public_key_fingerprint", "TEXT"); err != nil {
		return err
	}
	return nil
}

func (s *Store) UpsertImage(ctx context.Context, image model.Image) error {
	s.writeMu.Lock()
	defer s.writeMu.Unlock()
	const query = `
	INSERT INTO images (
		id, name, managed, artifact_dir, rootfs_path, work_seed_path, kernel_path, initrd_path,
		modules_dir, build_size, seeded_ssh_public_key_fingerprint, docker, created_at, updated_at
	) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
	ON CONFLICT(id) DO UPDATE SET
		name=excluded.name,
		managed=excluded.managed,
		artifact_dir=excluded.artifact_dir,
		rootfs_path=excluded.rootfs_path,
		work_seed_path=excluded.work_seed_path,
		kernel_path=excluded.kernel_path,
		initrd_path=excluded.initrd_path,
		modules_dir=excluded.modules_dir,
		build_size=excluded.build_size,
		seeded_ssh_public_key_fingerprint=excluded.seeded_ssh_public_key_fingerprint,
		docker=excluded.docker,
		updated_at=excluded.updated_at`
	_, err := s.db.ExecContext(ctx, query,
		image.ID,
		image.Name,
		boolToInt(image.Managed),
		image.ArtifactDir,
		image.RootfsPath,
		image.WorkSeedPath,
		image.KernelPath,
		image.InitrdPath,
		image.ModulesDir,
		image.BuildSize,
		image.SeededSSHPublicKeyFingerprint,
		boolToInt(image.Docker),
		image.CreatedAt.Format(time.RFC3339),
		image.UpdatedAt.Format(time.RFC3339),
	)
	return err
}

func (s *Store) GetImageByName(ctx context.Context, name string) (model.Image, error) {
	return s.getImage(ctx, "SELECT id, name, managed, artifact_dir, rootfs_path, work_seed_path, kernel_path, initrd_path, modules_dir, build_size, seeded_ssh_public_key_fingerprint, docker, created_at, updated_at FROM images WHERE name = ?", name)
}

func (s *Store) GetImageByID(ctx context.Context, id string) (model.Image, error) {
	return s.getImage(ctx, "SELECT id, name, managed, artifact_dir, rootfs_path, work_seed_path, kernel_path, initrd_path, modules_dir, build_size, seeded_ssh_public_key_fingerprint, docker, created_at, updated_at FROM images WHERE id = ?", id)
}

func (s *Store) ListImages(ctx context.Context) ([]model.Image, error) {
	rows, err := s.db.QueryContext(ctx, "SELECT id, name, managed, artifact_dir, rootfs_path, work_seed_path, kernel_path, initrd_path, modules_dir, build_size, seeded_ssh_public_key_fingerprint, docker, created_at, updated_at FROM images ORDER BY created_at ASC")
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var images []model.Image
	for rows.Next() {
		image, err := scanImage(rows)
		if err != nil {
			return nil, err
		}
		images = append(images, image)
	}
	return images, rows.Err()
}

func (s *Store) DeleteImage(ctx context.Context, id string) error {
	s.writeMu.Lock()
	defer s.writeMu.Unlock()
	_, err := s.db.ExecContext(ctx, "DELETE FROM images WHERE id = ?", id)
	return err
}

func (s *Store) UpsertVM(ctx context.Context, vm model.VMRecord) error {
	s.writeMu.Lock()
	defer s.writeMu.Unlock()
	specJSON, err := json.Marshal(vm.Spec)
	if err != nil {
		return err
	}
	runtimeJSON, err := json.Marshal(vm.Runtime)
	if err != nil {
		return err
	}
	statsJSON, err := json.Marshal(vm.Stats)
	if err != nil {
		return err
	}
	const query = `
	INSERT INTO vms (
		id, name, image_id, guest_ip, state, created_at, updated_at, last_touched_at,
		spec_json, runtime_json, stats_json
	) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
	ON CONFLICT(id) DO UPDATE SET
		name=excluded.name,
		image_id=excluded.image_id,
		guest_ip=excluded.guest_ip,
		state=excluded.state,
		updated_at=excluded.updated_at,
		last_touched_at=excluded.last_touched_at,
		spec_json=excluded.spec_json,
		runtime_json=excluded.runtime_json,
		stats_json=excluded.stats_json`
	_, err = s.db.ExecContext(ctx, query,
		vm.ID,
		vm.Name,
		vm.ImageID,
		vm.Runtime.GuestIP,
		string(vm.State),
		vm.CreatedAt.Format(time.RFC3339),
		vm.UpdatedAt.Format(time.RFC3339),
		vm.LastTouchedAt.Format(time.RFC3339),
		string(specJSON),
		string(runtimeJSON),
		string(statsJSON),
	)
	return err
}

func (s *Store) GetVM(ctx context.Context, idOrName string) (model.VMRecord, error) {
	const query = `
	SELECT id, name, image_id, guest_ip, state, created_at, updated_at, last_touched_at,
	       spec_json, runtime_json, stats_json
	FROM vms
	WHERE id = ? OR name = ?
	`
	row := s.db.QueryRowContext(ctx, query, idOrName, idOrName)
	return scanVMRow(row)
}

func (s *Store) GetVMByID(ctx context.Context, id string) (model.VMRecord, error) {
	row := s.db.QueryRowContext(ctx, `
	SELECT id, name, image_id, guest_ip, state, created_at, updated_at, last_touched_at,
	       spec_json, runtime_json, stats_json
	FROM vms WHERE id = ?`, id)
	return scanVMRow(row)
}

func (s *Store) ListVMs(ctx context.Context) ([]model.VMRecord, error) {
	rows, err := s.db.QueryContext(ctx, `
	SELECT id, name, image_id, guest_ip, state, created_at, updated_at, last_touched_at,
	       spec_json, runtime_json, stats_json
	FROM vms ORDER BY created_at ASC`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var vms []model.VMRecord
	for rows.Next() {
		vm, err := scanVMRows(rows)
		if err != nil {
			return nil, err
		}
		vms = append(vms, vm)
	}
	return vms, rows.Err()
}

func (s *Store) DeleteVM(ctx context.Context, id string) error {
	s.writeMu.Lock()
	defer s.writeMu.Unlock()
	_, err := s.db.ExecContext(ctx, "DELETE FROM vms WHERE id = ?", id)
	return err
}

func (s *Store) FindVMsUsingImage(ctx context.Context, imageID string) ([]model.VMRecord, error) {
	rows, err := s.db.QueryContext(ctx, `
	SELECT id, name, image_id, guest_ip, state, created_at, updated_at, last_touched_at,
	       spec_json, runtime_json, stats_json
	FROM vms WHERE image_id = ?`, imageID)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var vms []model.VMRecord
	for rows.Next() {
		vm, err := scanVMRows(rows)
		if err != nil {
			return nil, err
		}
		vms = append(vms, vm)
	}
	return vms, rows.Err()
}

func (s *Store) NextGuestIP(ctx context.Context, bridgeIPPrefix string) (string, error) {
	used := map[string]struct{}{}
	rows, err := s.db.QueryContext(ctx, "SELECT guest_ip FROM vms")
	if err != nil {
		return "", err
	}
	defer rows.Close()
	for rows.Next() {
		var ip string
		if err := rows.Scan(&ip); err != nil {
			return "", err
		}
		used[ip] = struct{}{}
	}
	if err := rows.Err(); err != nil {
		return "", err
	}
	for i := 2; i < 255; i++ {
		candidate := fmt.Sprintf("%s.%d", bridgeIPPrefix, i)
		if _, exists := used[candidate]; !exists {
			return candidate, nil
		}
	}
	return "", errors.New("no guest IPs available")
}

func (s *Store) getImage(ctx context.Context, query string, arg string) (model.Image, error) {
	row := s.db.QueryRowContext(ctx, query, arg)
	return scanImageRow(row)
}

func scanImage(rows scanner) (model.Image, error) {
	return scanImageRow(rows)
}

type scanner interface {
	Scan(dest ...any) error
}

func scanImageRow(row scanner) (model.Image, error) {
	var image model.Image
	var managed, docker int
	var workSeedPath sql.NullString
	var seededSSHPublicKeyFingerprint sql.NullString
	var createdAt, updatedAt string
	err := row.Scan(
		&image.ID,
		&image.Name,
		&managed,
		&image.ArtifactDir,
		&image.RootfsPath,
		&workSeedPath,
		&image.KernelPath,
		&image.InitrdPath,
		&image.ModulesDir,
		&image.BuildSize,
		&seededSSHPublicKeyFingerprint,
		&docker,
		&createdAt,
		&updatedAt,
	)
	if err != nil {
		return image, err
	}
	image.Managed = managed == 1
	image.Docker = docker == 1
	image.WorkSeedPath = workSeedPath.String
	image.SeededSSHPublicKeyFingerprint = seededSSHPublicKeyFingerprint.String
	image.CreatedAt, err = time.Parse(time.RFC3339, createdAt)
	if err != nil {
		return image, err
	}
	image.UpdatedAt, err = time.Parse(time.RFC3339, updatedAt)
	if err != nil {
		return image, err
	}
	return image, nil
}

func scanVMRow(row scanner) (model.VMRecord, error) {
	return scanVMInto(row)
}

func scanVMRows(rows scanner) (model.VMRecord, error) {
	return scanVMInto(rows)
}

func scanVMInto(row scanner) (model.VMRecord, error) {
	var vm model.VMRecord
	var state, createdAt, updatedAt, touchedAt, specJSON, runtimeJSON, statsJSON string
	err := row.Scan(
		&vm.ID,
		&vm.Name,
		&vm.ImageID,
		&vm.Runtime.GuestIP,
		&state,
		&createdAt,
		&updatedAt,
		&touchedAt,
		&specJSON,
		&runtimeJSON,
		&statsJSON,
	)
	if err != nil {
		return vm, err
	}
	vm.State = model.VMState(state)
	if err := json.Unmarshal([]byte(specJSON), &vm.Spec); err != nil {
		return vm, err
	}
	if err := json.Unmarshal([]byte(runtimeJSON), &vm.Runtime); err != nil {
		return vm, err
	}
	if statsJSON != "" {
		if err := json.Unmarshal([]byte(statsJSON), &vm.Stats); err != nil {
			return vm, err
		}
	}
	var parseErr error
	vm.CreatedAt, parseErr = time.Parse(time.RFC3339, createdAt)
	if parseErr != nil {
		return vm, parseErr
	}
	vm.UpdatedAt, parseErr = time.Parse(time.RFC3339, updatedAt)
	if parseErr != nil {
		return vm, parseErr
	}
	vm.LastTouchedAt, parseErr = time.Parse(time.RFC3339, touchedAt)
	if parseErr != nil {
		return vm, parseErr
	}
	return vm, nil
}

func ensureColumnExists(db *sql.DB, table, column, columnType string) error {
	rows, err := db.Query(fmt.Sprintf("PRAGMA table_info(%s)", table))
	if err != nil {
		return err
	}
	defer rows.Close()
	for rows.Next() {
		var (
			cid       int
			name      string
			valueType string
			notNull   int
			defaultV  sql.NullString
			pk        int
		)
		if err := rows.Scan(&cid, &name, &valueType, &notNull, &defaultV, &pk); err != nil {
			return err
		}
		if name == column {
			return nil
		}
	}
	if err := rows.Err(); err != nil {
		return err
	}
	_, err = db.Exec(fmt.Sprintf("ALTER TABLE %s ADD COLUMN %s %s", table, column, columnType))
	return err
}

func boolToInt(value bool) int {
	if value {
		return 1
	}
	return 0
}

func nullableTimeString(value time.Time) any {
	if value.IsZero() {
		return nil
	}
	return value.Format(time.RFC3339)
}

func nullableInt(value *int) any {
	if value == nil {
		return nil
	}
	return *value
}