db.go

package main

import (
	"context"
	"encoding/json"
	"fmt"
	"log"
	"time"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"
)

const pgSchema = `
CREATE TABLE IF NOT EXISTS blocks (
    slot         BIGINT PRIMARY KEY,
    epoch        INTEGER NOT NULL,
    block_hash   TEXT NOT NULL,
    vrf_output   BYTEA NOT NULL,
    nonce_value  BYTEA NOT NULL,
    created_at   TIMESTAMPTZ DEFAULT NOW()
);
CREATE INDEX IF NOT EXISTS idx_blocks_epoch ON blocks(epoch);

CREATE TABLE IF NOT EXISTS epoch_nonces (
    epoch            INTEGER PRIMARY KEY,
    evolving_nonce   BYTEA NOT NULL,
    candidate_nonce  BYTEA,
    final_nonce      BYTEA,
    block_count      INTEGER DEFAULT 0,
    source           TEXT DEFAULT 'chain_sync',
    updated_at       TIMESTAMPTZ DEFAULT NOW()
);

CREATE TABLE IF NOT EXISTS leader_schedules (
    epoch          INTEGER PRIMARY KEY,
    pool_stake     BIGINT NOT NULL,
    total_stake    BIGINT NOT NULL,
    epoch_nonce    TEXT NOT NULL,
    sigma          DOUBLE PRECISION,
    ideal_slots    DOUBLE PRECISION,
    slot_count     INTEGER DEFAULT 0,
    performance    DOUBLE PRECISION,
    slots          JSONB DEFAULT '[]',
    posted         BOOLEAN DEFAULT FALSE,
    history_classified BOOLEAN DEFAULT FALSE,
    calculated_at  TIMESTAMPTZ DEFAULT NOW()
);

CREATE TABLE IF NOT EXISTS slot_outcomes (
    epoch    INTEGER NOT NULL,
    slot     BIGINT NOT NULL,
    outcome  TEXT NOT NULL,
    opponent TEXT,
    PRIMARY KEY (epoch, slot)
);
CREATE INDEX IF NOT EXISTS idx_slot_outcomes_epoch ON slot_outcomes(epoch);
`

// PgStore implements Store using PostgreSQL via pgx.
type PgStore struct {
	pool *pgxpool.Pool
}

// NewPgStore connects to PostgreSQL and creates tables if they don't exist.
func NewPgStore(connString string) (*PgStore, error) {
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	pool, err := pgxpool.New(ctx, connString)
	if err != nil {
		return nil, fmt.Errorf("connecting to database: %w", err)
	}

	if err := pool.Ping(ctx); err != nil {
		pool.Close()
		return nil, fmt.Errorf("pinging database: %w", err)
	}

	if _, err := pool.Exec(ctx, pgSchema); err != nil {
		pool.Close()
		return nil, fmt.Errorf("creating schema: %w", err)
	}

	// Migrations for existing databases
	pool.Exec(ctx, `ALTER TABLE leader_schedules ADD COLUMN IF NOT EXISTS history_classified BOOLEAN DEFAULT FALSE`)

	log.Println("PostgreSQL connected and schema initialized")
	return &PgStore{pool: pool}, nil
}

func (s *PgStore) Close() error {
	s.pool.Close()
	return nil
}

func (s *PgStore) InsertBlock(ctx context.Context, slot uint64, epoch int, blockHash string, vrfOutput, nonceValue []byte) (bool, error) {
	result, err := s.pool.Exec(ctx,
		`INSERT INTO blocks (slot, epoch, block_hash, vrf_output, nonce_value)
		 VALUES ($1, $2, $3, $4, $5)
		 ON CONFLICT (slot) DO NOTHING`,
		int64(slot), epoch, blockHash, vrfOutput, nonceValue,
	)
	if err != nil {
		return false, err
	}
	return result.RowsAffected() > 0, nil
}

func (s *PgStore) UpsertEvolvingNonce(ctx context.Context, epoch int, nonce []byte, blockCount int) error {
	_, err := s.pool.Exec(ctx,
		`INSERT INTO epoch_nonces (epoch, evolving_nonce, block_count, updated_at)
		 VALUES ($1, $2, $3, NOW())
		 ON CONFLICT (epoch) DO UPDATE SET
		   evolving_nonce = EXCLUDED.evolving_nonce,
		   block_count = EXCLUDED.block_count,
		   updated_at = NOW()`,
		epoch, nonce, blockCount,
	)
	return err
}

