backend/internal/store/postgres.go

package store

import (
	"database/sql"
	"encoding/json"
	"errors"
	"strings"
	"time"

	_ "github.com/jackc/pgx/v5/stdlib"
)

type PostgresStore struct {
	db *sql.DB
}

func NewPostgresStore(databaseURL string) (*PostgresStore, error) {
	db, err := sql.Open("pgx", databaseURL)
	if err != nil {
		return nil, err
	}
	if err := db.Ping(); err != nil {
		return nil, err
	}
	return &PostgresStore{db: db}, nil
}

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

func (s *PostgresStore) UpsertUser(user User) {
	_, _ = s.db.Exec(
		`INSERT INTO users (public_key, inviter, created_at) VALUES ($1, $2, $3)
		 ON CONFLICT (public_key) DO UPDATE SET inviter = EXCLUDED.inviter`,
		user.PublicKey, nullStr(user.Inviter), user.CreatedAt,
	)
}

func (s *PostgresStore) GetUser(publicKey string) (User, error) {
	var u User
	err := s.db.QueryRow(
		`SELECT public_key, COALESCE(inviter,''), created_at FROM users WHERE public_key = $1`,
		publicKey,
	).Scan(&u.PublicKey, &u.Inviter, &u.CreatedAt)
	if errors.Is(err, sql.ErrNoRows) {
		return User{}, ErrNotFound
	}
	if err != nil {
		return User{}, err
	}
	return u, nil
}

func (s *PostgresStore) CreateChallenge(ch Challenge) error {
	_, err := s.db.Exec(
		`INSERT INTO challenges (nonce, public_key, expires_at, used) VALUES ($1, $2, $3, $4)`,
		ch.Nonce, ch.PublicKey, ch.ExpiresAt, ch.Used,
	)
	if err != nil {
		if isUniqueViolation(err) {
			return ErrAlreadyExists
		}
		return err
	}
	return nil
}

func (s *PostgresStore) GetChallenge(nonce string) (Challenge, error) {
	var ch Challenge
	err := s.db.QueryRow(
		`SELECT nonce, public_key, expires_at, used FROM challenges WHERE nonce = $1`,
		nonce,
	).Scan(&ch.Nonce, &ch.PublicKey, &ch.ExpiresAt, &ch.Used)
	if errors.Is(err, sql.ErrNoRows) {
		return Challenge{}, ErrNotFound
	}
	if err != nil {
		return Challenge{}, err
	}
	return ch, nil
}

func (s *PostgresStore) MarkChallengeUsed(nonce string) error {
	res, err := s.db.Exec(`UPDATE challenges SET used = true WHERE nonce = $1`, nonce)
	if err != nil {
		return err
	}
	n, _ := res.RowsAffected()
	if n == 0 {
		return ErrNotFound
	}
	return nil
}

func (s *PostgresStore) SaveSession(session Session) {
	_, _ = s.db.Exec(
		`INSERT INTO sessions (token, public_key, expires_at) VALUES ($1, $2, $3)
		 ON CONFLICT (token) DO UPDATE SET public_key = EXCLUDED.public_key, expires_at = EXCLUDED.expires_at`,
		session.Token, session.PublicKey, session.ExpiresAt,
	)
}

func (s *PostgresStore) GetSession(token string) (Session, error) {
	var sess Session
	err := s.db.QueryRow(
		`SELECT token, public_key, expires_at FROM sessions WHERE token = $1`,
		token,
	).Scan(&sess.Token, &sess.PublicKey, &sess.ExpiresAt)
	if errors.Is(err, sql.ErrNoRows) {
		return Session{}, ErrNotFound
	}
	if err != nil {
		return Session{}, err
	}
	return sess, nil
}

func (s *PostgresStore) SaveInvitation(inv Invitation) error {
	_, err := s.db.Exec(
		`INSERT INTO invitations (jti, inviter_public_key, invitee_public_key, expires_at, max_uses, used_count)
		 VALUES ($1, $2, $3, $4, $5, $6)`,
		inv.JTI, inv.InviterPublicKey, nullStr(inv.InviteePublicKey), inv.ExpiresAt, inv.MaxUses, inv.UsedCount,
	)
	if err != nil {
		if isUniqueViolation(err) {
			return ErrAlreadyExists
		}
		return err
	}
	return nil
}

func (s *PostgresStore) GetInvitation(jti string) (Invitation, error) {
	var inv Invitation
	var invitee sql.NullString
	err := s.db.QueryRow(
		`SELECT jti, inviter_public_key, invitee_public_key, expires_at, max_uses, used_count
		 FROM invitations WHERE jti = $1`,
		jti,
	).Scan(&inv.JTI, &inv.InviterPublicKey, &invitee, &inv.ExpiresAt, &inv.MaxUses, &inv.UsedCount)
	if errors.Is(err, sql.ErrNoRows) {
		return Invitation{}, ErrNotFound
	}
	if err != nil {
		return Invitation{}, err
	}
	inv.InviteePublicKey = invitee.String
	return inv, nil
}

func (s *PostgresStore) IncrementInvitationUsage(jti string) error {
	res, err := s.db.Exec(`UPDATE invitations SET used_count = used_count + 1 WHERE jti = $1`, jti)
	if err != nil {
		return err
	}
	n, _ := res.RowsAffected()
	if n == 0 {
		return ErrNotFound
	}
	return nil
}

