package db

import (
	"context"
	"database/sql"
	"fmt"
	"time"

	"tangled.org/core/spindle/quota"
)

type CacheEntry struct {
	ID           string
	StorageKey   string
	OwnerDID     string
	RepoDID      string
	Engine       string
	CacheKey     string
	CacheHash    string
	Checksum     string
	SizeBytes    int64
	RestoreCount int64
	State        string
	CreatedAt    time.Time
	LastUsedAt   time.Time
}

const cacheEntryColumns = `
	id, storage_key, owner_did, repo_did, engine, cache_key, cache_hash, checksum,
	size_bytes, restore_count, state, created_at, last_used_at`

func (d *DB) InsertCacheEntry(ctx context.Context, entry CacheEntry) error {
	_, err := d.ExecContext(ctx, `
		insert into cache_entries (
			id, storage_key, owner_did, repo_did, engine, cache_key, cache_hash, checksum,
			size_bytes, restore_count, state, created_at, last_used_at
		) values (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`,
		entry.ID,
		entry.StorageKey,
		entry.OwnerDID,
		entry.RepoDID,
		entry.Engine,
		entry.CacheKey,
		entry.CacheHash,
		entry.Checksum,
		entry.SizeBytes,
		entry.RestoreCount,
		entry.State,
		entry.CreatedAt.UnixNano(),
		entry.LastUsedAt.UnixNano(),
	)
	return err
}
func (d *DB) InsertCacheEntryWithinQuota(ctx context.Context, entry CacheEntry, maxEntries int64) (bool, error) {
	result, err := d.ExecContext(ctx, `
		insert into cache_entries (
			id, storage_key, owner_did, repo_did, engine, cache_key, cache_hash, checksum,
			size_bytes, restore_count, state, created_at, last_used_at
		)
		select ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?
		where ? <= 0 or (
			select count(*) from cache_entries
			where owner_did = ? and state in ('pending', 'ready', 'deleting')
		) < ?`,
		entry.ID,
		entry.StorageKey,
		entry.OwnerDID,
		entry.RepoDID,
		entry.Engine,
		entry.CacheKey,
		entry.CacheHash,
		entry.Checksum,
		entry.SizeBytes,
		entry.RestoreCount,
		entry.State,
		entry.CreatedAt.UnixNano(),
		entry.LastUsedAt.UnixNano(),
		maxEntries,
		entry.OwnerDID,
		maxEntries,
	)
	if err != nil {
		return false, err
	}
	inserted, err := result.RowsAffected()
	return inserted == 1, err
}

func (d *DB) MarkCacheEntryReady(ctx context.Context, id string, sizeBytes int64, qs *QuotaStore, now time.Time) ([]CacheEntry, error) {
	return d.MarkCacheEntryReadyWithChecksum(ctx, id, sizeBytes, "", qs, 0, now)
}

// a nil qs charges nothing
func (d *DB) MarkCacheEntryReadyWithChecksum(ctx context.Context, id string, sizeBytes int64, checksum string, qs *QuotaStore, maxTotalBytes int64, now time.Time) ([]CacheEntry, error) {
	tx, err := d.BeginTx(ctx, nil)
	if err != nil {
		return nil, err
	}
	defer tx.Rollback()
	superseded, ready, err := markCacheEntryReady(ctx, tx, id, sizeBytes, checksum, qs, maxTotalBytes, now)
	if err != nil {
		return nil, err
	}
	if !ready {
		return nil, sql.ErrNoRows
	}
	if err := tx.Commit(); err != nil {
		return nil, err
	}
	return superseded, nil
}

type execQuerier interface {
	ExecContext(context.Context, string, ...any) (sql.Result, error)
	QueryRowContext(context.Context, string, ...any) *sql.Row
	QueryContext(context.Context, string, ...any) (*sql.Rows, error)
}