func (s *PgStore) SetCandidateNonce(ctx context.Context, epoch int, nonce []byte) error {
	_, err := s.pool.Exec(ctx,
		`INSERT INTO epoch_nonces (epoch, evolving_nonce, candidate_nonce, updated_at)
		 VALUES ($1, $2, $2, NOW())
		 ON CONFLICT (epoch) DO UPDATE SET
		   candidate_nonce = EXCLUDED.candidate_nonce,
		   updated_at = NOW()`,
		epoch, nonce,
	)
	return err
}

func (s *PgStore) DeleteCandidateNonce(ctx context.Context, epoch int) error {
	_, err := s.pool.Exec(ctx,
		`UPDATE epoch_nonces SET candidate_nonce = NULL, updated_at = NOW() WHERE epoch = $1`,
		epoch,
	)
	return err
}

func (s *PgStore) SetFinalNonce(ctx context.Context, epoch int, nonce []byte, source string) error {
	_, err := s.pool.Exec(ctx,
		`INSERT INTO epoch_nonces (epoch, evolving_nonce, final_nonce, source, updated_at)
		 VALUES ($1, $2, $2, $3, NOW())
		 ON CONFLICT (epoch) DO UPDATE SET
		   final_nonce = EXCLUDED.final_nonce,
		   source = EXCLUDED.source,
		   updated_at = NOW()`,
		epoch, nonce, source,
	)
	return err
}

func (s *PgStore) GetFinalNonce(ctx context.Context, epoch int) ([]byte, error) {
	var nonce []byte
	err := s.pool.QueryRow(ctx,
		`SELECT final_nonce FROM epoch_nonces WHERE epoch = $1 AND final_nonce IS NOT NULL`,
		epoch,
	).Scan(&nonce)
	if err != nil {
		return nil, err
	}
	return nonce, nil
}

func (s *PgStore) GetEvolvingNonce(ctx context.Context, epoch int) ([]byte, int, error) {
	var nonce []byte
	var blockCount int
	err := s.pool.QueryRow(ctx,
		`SELECT evolving_nonce, block_count FROM epoch_nonces WHERE epoch = $1`,
		epoch,
	).Scan(&nonce, &blockCount)
	if err != nil {
		return nil, 0, err
	}
	return nonce, blockCount, nil
}

func (s *PgStore) InsertLeaderSchedule(ctx context.Context, schedule *LeaderSchedule) error {
	slotsJSON, err := json.Marshal(schedule.AssignedSlots)
	if err != nil {
		return fmt.Errorf("marshaling slots: %w", err)
	}

	performance := 0.0
	if schedule.IdealSlots > 0 {
		performance = float64(len(schedule.AssignedSlots)) / schedule.IdealSlots * 100
	}

	_, err = s.pool.Exec(ctx,
		`INSERT INTO leader_schedules (epoch, pool_stake, total_stake, epoch_nonce, sigma, ideal_slots, slot_count, performance, slots, calculated_at)
		 VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)
		 ON CONFLICT (epoch) DO UPDATE SET
		   pool_stake = EXCLUDED.pool_stake,
		   total_stake = EXCLUDED.total_stake,
		   epoch_nonce = EXCLUDED.epoch_nonce,
		   sigma = EXCLUDED.sigma,
		   ideal_slots = EXCLUDED.ideal_slots,
		   slot_count = EXCLUDED.slot_count,
		   performance = EXCLUDED.performance,
		   slots = EXCLUDED.slots,
		   calculated_at = EXCLUDED.calculated_at`,
		schedule.Epoch,
		int64(schedule.PoolStake),
		int64(schedule.TotalStake),
		schedule.EpochNonce,
		schedule.Sigma,
		schedule.IdealSlots,
		len(schedule.AssignedSlots),
		performance,
		slotsJSON,
		schedule.CalculatedAt,
	)
	return err
}

func (s *PgStore) IsSchedulePosted(ctx context.Context, epoch int) bool {
	var posted bool
	err := s.pool.QueryRow(ctx,
		`SELECT posted FROM leader_schedules WHERE epoch = $1`,
		epoch,
	).Scan(&posted)
	if err != nil {
		return false
	}
	return posted
}