func (s *PostgresStore) SaveCollection(c Collection) {
	_, _ = s.db.Exec(
		`INSERT INTO collections (id, owner_key, name, created_at) VALUES ($1, $2, $3, $4)
		 ON CONFLICT (id) DO UPDATE SET name = EXCLUDED.name`,
		c.ID, c.OwnerKey, c.Name, c.CreatedAt,
	)
}

func (s *PostgresStore) ListCollectionsByOwner(owner string) []Collection {
	rows, err := s.db.Query(
		`SELECT id, owner_key, name, created_at FROM collections WHERE owner_key = $1 ORDER BY created_at`,
		owner,
	)
	if err != nil {
		return nil
	}
	defer rows.Close()
	var result []Collection
	for rows.Next() {
		var c Collection
		if err := rows.Scan(&c.ID, &c.OwnerKey, &c.Name, &c.CreatedAt); err != nil {
			return result
		}
		result = append(result, c)
	}
	return result
}

func (s *PostgresStore) GetCollection(id string) (Collection, error) {
	var c Collection
	err := s.db.QueryRow(
		`SELECT id, owner_key, name, created_at FROM collections WHERE id = $1`,
		id,
	).Scan(&c.ID, &c.OwnerKey, &c.Name, &c.CreatedAt)
	if errors.Is(err, sql.ErrNoRows) {
		return Collection{}, ErrNotFound
	}
	if err != nil {
		return Collection{}, err
	}
	return c, nil
}

func (s *PostgresStore) DeleteCollection(id string) error {
	res, err := s.db.Exec(`DELETE FROM collections WHERE id = $1`, id)
	if err != nil {
		return err
	}
	n, _ := res.RowsAffected()
	if n == 0 {
		return ErrNotFound
	}
	return nil
}

func (s *PostgresStore) SaveFeature(f Feature) {
	geom, _ := json.Marshal(f.Geometry)
	props, _ := json.Marshal(f.Properties)
	_, _ = s.db.Exec(
		`INSERT INTO features (id, collection_id, owner_key, type, geometry, properties, created_at, updated_at)
		 VALUES ($1, $2, $3, $4, $5, $6, $7, $8)
		 ON CONFLICT (id) DO UPDATE SET geometry = EXCLUDED.geometry, properties = EXCLUDED.properties, updated_at = EXCLUDED.updated_at`,
		f.ID, f.CollectionID, f.OwnerKey, f.Type, geom, props, f.CreatedAt, f.UpdatedAt,
	)
}

func (s *PostgresStore) ListFeaturesByCollection(collectionID string) []Feature {
	rows, err := s.db.Query(
		`SELECT id, collection_id, owner_key, type, geometry, properties, created_at, updated_at
		 FROM features WHERE collection_id = $1 ORDER BY created_at`,
		collectionID,
	)
	if err != nil {
		return nil
	}
	defer rows.Close()
	var result []Feature
	for rows.Next() {
		var f Feature
		var geom, props []byte
		if err := rows.Scan(&f.ID, &f.CollectionID, &f.OwnerKey, &f.Type, &geom, &props, &f.CreatedAt, &f.UpdatedAt); err != nil {
			return result
		}
		_ = json.Unmarshal(geom, &f.Geometry)
		_ = json.Unmarshal(props, &f.Properties)
		result = append(result, f)
	}
	return result
}

func (s *PostgresStore) GetFeature(featureID string) (Feature, error) {
	var f Feature
	var geom, props []byte
	err := s.db.QueryRow(
		`SELECT id, collection_id, owner_key, type, geometry, properties, created_at, updated_at
		 FROM features WHERE id = $1`,
		featureID,
	).Scan(&f.ID, &f.CollectionID, &f.OwnerKey, &f.Type, &geom, &props, &f.CreatedAt, &f.UpdatedAt)
	if errors.Is(err, sql.ErrNoRows) {
		return Feature{}, ErrNotFound
	}
	if err != nil {
		return Feature{}, err
	}
	_ = json.Unmarshal(geom, &f.Geometry)
	_ = json.Unmarshal(props, &f.Properties)
	return f, nil
}

func (s *PostgresStore) DeleteFeature(featureID string) error {
	res, err := s.db.Exec(`DELETE FROM features WHERE id = $1`, featureID)
	if err != nil {
		return err
	}
	n, _ := res.RowsAffected()
	if n == 0 {
		return ErrNotFound
	}
	return nil
}

func (s *PostgresStore) PruneExpired(now time.Time) {
	_, _ = s.db.Exec(`DELETE FROM challenges WHERE expires_at < $1`, now)
	_, _ = s.db.Exec(`DELETE FROM sessions WHERE expires_at < $1`, now)
	_, _ = s.db.Exec(`DELETE FROM invitations WHERE expires_at < $1`, now)
}

func nullStr(s string) interface{} {
	if s == "" {
		return nil
	}
	return s
}

func isUniqueViolation(err error) bool {
	if err == nil {
		return false
	}
	s := err.Error()
	return strings.Contains(s, "duplicate key") || strings.Contains(s, "23505")
}