func markCacheEntryReady(ctx context.Context, tx execQuerier, id string, sizeBytes int64, checksum string, qs *QuotaStore, maxTotalBytes int64, now time.Time) ([]CacheEntry, bool, error) {
	current, err := scanCacheEntry(tx.QueryRowContext(ctx, `select `+cacheEntryColumns+` from cache_entries where id = ?`, id))
	if err == sql.ErrNoRows {
		return nil, false, nil
	}
	if err != nil {
		return nil, false, err
	}
	if current.State != "pending" {
		if current.State == "ready" {
			if current.SizeBytes != sizeBytes {
				return nil, false, fmt.Errorf("cache entry %q size changed from %d to %d", id, current.SizeBytes, sizeBytes)
			}
			return nil, true, nil
		}
		return nil, false, nil
	}

	rows, err := tx.QueryContext(ctx, `
		update cache_entries
		set state = 'deleting'
		where repo_did = ? and engine = ? and cache_key = ? and cache_hash = ?
			and state = 'ready'
		returning `+cacheEntryColumns, current.RepoDID, current.Engine, current.CacheKey, current.CacheHash)
	if err != nil {
		return nil, false, err
	}
	var superseded []CacheEntry
	for rows.Next() {
		entry, err := scanCacheEntry(rows)
		if err != nil {
			rows.Close()
			return nil, false, err
		}
		superseded = append(superseded, *entry)
	}
	if err := rows.Close(); err != nil {
		return nil, false, err
	}
	if err := rows.Err(); err != nil {
		return nil, false, err
	}

	if maxTotalBytes > 0 {
		evicted, fits, err := evictForTotal(ctx, tx, sizeBytes, maxTotalBytes)
		if err != nil {
			return nil, false, err
		}
		if !fits {
			return nil, false, unmarkCacheEntries(ctx, tx, superseded)
		}
		superseded = append(superseded, evicted...)
	}

	if qs != nil {
		// what this save replaces stops counting against it, and counts again if the charge is refused
		if _, err := tx.ExecContext(ctx, `savepoint replace_charge`); err != nil {
			return nil, false, err
		}
		for _, entry := range superseded {
			if err := unchargeCacheEntry(ctx, tx, entry.RepoDID, entry.ID); err != nil {
				return nil, false, err
			}
		}
		res, err := qs.ChargeTx(ctx, tx, quota.ReserveRequest{
			Kind:      quota.KindGenericCache,
			Key:       current.ID,
			Identity:  quota.Identity{OwnerDID: current.OwnerDID, RepoDID: current.RepoDID},
			Resources: quota.Resources{quota.ResourceCacheStorageBytes: sizeBytes},
		})
		if err != nil {
			return nil, false, err
		}
		if !res.Allowed {
			if _, err := tx.ExecContext(ctx, `rollback to replace_charge`); err != nil {
				return nil, false, err
			}
		}
		if _, err := tx.ExecContext(ctx, `release replace_charge`); err != nil {
			return nil, false, err
		}
		if !res.Allowed {
			return nil, false, unmarkCacheEntries(ctx, tx, superseded)
		}
	}

	updated, err := tx.QueryContext(ctx, `
		update cache_entries
		set state = 'ready', checksum = ?, size_bytes = ?, last_used_at = ?
		where id = ? and state = 'pending'
		returning id`, checksum, sizeBytes, now.UnixNano(), id)
	if err != nil {
		return nil, false, err
	}
	ready := updated.Next()
	updated.Close()
	return superseded, ready, nil
}

// a refused save still commits inside the mill's batch tx
func unmarkCacheEntries(ctx context.Context, tx execQuerier, entries []CacheEntry) error {
	for _, entry := range entries {
		if _, err := tx.ExecContext(ctx, `update cache_entries set state = 'ready' where id = ? and state = 'deleting'`, entry.ID); err != nil {
			return err
		}
	}
	return nil
}

const (
	readyCacheBytesQuery = `select coalesce(sum(size_bytes), 0) from cache_entries where state = 'ready'`
	// entries a restore holds stay out of eviction
	evictableCacheEntriesQuery = `select ` + cacheEntryColumns + ` from cache_entries
		where state = 'ready' and restore_count = 0
		order by last_used_at, created_at`
)

// evicts nothing when even every evictable entry would not make room
func evictForTotal(ctx context.Context, tx execQuerier, sizeBytes, maxTotalBytes int64) ([]CacheEntry, bool, error) {
	var used int64
	if err := tx.QueryRowContext(ctx, readyCacheBytesQuery).Scan(&used); err != nil {
		return nil, false, err
	}
	need := used + sizeBytes - maxTotalBytes
	if need <= 0 {
		return nil, true, nil
	}
	rows, err := tx.QueryContext(ctx, evictableCacheEntriesQuery)
	if err != nil {
		return nil, false, err
	}
	var evicted []CacheEntry
	for need > 0 && rows.Next() {
		entry, err := scanCacheEntry(rows)
		if err != nil {
			rows.Close()
			return nil, false, err
		}
		evicted = append(evicted, *entry)
		need -= entry.SizeBytes
	}
	if err := rows.Close(); err != nil {
		return nil, false, err
	}
	if err := rows.Err(); err != nil {
		return nil, false, err
	}
	if need > 0 {
		return nil, false, nil
	}
	for _, entry := range evicted {
		if _, err := tx.ExecContext(ctx, `update cache_entries set state = 'deleting' where id = ?`, entry.ID); err != nil {
			return nil, false, err
		}
	}
	return evicted, true, nil
}

