package db

import (
	"context"
	"sort"
	"time"
)

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

const cacheEntryColumns = `
	id, storage_key, owner_did, repo_did, engine, cache_key, cache_hash,
	size_bytes, 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,
			size_bytes, state, created_at, last_used_at
		) values (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`,
		entry.ID,
		entry.StorageKey,
		entry.OwnerDID,
		entry.RepoDID,
		entry.Engine,
		entry.CacheKey,
		entry.CacheHash,
		entry.SizeBytes,
		entry.State,
		entry.CreatedAt.UnixNano(),
		entry.LastUsedAt.UnixNano(),
	)
	return err
}

func (d *DB) MarkCacheEntryReady(ctx context.Context, id string, sizeBytes int64, now time.Time) ([]CacheEntry, error) {
	tx, err := d.BeginTx(ctx, nil)
	if err != nil {
		return nil, err
	}
	defer tx.Rollback()

	var repoDID, engine, key, hash string
	if err := tx.QueryRowContext(ctx, `
		update cache_entries
		set state = 'ready', size_bytes = ?, last_used_at = ?
		where id = ? and state = 'pending'
		returning repo_did, engine, cache_key, cache_hash`,
		sizeBytes, now.UnixNano(), id).Scan(&repoDID, &engine, &key, &hash); err != nil {
		return nil, err
	}

	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' and id <> ?
		returning `+cacheEntryColumns, repoDID, engine, key, hash, id)
	if err != nil {
		return nil, err
	}
	var superseded []CacheEntry
	for rows.Next() {
		entry, err := scanCacheEntry(rows)
		if err != nil {
			rows.Close()
			return nil, err
		}
		superseded = append(superseded, *entry)
	}
	if err := rows.Close(); err != nil {
		return nil, err
	}
	if err := rows.Err(); err != nil {
		return nil, err
	}
	sort.Slice(superseded, func(i, j int) bool {
		return superseded[i].CreatedAt.After(superseded[j].CreatedAt)
	})
	if err := tx.Commit(); err != nil {
		return nil, err
	}
	return superseded, nil
}

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) TouchCacheEntry(ctx context.Context, id string, now time.Time) error {
	_, err := d.ExecContext(ctx, `
		update cache_entries set last_used_at = ? where id = ? and state = 'ready'`, now.UnixNano(), id)
	return 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 = ?`,
		id, expectedState, expectedLastUsed.UnixNano())
	if err != nil {
		return false, err
	}
	changed, err := result.RowsAffected()
	return changed == 1, err
}

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) {
	ready, err := d.queryCacheEntries(ctx, `
		select `+cacheEntryColumns+`
		from cache_entries
		where state = 'ready' and last_used_at < ?
		order by last_used_at
		limit ?`, readyBefore.UnixNano(), limit)
	if err != nil {
		return nil, err
	}
	recovery, err := d.queryCacheEntries(ctx, `
		select `+cacheEntryColumns+`
		from cache_entries
		where state in ('pending', 'deleting') and created_at < ?
		order by created_at
		limit ?`, pendingBefore.UnixNano(), limit)
	if err != nil {
		return nil, err
	}
	entries := append(ready, recovery...)
	sort.Slice(entries, func(i, j int) bool {
		left, right := entries[i].CreatedAt, entries[j].CreatedAt
		if entries[i].State == "ready" {
			left = entries[i].LastUsedAt
		}
		if entries[j].State == "ready" {
			right = entries[j].LastUsedAt
		}
		return left.Before(right)
	})
	if len(entries) > limit {
		entries = entries[:limit]
	}
	return entries, nil
}

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) DeleteCacheEntry(ctx context.Context, id string) error {
	_, err := d.ExecContext(ctx, `delete from cache_entries where id = ?`, id)
	return err
}

func (d *DB) CacheUsageByOwner(ctx context.Context, ownerDID string) (bytes, count int64, err error) {
	err = d.QueryRowContext(ctx, `
		select coalesce(sum(size_bytes), 0), count(*)
		from cache_entries
		where owner_did = ? and state in ('ready', 'deleting')`, ownerDID).Scan(&bytes, &count)
	return bytes, count, 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.SizeBytes,
		&entry.State,
		&createdAt,
		&lastUsedAt,
	); err != nil {
		return nil, err
	}
	entry.CreatedAt = time.Unix(0, createdAt)
	entry.LastUsedAt = time.Unix(0, lastUsedAt)
	return &entry, nil
}
