package db

import (
	"context"
	"database/sql"
	"encoding/json"
	"slices"
	"strings"
	"time"

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

var _ quota.NixUseRecorder = (*QuotaStore)(nil)

const (
	nixUseBatch  = 500
	nixDropBatch = 500
)

type NixObjectRow struct {
	Cache, StoreHash, RepoDID, OwnerDID string
	NarSize                             int64
	LastUsedAt                          time.Time
	Gen                                 int64
}

type NixEvictFilter struct {
	RepoDID, OwnerDID string
}

const nixObjectColumns = `cache, store_hash, repo_did, owner_did, nar_size, last_used_at, gen`

func scanNixObjects(rows *sql.Rows) ([]NixObjectRow, error) {
	defer rows.Close()
	var out []NixObjectRow
	for rows.Next() {
		var r NixObjectRow
		var lastUsed int64
		if err := rows.Scan(&r.Cache, &r.StoreHash, &r.RepoDID, &r.OwnerDID, &r.NarSize, &lastUsed, &r.Gen); err != nil {
			return nil, err
		}
		r.LastUsedAt = time.Unix(lastUsed, 0)
		out = append(out, r)
	}
	return out, rows.Err()
}

func (d *QuotaStore) upsertNixObject(ctx context.Context, conn execQuerier, repoDID, ownerDID, encoded string) error {
	var obj quota.NixObject
	if err := json.Unmarshal([]byte(encoded), &obj); err != nil {
		return err
	}
	now := d.now().Unix()
	if _, err := conn.ExecContext(ctx, `
		INSERT INTO nix_objects (cache, store_hash, repo_did, owner_did, store_path, nar_size, seq, last_used_at, state, gen)
		VALUES (?, ?, ?, ?, ?, ?, (SELECT coalesce(max(seq), 0) + 1 FROM nix_objects), ?, 'live', 1)
		ON CONFLICT(cache, store_hash) DO UPDATE SET
			repo_did = excluded.repo_did,
			owner_did = excluded.owner_did,
			store_path = excluded.store_path,
			nar_size = excluded.nar_size,
			last_used_at = excluded.last_used_at,
			state = 'live',
			gen = gen + 1
	`, obj.Cache, obj.StoreHash, repoDID, ownerDID, obj.StorePath, obj.NarSize, now); err != nil {
		return err
	}
	if _, err := conn.ExecContext(ctx, `DELETE FROM nix_refs WHERE cache = ? AND from_hash = ?`, obj.Cache, obj.StoreHash); err != nil {
		return err
	}
	for _, ref := range obj.References {
		hash, _, _ := strings.Cut(ref, "-")
		if hash == "" || hash == obj.StoreHash {
			continue
		}
		if _, err := conn.ExecContext(ctx, `
			INSERT OR IGNORE INTO nix_refs (cache, from_hash, to_hash) VALUES (?, ?, ?)
		`, obj.Cache, obj.StoreHash, hash); err != nil {
			return err
		}
		// bumping gen makes the pending finish skip it. if attic already deleted it, the reference dangles and nix builds it
		if _, err := conn.ExecContext(ctx, `
			UPDATE nix_objects SET state = 'live', gen = gen + 1
			WHERE cache = ? AND store_hash = ? AND state = 'evicting'
		`, obj.Cache, hash); err != nil {
			return err
		}
	}
	return nil
}

func (d *QuotaStore) RecordNixUse(ctx context.Context, cache string, storeHashes []string) error {
	return d.withTx(ctx, func(conn *sql.Conn) error {
		now := d.now().Unix()
		for batch := range slices.Chunk(storeHashes, nixUseBatch) {
			args := []any{now, cache}
			for _, hash := range batch {
				args = append(args, hash)
			}
			if _, err := conn.ExecContext(ctx, `
				UPDATE nix_objects SET last_used_at = max(last_used_at, ?)
				WHERE cache = ? AND store_hash IN (`+placeholders(len(batch))+`)
			`, args...); err != nil {
				return err
			}
		}
		return nil
	})
}

// NixCacheUsage counts evicting rows too because their bytes stay in the cache until the delete lands
func (d *QuotaStore) NixCacheUsage(ctx context.Context) (byRepo, byOwner map[string]int64, total int64, err error) {
	byRepo, byOwner = map[string]int64{}, map[string]int64{}
	err = d.withReadTx(ctx, func(tx *sql.Tx) error {
		rows, err := tx.QueryContext(ctx, `SELECT repo_did, owner_did, SUM(nar_size) FROM nix_objects GROUP BY repo_did, owner_did`)
		if err != nil {
			return err
		}
		defer rows.Close()
		for rows.Next() {
			var repo, owner string
			var size int64
			if err := rows.Scan(&repo, &owner, &size); err != nil {
				return err
			}
			byRepo[repo] += size
			byOwner[owner] += size
			total += size
		}
		return rows.Err()
	})
	return byRepo, byOwner, total, err
}

// a referrer pins its dependencies only if committed after them. a cycle can't pin itself that way
// an out of order upload can leave a dangling reference that nix builds
const noReferrer = `NOT EXISTS (
	SELECT 1 FROM nix_refs r
	JOIN nix_objects p ON p.cache = r.cache AND p.store_hash = r.from_hash AND p.seq > o.seq
	WHERE r.cache = o.cache AND r.to_hash = o.store_hash
)`

func (d *QuotaStore) NixEvictable(ctx context.Context, f NixEvictFilter, usedBefore time.Time, limit int) ([]NixObjectRow, error) {
	var out []NixObjectRow
	err := d.withReadTx(ctx, func(tx *sql.Tx) error {
		where, args := `state = 'live' AND last_used_at < ?`, []any{usedBefore.Unix()}
		if f.RepoDID != "" {
			where += ` AND repo_did = ?`
			args = append(args, f.RepoDID)
		}
		if f.OwnerDID != "" {
			where += ` AND owner_did = ?`
			args = append(args, f.OwnerDID)
		}
		rows, err := tx.QueryContext(ctx, `
			SELECT `+nixObjectColumns+` FROM nix_objects o
			WHERE `+where+` AND `+noReferrer+`
			ORDER BY last_used_at, cache, store_hash
			LIMIT ?
		`, append(args, limit)...)
		if err != nil {
			return err
		}
		out, err = scanNixObjects(rows)
		return err
	})
	return out, err
}

func (d *QuotaStore) MarkNixEvicting(ctx context.Context, rows []NixObjectRow, usedBefore time.Time) ([]NixObjectRow, error) {
	var marked []NixObjectRow
	err := d.withTx(ctx, func(conn *sql.Conn) error {
		marked = nil
		for _, r := range rows {
			res, err := conn.ExecContext(ctx, `
				UPDATE nix_objects AS o SET state = 'evicting'
				WHERE cache = ? AND store_hash = ? AND state = 'live' AND gen = ? AND last_used_at < ? AND `+noReferrer,
				r.Cache, r.StoreHash, r.Gen, usedBefore.Unix())
			if err != nil {
				return err
			}
			if n, err := res.RowsAffected(); err != nil {
				return err
			} else if n == 1 {
				marked = append(marked, r)
			}
		}
		return nil
	})
	if err != nil {
		return nil, err
	}
	return marked, nil
}

func (d *QuotaStore) NixEvicting(ctx context.Context) ([]NixObjectRow, error) {
	var out []NixObjectRow
	err := d.withReadTx(ctx, func(tx *sql.Tx) error {
		rows, err := tx.QueryContext(ctx, `
			SELECT `+nixObjectColumns+` FROM nix_objects WHERE state = 'evicting' ORDER BY cache, store_hash
		`)
		if err != nil {
			return err
		}
		out, err = scanNixObjects(rows)
		return err
	})
	return out, err
}

// FinishNixEviction skips a row that was uploaded again, because that bumps its gen
// one tx per chunk, and a failure leaves the rest evicting for the next pass
func (d *QuotaStore) FinishNixEviction(ctx context.Context, cache string, gone []NixObjectRow) error {
	for chunk := range slices.Chunk(gone, nixDropBatch) {
		err := d.withTx(ctx, func(conn *sql.Conn) error {
			for _, r := range chunk {
				res, err := conn.ExecContext(ctx, `
					DELETE FROM nix_objects WHERE cache = ? AND store_hash = ? AND state = 'evicting' AND gen = ?`,
					cache, r.StoreHash, r.Gen)
				if err != nil {
					return err
				}
				if n, err := res.RowsAffected(); err != nil {
					return err
				} else if n != 1 {
					continue
				}
				if _, err := conn.ExecContext(ctx, `DELETE FROM nix_refs WHERE cache = ? AND from_hash = ?`, cache, r.StoreHash); err != nil {
					return err
				}
				if _, err := conn.ExecContext(ctx, `DELETE FROM quota_allocations WHERE repo_did = ? AND kind = ? AND key = ?`,
					r.RepoDID, quota.KindNixCache, quota.NixCacheKey(cache, r.StoreHash)); err != nil {
					return err
				}
			}
			return nil
		})
		if err != nil {
			return err
		}
	}
	return nil
}

func placeholders(n int) string {
	return strings.TrimSuffix(strings.Repeat("?,", n), ",")
}

func (d *QuotaStore) RecordNixGrant(ctx context.Context, expires time.Time, life time.Duration) error {
	return d.withTx(ctx, func(conn *sql.Conn) error {
		if _, err := conn.ExecContext(ctx, `DELETE FROM nix_grants WHERE expires_at <= ?`, d.now().Unix()); err != nil {
			return err
		}
		_, err := conn.ExecContext(ctx, `INSERT INTO nix_grants (expires_at, life) VALUES (?, ?)`,
			expires.Add(time.Second-1).Unix(), int64((life+time.Second-1)/time.Second))
		return err
	})
}

func (d *QuotaStore) LongestLiveNixGrant(ctx context.Context, now time.Time) (time.Duration, error) {
	var life int64
	err := d.withReadTx(ctx, func(tx *sql.Tx) error {
		return tx.QueryRowContext(ctx, `SELECT coalesce(max(life), 0) FROM nix_grants WHERE expires_at > ?`, now.Unix()).Scan(&life)
	})
	return time.Duration(life) * time.Second, err
}