func (tx *EventBatchTx) MarkCacheEntryReadyWithChecksum(ctx context.Context, id string, sizeBytes int64, checksum string, qs *QuotaStore, maxTotalBytes int64, now time.Time) ([]CacheEntry, bool, error) {
	return markCacheEntryReady(ctx, tx.tx, id, sizeBytes, checksum, qs, maxTotalBytes, now)
}

func (tx *EventBatchTx) TouchCacheEntry(ctx context.Context, id string, now time.Time) error {
	_, err := tx.tx.ExecContext(ctx, `
		update cache_entries
		set last_used_at = max(last_used_at, ?)
		where id = ? and state = 'ready'`, now.UnixNano(), id)
	return err
}

func (tx *EventBatchTx) DiscardCacheEntry(ctx context.Context, id string, pendingOnly bool) (*CacheEntry, error) {
	stateFilter := ""
	if pendingOnly {
		stateFilter = " and state in ('pending', 'deleting')"
	}
	current, err := scanCacheEntry(tx.tx.QueryRowContext(ctx, `
		update cache_entries
		set state = 'deleting'
		where id = ?`+stateFilter+`
		returning `+cacheEntryColumns, id))
	if err == sql.ErrNoRows {
		return nil, nil
	}
	return current, err
}

func (d *DB) FindCacheEntry(ctx context.Context, repoDID, engine, key, hash string) (*CacheEntry, error) {
	return scanCacheEntry(d.QueryRowContext(ctx, `
		select `+cacheEntryColumns+`
		from cache_entries
		where repo_did = ? and engine = ? and cache_key = ? and cache_hash = ? and state = 'ready'
		order by created_at desc
		limit 1`, repoDID, engine, key, hash))
}

func (d *DB) FindFallbackCacheEntry(ctx context.Context, repoDID, engine, key, excludeHash string) (*CacheEntry, error) {
	return scanCacheEntry(d.QueryRowContext(ctx, `
		select `+cacheEntryColumns+`
		from cache_entries
		where repo_did = ? and engine = ? and cache_key = ?
			and cache_hash <> ? and cache_hash <> '' and state = 'ready'
		order by created_at desc
		limit 1`, repoDID, engine, key, excludeHash))
}

func (d *DB) BeginCacheRestore(ctx context.Context, id string) (bool, error) {
	result, err := d.ExecContext(ctx, `update cache_entries set restore_count = restore_count + 1 where id = ? and state = 'ready'`, id)
	if err != nil {
		return false, err
	}
	n, err := result.RowsAffected()
	return n == 1, err
}

func (d *DB) EndCacheRestore(ctx context.Context, id string) error {
	_, err := d.ExecContext(ctx, `update cache_entries set restore_count = restore_count - 1 where id = ? and restore_count > 0`, id)
	return err
}

func (d *DB) CacheEntryPinned(ctx context.Context, id string) (bool, error) {
	var count int64
	err := d.QueryRowContext(ctx, `select restore_count from cache_entries where id = ?`, id).Scan(&count)
	if err == sql.ErrNoRows {
		return false, nil
	}
	return count > 0, err
}

func (d *DB) TouchCacheEntry(ctx context.Context, id string, now time.Time) error {
	_, err := d.ExecContext(ctx, `
		update cache_entries set last_used_at = max(last_used_at, ?) where id = ? and state = 'ready'`, now.UnixNano(), id)
	return err
}

func (d *DB) ClaimPendingCacheEntry(ctx context.Context, id string) (bool, error) {
	result, err := d.ExecContext(ctx, `
		update cache_entries set state = 'deleting' where id = ? and state = 'pending'`, id)
	if err != nil {
		return false, err
	}
	changed, err := result.RowsAffected()
	return changed == 1, err
}

func (d *DB) ClaimCacheEntry(ctx context.Context, id, expectedState string, expectedLastUsed time.Time) (bool, error) {
	result, err := d.ExecContext(ctx, `
		update cache_entries
		set state = 'deleting'
		where id = ? and state = ? and last_used_at = ? and restore_count = 0`,
		id, expectedState, expectedLastUsed.UnixNano())
	if err != nil {
		return false, err
	}
	changed, err := result.RowsAffected()
	return changed == 1, err
}

