// Run it twice: the 1st creates the in-memory database and writes the snapshot
// to disk; the 2nd reads the snapshot back into RAM and lists.
package main

import (
	"context"
	"database/sql"
	"errors"
	"fmt"
	"os"
	"text/tabwriter"

	_ "modernc.org/sqlite"
)

const snapFile = "db.snap"

type serializer interface {
	Serialize() ([]byte, error)
	Deserialize([]byte) error
}

func mustExec(db *sql.DB, q string, args ...any) {
	if _, err := db.Exec(q, args...); err != nil {
		fatal(err)
	}
}

func fileExists(path string) bool {
	_, err := os.Stat(path)
	return err == nil
}

func fatal(err error) {
	fmt.Fprintln(os.Stderr, "error:", err)
	os.Exit(1)
}

// snapshot returns the raw bytes of the in-memory database — exactly what
// would be in the file if it were saved to disk.
func snapshot(db *sql.DB) ([]byte, error) {
	conn, err := db.Conn(context.Background())
	if err != nil {
		return nil, err
	}
	defer func() { _ = conn.Close() }()
	var buf []byte
	err = conn.Raw(func(dc any) error {
		s, ok := dc.(serializer)
		if !ok {
			return errors.New("the driver does not expose Serialize")
		}
		var serr error
		buf, serr = s.Serialize()
		return serr
	})
	return buf, err
}

// restore reloads a snapshot into the in-memory database. After this, every
// table and row is back in RAM — no INSERT, no rebuild.
func restore(db *sql.DB, buf []byte) error {
	conn, err := db.Conn(context.Background())
	if err != nil {
		return err
	}
	defer func() { _ = conn.Close() }()
	return conn.Raw(func(dc any) error {
		s, ok := dc.(serializer)
		if !ok {
			return errors.New("the driver does not expose Deserialize")
		}
		return s.Deserialize(buf)
	})
}

func list(db *sql.DB) {
	rows, err := db.Query(`SELECT id, text FROM notes ORDER BY id`)
	if err != nil {
		fatal(err)
	}
	defer func() { _ = rows.Close() }()
	w := tabwriter.NewWriter(os.Stdout, 0, 2, 2, ' ', 0)
	fmt.Fprintln(w, "ID\tTEXT")
	for rows.Next() {
		var id int
		var text string
		if err := rows.Scan(&id, &text); err != nil {
			fatal(err)
		}
		fmt.Fprintf(w, "%d\t%s\n", id, text)
	}
	_ = w.Flush()
	if err := rows.Err(); err != nil {
		fatal(err)
	}
}

func main() {
	// Database ONLY in RAM. MaxOpenConns(1) is essential: each new connection
	// to a ":memory:" DB would open a fresh, EMPTY one — so we pin everything
	// to a single connection for the life of the process.
	db, err := sql.Open("sqlite", ":memory:")
	if err != nil {
		fatal(err)
	}
	defer func() { _ = db.Close() }()
	db.SetMaxOpenConns(1)

	if fileExists(snapFile) {
		blob, err := os.ReadFile(snapFile)
		if err != nil {
			fatal(err)
		}
		if err := restore(db, blob); err != nil {
			fatal(err)
		}
		fmt.Println("snapshot loaded into RAM:")
		list(db)
		return
	}

	fmt.Println("new database — creating and inserting…")
	mustExec(db, `CREATE TABLE notes (id INTEGER PRIMARY KEY, text TEXT);`)
	mustExec(db, `INSERT INTO notes (text) VALUES (?),(?),(?)`,
		"the database lives in memory",
		"the snapshot is the SQLite bytes themselves",
		"run it again to read it back",
	)
	list(db)

	blob, err := snapshot(db)
	if err != nil {
		fatal(err)
	}
	if err := os.WriteFile(snapFile, blob, 0o600); err != nil {
		fatal(err)
	}
	fmt.Printf("\nsnapshot saved to %s (%d bytes). run it again.\n", snapFile, len(blob))
}