func (s *PgStore) MarkSchedulePosted(ctx context.Context, epoch int) error {
	_, err := s.pool.Exec(ctx,
		`UPDATE leader_schedules SET posted = TRUE WHERE epoch = $1`,
		epoch,
	)
	return err
}

func (s *PgStore) GetLastSyncedSlot(ctx context.Context) (uint64, error) {
	var slot *int64
	err := s.pool.QueryRow(ctx,
		`SELECT MAX(slot) FROM blocks`,
	).Scan(&slot)
	if err != nil {
		return 0, err
	}
	if slot == nil {
		return 0, nil
	}
	return uint64(*slot), nil
}

func (s *PgStore) GetBlockHash(ctx context.Context, slot uint64) (string, error) {
	var hash string
	err := s.pool.QueryRow(ctx,
		`SELECT block_hash FROM blocks WHERE slot = $1`,
		int64(slot),
	).Scan(&hash)
	return hash, err
}

func (s *PgStore) GetForgedSlots(ctx context.Context, epoch int) ([]uint64, error) {
	rows, err := s.pool.Query(ctx,
		`SELECT slot FROM blocks WHERE epoch = $1 ORDER BY slot`, epoch)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var slots []uint64
	for rows.Next() {
		var slot int64
		if err := rows.Scan(&slot); err != nil {
			return nil, err
		}
		slots = append(slots, uint64(slot))
	}
	return slots, rows.Err()
}

func (s *PgStore) InsertBlockBatch(ctx context.Context, blocks []BlockData) (int, error) {
	tx, err := s.pool.Begin(ctx)
	if err != nil {
		return 0, err
	}
	defer tx.Rollback(ctx)

	// Temp table with no constraints — CopyFrom always succeeds even with duplicates
	_, err = tx.Exec(ctx, `CREATE TEMP TABLE blocks_staging (
		slot BIGINT, epoch INT, block_hash TEXT, vrf_output BYTEA, nonce_value BYTEA
	) ON COMMIT DROP`)
	if err != nil {
		return 0, err
	}

	rows := make([][]interface{}, len(blocks))
	for i, b := range blocks {
		nonceValue := vrfNonceValueForEpoch(b.VrfOutput, b.Epoch, b.NetworkMagic)
		rows[i] = []interface{}{int64(b.Slot), b.Epoch, b.BlockHash, b.VrfOutput, nonceValue}
	}

	// COPY into staging (no constraints = no failures on duplicate keys)
	_, err = tx.CopyFrom(ctx,
		pgx.Identifier{"blocks_staging"},
		[]string{"slot", "epoch", "block_hash", "vrf_output", "nonce_value"},
		pgx.CopyFromRows(rows),
	)
	if err != nil {
		return 0, err
	}

	// Merge into blocks — duplicates silently skipped, count actually inserted
	result, err := tx.Exec(ctx, `INSERT INTO blocks (slot, epoch, block_hash, vrf_output, nonce_value)
		SELECT slot, epoch, block_hash, vrf_output, nonce_value FROM blocks_staging
		ON CONFLICT (slot) DO NOTHING`)
	if err != nil {
		return 0, err
	}

	inserted := int(result.RowsAffected())
	return inserted, tx.Commit(ctx)
}


func (s *PgStore) GetBlockByHash(ctx context.Context, hashPrefix string) ([]BlockRecord, error) {
	rows, err := s.pool.Query(ctx,
		`SELECT slot, epoch, block_hash FROM blocks WHERE block_hash LIKE $1 ORDER BY slot`,
		hashPrefix+"%",
	)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var records []BlockRecord
	for rows.Next() {
		var r BlockRecord
		if err := rows.Scan(&r.Slot, &r.Epoch, &r.BlockHash); err != nil {
			return nil, err
		}
		records = append(records, r)
	}
	if err := rows.Err(); err != nil {
		return nil, err
	}
	return records, nil
}