func (d *DB) ReadyCacheBytes(ctx context.Context) (int64, error) {
	var used int64
	err := d.QueryRowContext(ctx, readyCacheBytesQuery).Scan(&used)
	return used, err
}

func (d *DB) EvictableCacheEntries(ctx context.Context, limit int) ([]CacheEntry, error) {
	return d.queryCacheEntries(ctx, evictableCacheEntriesQuery+` limit ?`, limit)
}

func (d *DB) EvictableCacheEntriesOf(ctx context.Context, repoDID, ownerDID string, limit int) ([]CacheEntry, error) {
	return d.queryCacheEntries(ctx, `select `+cacheEntryColumns+` from cache_entries
		where state = 'ready' and restore_count = 0 and (repo_did = ? or owner_did = ?)
		order by last_used_at, created_at
		limit ?`, repoDID, ownerDID, limit)
}

func (d *DB) GetCacheEntry(ctx context.Context, id string) (*CacheEntry, error) {
	return scanCacheEntry(d.QueryRowContext(ctx, `select `+cacheEntryColumns+` from cache_entries where id = ?`, id))
}

func (d *DB) RestoreCacheEntryState(ctx context.Context, id, state string) error {
	_, err := d.ExecContext(ctx, `
		update cache_entries set state = ? where id = ? and state = 'deleting'`, state, id)
	return err
}

func (d *DB) ExpiredCacheEntries(ctx context.Context, readyBefore, pendingBefore time.Time, limit int) ([]CacheEntry, error) {
	return d.queryCacheEntries(ctx, `
		select `+cacheEntryColumns+`
		from cache_entries
		where (state = 'ready' and last_used_at < ?)
			or (state in ('pending', 'deleting') and last_used_at < ?)
		order by last_used_at
		limit ?`,
		readyBefore.UnixNano(),
		pendingBefore.UnixNano(),
		limit,
	)
}

// discardPendingCacheEntriesForLease fences uploads that can arrive after a lease dies.
// Their tombstones let the cache pruner remove a late object without metadata.
func (d *DB) DiscardPendingCacheEntriesForLease(leaseID string, now time.Time) ([]string, error) {
	tx, err := d.Begin()
	if err != nil {
		return nil, err
	}
	defer tx.Rollback()

	rows, err := tx.Query(`
		select c.id, c.storage_key
		from cache_entries c
		join mill_cache_capabilities cap on cap.cache_id = c.id
		where cap.lease_id = ? and cap.action = 'save' and c.state = 'pending'`, leaseID)
	if err != nil {
		return nil, err
	}
	var keys []string
	var ids []string
	for rows.Next() {
		var id, key string
		if err := rows.Scan(&id, &key); err != nil {
			_ = rows.Close()
			return nil, err
		}
		ids = append(ids, id)
		keys = append(keys, key)
	}
	if err := rows.Err(); err != nil {
		_ = rows.Close()
		return nil, err
	}
	if err := rows.Close(); err != nil {
		return nil, err
	}
	for i, id := range ids {
		if _, err := tx.Exec(`update cache_entries set state = 'deleting' where id = ? and state = 'pending'`, id); err != nil {
			return nil, err
		}
		if _, err := tx.Exec(`
			insert into cache_object_deletions (storage_key, created_at)
			values (?, ?)
			on conflict(storage_key) do nothing`, keys[i], now.UnixNano()); err != nil {
			return nil, err
		}
	}
	if err := tx.Commit(); err != nil {
		return nil, err
	}
	return keys, nil
}

func (d *DB) DeleteCacheEntry(ctx context.Context, id string) error {
	return d.deleteCacheEntries(ctx, `id = ?`, id)
}

func (d *DB) DeleteCacheEntriesByStorageKey(ctx context.Context, key string) error {
	return d.deleteCacheEntries(ctx, `storage_key = ? and restore_count = 0`, key)
}

func (d *DB) deleteCacheEntries(ctx context.Context, where string, args ...any) error {
	tx, err := d.BeginTx(ctx, nil)
	if err != nil {
		return err
	}
	defer tx.Rollback()
	rows, err := tx.QueryContext(ctx, `delete from cache_entries where `+where+` returning repo_did, id`, args...)
	if err != nil {
		return err
	}
	var gone [][2]string
	for rows.Next() {
		var repo, id string
		if err := rows.Scan(&repo, &id); err != nil {
			rows.Close()
			return err
		}
		gone = append(gone, [2]string{repo, id})
	}
	if err := rows.Close(); err != nil {
		return err
	}
	for _, entry := range gone {
		if err := unchargeCacheEntry(ctx, tx, entry[0], entry[1]); err != nil {
			return err
		}
	}
	return tx.Commit()
}

func unchargeCacheEntry(ctx context.Context, tx execQuerier, repoDID, id string) error {
	_, err := tx.ExecContext(ctx, `
		delete from quota_allocations where repo_did = ? and kind = ? and key = ?`,
		repoDID, quota.KindGenericCache, id)
	return err
}

type MillCacheCapability struct {
	Action     string
	CacheID    string
	StorageKey string
}

func (d *DB) SaveMillCacheCapabilities(leaseID string, capabilities []MillCacheCapability) error {
	tx, err := d.Begin()
	if err != nil {
		return err
	}
	defer tx.Rollback()
	if _, err := tx.Exec(`delete from mill_cache_capabilities where lease_id = ?`, leaseID); err != nil {
		return err
	}
	for _, capability := range capabilities {
		if capability.Action != "restore" && capability.Action != "save" {
			return fmt.Errorf("cache capability has invalid action %q", capability.Action)
		}
		if capability.CacheID == "" || capability.StorageKey == "" {
			return fmt.Errorf("cache capability has empty identity")
		}
		if _, err := tx.Exec(`
			insert into mill_cache_capabilities (
				lease_id, action, cache_id, storage_key
			) values (?, ?, ?, ?)`,
			leaseID,
			capability.Action,
			capability.CacheID,
			capability.StorageKey,
		); err != nil {
			return err
		}
	}
	return tx.Commit()
}

func (tx *EventBatchTx) ConsumeMillCacheCapability(ctx context.Context, leaseID, action, cacheID string) (string, bool, error) {
	var storageKey string
	err := tx.tx.QueryRowContext(ctx, `
		delete from mill_cache_capabilities
		where lease_id = ? and action = ? and cache_id = ?
		returning storage_key`,
		leaseID, action, cacheID,
	).Scan(&storageKey)
	if err == sql.ErrNoRows {
		return "", false, nil
	}
	if err != nil {
		return "", false, err
	}
	return storageKey, true, nil
}

func (tx *EventBatchTx) QueueCacheObjectDeletion(ctx context.Context, storageKey string, now time.Time) error {
	_, err := tx.tx.ExecContext(ctx, `
		insert into cache_object_deletions (storage_key, created_at)
		values (?, ?)
		on conflict(storage_key) do nothing`,
		storageKey, now.UnixNano(),
	)
	return err
}

func (d *DB) PendingCacheObjectDeletions(ctx context.Context, limit int) ([]string, error) {
	rows, err := d.QueryContext(ctx, `
		select storage_key
		from cache_object_deletions
		order by created_at
		limit ?`, limit)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var keys []string
	for rows.Next() {
		var key string
		if err := rows.Scan(&key); err != nil {
			return nil, err
		}
		keys = append(keys, key)
	}
	return keys, rows.Err()
}

func (d *DB) CacheEntriesByStorageKey(ctx context.Context, key string) ([]CacheEntry, error) {
	return d.queryCacheEntries(ctx, `select `+cacheEntryColumns+` from cache_entries where storage_key = ?`, key)
}

func (d *DB) queryCacheEntries(ctx context.Context, query string, args ...any) ([]CacheEntry, error) {
	rows, err := d.QueryContext(ctx, query, args...)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var entries []CacheEntry
	for rows.Next() {
		entry, err := scanCacheEntry(rows)
		if err != nil {
			return nil, err
		}
		entries = append(entries, *entry)
	}
	return entries, rows.Err()
}

func (d *DB) CompleteCacheObjectDeletion(ctx context.Context, storageKey string) error {
	_, err := d.ExecContext(ctx, `delete from cache_object_deletions where storage_key = ?`, storageKey)
	return err
}

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

func scanCacheEntry(row cacheEntryScanner) (*CacheEntry, error) {
	var entry CacheEntry
	var createdAt, lastUsedAt int64
	if err := row.Scan(
		&entry.ID,
		&entry.StorageKey,
		&entry.OwnerDID,
		&entry.RepoDID,
		&entry.Engine,
		&entry.CacheKey,
		&entry.CacheHash,
		&entry.Checksum,
		&entry.SizeBytes,
		&entry.RestoreCount,
		&entry.State,
		&createdAt,
		&lastUsedAt,
	); err != nil {
		return nil, err
	}
	entry.CreatedAt = time.Unix(0, createdAt)
	entry.LastUsedAt = time.Unix(0, lastUsedAt)
	return &entry, nil
}