func (s *PgStore) GetLeaderSchedule(ctx context.Context, epoch int) (*LeaderSchedule, error) {
	var (
		poolStake    int64
		totalStake   int64
		epochNonce   string
		sigma        float64
		idealSlots   float64
		slotsJSON    []byte
		calculatedAt time.Time
	)

	err := s.pool.QueryRow(ctx,
		`SELECT pool_stake, total_stake, epoch_nonce, sigma, ideal_slots, slots, calculated_at
		 FROM leader_schedules WHERE epoch = $1`,
		epoch,
	).Scan(&poolStake, &totalStake, &epochNonce, &sigma, &idealSlots, &slotsJSON, &calculatedAt)
	if err != nil {
		return nil, err
	}

	var slots []LeaderSlot
	if err := json.Unmarshal(slotsJSON, &slots); err != nil {
		return nil, fmt.Errorf("unmarshaling slots: %w", err)
	}

	return &LeaderSchedule{
		Epoch:         epoch,
		EpochNonce:    epochNonce,
		PoolStake:     uint64(poolStake),
		TotalStake:    uint64(totalStake),
		Sigma:         sigma,
		IdealSlots:    idealSlots,
		AssignedSlots: slots,
		CalculatedAt:  calculatedAt,
	}, nil
}


func (s *PgStore) StreamBlockVrfOutputs(ctx context.Context) (BlockVrfRows, error) {
	rows, err := s.pool.Query(ctx,
		`SELECT epoch, slot, vrf_output, nonce_value, block_hash FROM blocks ORDER BY slot`,
	)
	if err != nil {
		return nil, err
	}
	return &pgBlockVrfRows{rows: rows}, nil
}

type pgBlockVrfRows struct {
	rows       pgx.Rows
	epoch      int
	slot       uint64
	vrfOutput  []byte
	nonceValue []byte
	blockHash  string
	err        error
	closed     bool
}

func (r *pgBlockVrfRows) Next() bool {
	if r.closed {
		return false
	}
	if !r.rows.Next() {
		r.err = r.rows.Err()
		r.closed = true
		return false
	}
	var slotInt64 int64
	r.err = r.rows.Scan(&r.epoch, &slotInt64, &r.vrfOutput, &r.nonceValue, &r.blockHash)
	r.slot = uint64(slotInt64)
	return r.err == nil
}

func (r *pgBlockVrfRows) Scan() (epoch int, slot uint64, vrfOutput []byte, nonceValue []byte, blockHash string, err error) {
	return r.epoch, r.slot, r.vrfOutput, r.nonceValue, r.blockHash, r.err
}

func (r *pgBlockVrfRows) Close() {
	if !r.closed {
		r.rows.Close()
		r.closed = true
	}
}

func (r *pgBlockVrfRows) Err() error {
	return r.err
}

func (s *PgStore) GetLastNBlocks(ctx context.Context, n int) ([]BlockRecord, error) {
	rows, err := s.pool.Query(ctx,
		`SELECT slot, epoch, block_hash FROM blocks ORDER BY slot DESC LIMIT $1`, n)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var records []BlockRecord
	for rows.Next() {
		var r BlockRecord
		if err := rows.Scan(&r.Slot, &r.Epoch, &r.BlockHash); err != nil {
			return nil, err
		}
		records = append(records, r)
	}
	return records, rows.Err()
}

func (s *PgStore) GetBlockCountForEpoch(ctx context.Context, epoch int) (int, error) {
	var count int
	err := s.pool.QueryRow(ctx,
		`SELECT COUNT(*) FROM blocks WHERE epoch = $1`, epoch).Scan(&count)
	return count, err
}


func (s *PgStore) GetVrfOutputsForEpoch(ctx context.Context, epoch int) ([]VrfBlock, error) {
	rows, err := s.pool.Query(ctx,
		`SELECT epoch, vrf_output FROM blocks WHERE epoch = $1 ORDER BY slot`, epoch)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var blocks []VrfBlock
	for rows.Next() {
		var b VrfBlock
		if err := rows.Scan(&b.Epoch, &b.VrfOutput); err != nil {
			return nil, err
		}
		blocks = append(blocks, b)
	}
	return blocks, rows.Err()
}

func (s *PgStore) GetCandidateNonce(ctx context.Context, epoch int) ([]byte, error) {
	var nonce []byte
	err := s.pool.QueryRow(ctx,
		`SELECT candidate_nonce FROM epoch_nonces WHERE epoch = $1 AND candidate_nonce IS NOT NULL`, epoch).Scan(&nonce)
	if err != nil {
		return nil, err
	}
	return nonce, nil
}


// GetPrevHashOfLastBlock returns the block hash of the second-to-last block
// in the given epoch. This is the prevHash of the last block, which is what
// the Cardano node uses for praosStateLabNonce (η_ph in the TICKN rule).
func (s *PgStore) GetPrevHashOfLastBlock(ctx context.Context, epoch int) (string, error) {
	var hash string
	err := s.pool.QueryRow(ctx,
		`SELECT block_hash FROM blocks WHERE epoch = $1 ORDER BY slot DESC LIMIT 1 OFFSET 1`, epoch).Scan(&hash)
	if err != nil {
		return "", err
	}
	return hash, nil
}

func (s *PgStore) UpsertSlotOutcomes(ctx context.Context, epoch int, outcomes []SlotOutcome) error {
	tx, err := s.pool.Begin(ctx)
	if err != nil {
		return fmt.Errorf("begin tx: %w", err)
	}
	defer tx.Rollback(ctx)

	for _, o := range outcomes {
		_, err := tx.Exec(ctx,
			`INSERT INTO slot_outcomes (epoch, slot, outcome, opponent)
			 VALUES ($1, $2, $3, $4)
			 ON CONFLICT (epoch, slot) DO UPDATE SET
			   outcome = EXCLUDED.outcome,
			   opponent = EXCLUDED.opponent`,
			o.Epoch, int64(o.Slot), o.Outcome, o.Opponent)
		if err != nil {
			return fmt.Errorf("upsert slot %d: %w", o.Slot, err)
		}
	}
	return tx.Commit(ctx)
}

func (s *PgStore) GetSlotOutcomes(ctx context.Context, epoch int) ([]SlotOutcome, error) {
	rows, err := s.pool.Query(ctx,
		`SELECT epoch, slot, outcome, COALESCE(opponent, '') FROM slot_outcomes WHERE epoch = $1 ORDER BY slot`, epoch)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var outcomes []SlotOutcome
	for rows.Next() {
		var o SlotOutcome
		if err := rows.Scan(&o.Epoch, &o.Slot, &o.Outcome, &o.Opponent); err != nil {
			return nil, err
		}
		outcomes = append(outcomes, o)
	}
	return outcomes, rows.Err()
}

func (s *PgStore) IsEpochClassified(ctx context.Context, epoch int) bool {
	var classified bool
	err := s.pool.QueryRow(ctx,
		`SELECT history_classified FROM leader_schedules WHERE epoch = $1`, epoch).Scan(&classified)
	if err != nil {
		return false
	}
	return classified
}

func (s *PgStore) MarkEpochClassified(ctx context.Context, epoch int) error {
	_, err := s.pool.Exec(ctx,
		`UPDATE leader_schedules SET history_classified = TRUE WHERE epoch = $1`, epoch)
	return err
}

func (s *PgStore) DeleteSlotOutcomesBefore(ctx context.Context, epoch int) (int64, error) {
	tag, err := s.pool.Exec(ctx,
		`DELETE FROM slot_outcomes WHERE epoch < $1`, epoch)
	if err != nil {
		return 0, err
	}
	// Also unmark those epochs as classified so they won't be skipped
	_, _ = s.pool.Exec(ctx,
		`UPDATE leader_schedules SET history_classified = FALSE WHERE epoch < $1`, epoch)
	return tag.RowsAffected(), nil
}

func (s *PgStore) HasBlockAtSlot(ctx context.Context, slot uint64) (bool, error) {
	var exists bool
	err := s.pool.QueryRow(ctx, `SELECT EXISTS(SELECT 1 FROM blocks WHERE slot = $1)`, slot).Scan(&exists)
	return exists, err
}

func (s *PgStore) DeleteBlocksAfterSlot(ctx context.Context, slot uint64) (int64, error) {
	tag, err := s.pool.Exec(ctx, `DELETE FROM blocks WHERE slot > $1`, slot)
	if err != nil {
		return 0, err
	}
	return tag.RowsAffected(), nil
}

func (s *PgStore) TruncateAll(ctx context.Context) error {
	_, err := s.pool.Exec(ctx, `TRUNCATE blocks, epoch_nonces, leader_schedules, slot_outcomes`)
	return err
}