package engine

import (
	"bufio"
	"bytes"
	"context"
	"crypto/sha256"
	"database/sql"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"os/exec"
	"strings"
	"sync"
	"time"

	"github.com/google/uuid"
	"tangled.org/core/spindle/db"

	"tangled.org/core/spindle/models"
	"tangled.org/core/spindle/storage"
)

// cache log steps live below the setup step (-1)
const (
	CacheRestoreStepIdx = -2
	CacheSaveStepIdx    = -3
)

type cacheStep struct {
	name    string
	command string
}

func (s cacheStep) Name() string          { return s.name }
func (s cacheStep) Command() string       { return s.command }
func (s cacheStep) Kind() models.StepKind { return models.StepKindSystem }

var (
	CacheRestoreStep models.Step = cacheStep{name: "restore cache", command: "restore cached paths"}
	CacheSaveStep    models.Step = cacheStep{name: "save cache", command: "persist changed paths"}
)

type CacheRunner interface {
	RestoreCache(ctx context.Context, wid models.WorkflowId, wf *models.Workflow, store storage.Storage, caches []ResolvedCache, wfLogger models.WorkflowLogger) error
	SaveCache(ctx context.Context, wid models.WorkflowId, wf *models.Workflow, store storage.Storage, caches []ResolvedCache, wfLogger models.WorkflowLogger) error
}

type ResolvedCache struct {
	Paths            []string
	Key              string
	Hash             string
	SaveKey          string
	RestoreID        string
	RestoreKey       string
	RestoreName      string
	CompressionLevel int
	When             string
}

func (rc ResolvedCache) saveOn(failed bool) bool {
	return rc.When == "always" || !failed
}

const CacheExitNoPaths = 42

// avoids storing an empty archive when zstd is missing
const CacheExitNoCompressor = 43

func AbsolutizePaths(paths []string, workspaceRoot string) []string {
	abs := make([]string, 0, len(paths))
	for _, p := range paths {
		if !strings.HasPrefix(p, "/") {
			p = workspaceRoot + "/" + p
		}
		abs = append(abs, p)
	}
	return abs
}

// run tar from / so absolute paths survive extraction
// tar exits nonzero on missing paths, so only existing ones reach it
func CacheSaveScript(paths []string, workspaceRoot string, compressionLevel int) string {
	trimmed := make([]string, 0, len(paths))
	for _, p := range AbsolutizePaths(paths, workspaceRoot) {
		trimmed = append(trimmed, strings.TrimPrefix(p, "/"))
	}
	tail := fmt.Sprintf(`tar -cf - -C / "$@" | %s`, CacheCompressCmd(compressionLevel))
	return fmt.Sprintf(`set -o pipefail
command -v zstd >/dev/null 2>&1 || { echo "zstd not found in image; cannot save cache" >&2; exit %d; }
set --
for p in %s; do [ -e "/$p" ] && set -- "$@" "$p"; done
if [ $# -eq 0 ]; then echo "no cache paths exist; skipping" >&2; exit %d; fi
%s`, CacheExitNoCompressor, strings.Join(trimmed, " "), CacheExitNoPaths, tail)
}

func CacheCompressCmd(level int) string {
	if level == 0 {
		return "zstd -T0 -5"
	}
	return fmt.Sprintf("zstd -T0 -%d", level)
}

// old entries might still be gzip, so detect them instead of trusting config
func CacheDecompressCmd(br *bufio.Reader) string {
	head, _ := br.Peek(4)
	if bytes.HasPrefix(head, []byte{0x1f, 0x8b}) {
		return "gzip -dc"
	}
	return "zstd -dc"
}

// Put keeps draining the pipe after it returns, so the guest writer never
// blocks on a full pipe.
type CacheUpload struct {
	Writer *io.PipeWriter
	done   chan error
}

func NewCacheUpload(ctx context.Context, store storage.Storage, key string) *CacheUpload {
	pr, pw := io.Pipe()
	u := &CacheUpload{Writer: pw, done: make(chan error, 1)}
	go func() {
		err := store.Put(ctx, key, pr)
		_, _ = io.Copy(io.Discard, pr)
		u.done <- err
	}()
	return u
}

func (u *CacheUpload) Abort(err error) {
	u.Writer.CloseWithError(err)
	<-u.done
}

// storage treats EOF as a complete archive, so only a clean exec may Finish
func (u *CacheUpload) Finish() error {
	u.Writer.Close()
	return <-u.done
}

type indexedCacheStore struct {
	storage.Storage
	index   *db.DB
	logger  *slog.Logger
	mu      sync.Mutex
	entries map[string]string // storage key -> cache entry id
}

func (s *indexedCacheStore) Get(ctx context.Context, key string) (io.ReadCloser, error) {
	s.mu.Lock()
	id, ok := s.entries[key]
	s.mu.Unlock()
	r, err := s.Storage.Get(ctx, key)
	if err != nil {
		if ok && errors.Is(err, storage.ErrNotExist) {
			_ = s.index.DeleteCacheEntry(context.WithoutCancel(ctx), id)
		}
		return nil, err
	}
	if ok {
		if err := s.index.TouchCacheEntry(ctx, id, time.Now()); err != nil {
			s.logger.Warn("cache usage update failed", "id", id, "err", err)
		}
	}
	return r, nil
}

func (s *indexedCacheStore) Put(ctx context.Context, key string, r io.Reader) error {
	s.mu.Lock()
	id, ok := s.entries[key]
	s.mu.Unlock()
	if !ok {
		return fmt.Errorf("cache metadata missing for %q", key)
	}
	counted := &countingReader{r: r}
	if err := s.Storage.Put(ctx, key, counted); err != nil {
		return err
	}
	superseded, err := s.index.MarkCacheEntryReady(ctx, id, counted.n, time.Now())
	if err != nil {
		_ = s.Storage.Delete(context.WithoutCancel(ctx), key)
		return fmt.Errorf("mark cache ready: %w", err)
	}
	// completed saves leave the pending map so cleanup keeps their object
	s.mu.Lock()
	delete(s.entries, key)
	s.mu.Unlock()
	cleanupCtx := context.WithoutCancel(ctx)
	for _, old := range superseded {
		s.deleteEntry(cleanupCtx, old.StorageKey, old.ID, "replaced")
	}
	return nil
}

func (s *indexedCacheStore) cleanup(ctx context.Context) {
	s.mu.Lock()
	defer s.mu.Unlock()
	for key, id := range s.entries {
		s.deleteEntry(ctx, key, id, "incomplete")
	}
}

func (s *indexedCacheStore) deleteEntry(ctx context.Context, key, id, reason string) {
	if err := s.Storage.Delete(ctx, key); err != nil {
		s.logger.Warn("delete "+reason+" cache failed", "key", key, "err", err)
		return
	}
	if err := s.index.DeleteCacheEntry(ctx, id); err != nil {
		s.logger.Warn("delete "+reason+" cache metadata failed", "id", id, "err", err)
	}
}

type countingReader struct {
	r io.Reader
	n int64
}

func (r *countingReader) Read(p []byte) (int, error) {
	n, err := r.r.Read(p)
	r.n += int64(n)
	return n, err
}

func cacheStoreForRestore(base storage.Storage, index *db.DB, l *slog.Logger, entries []ResolvedCache) storage.Storage {
	mapped := make(map[string]string, len(entries))
	for _, entry := range entries {
		if entry.RestoreKey != "" {
			mapped[entry.RestoreKey] = entry.RestoreID
		}
	}
	return &indexedCacheStore{Storage: base, index: index, logger: l, entries: mapped}
}

func prepareCacheSaves(ctx context.Context, base storage.Storage, index *db.DB, l *slog.Logger, ownerDID, repoDID, engineName string, entries []ResolvedCache) (*indexedCacheStore, error) {
	saveStore := &indexedCacheStore{
		Storage: base,
		index:   index,
		logger:  l,
		entries: make(map[string]string, len(entries)),
	}
	now := time.Now()
	for i := range entries {
		id := uuid.NewString()
		entries[i].SaveKey = "objects/" + id
		if err := index.InsertCacheEntry(ctx, db.CacheEntry{
			ID:         id,
			StorageKey: entries[i].SaveKey,
			OwnerDID:   ownerDID,
			RepoDID:    repoDID,
			Engine:     engineName,
			CacheKey:   entries[i].Key,
			CacheHash:  entries[i].Hash,
			State:      "pending",
			CreatedAt:  now,
			LastUsedAt: now,
		}); err != nil {
			saveStore.cleanup(context.WithoutCancel(ctx))
			return nil, err
		}
		saveStore.entries[entries[i].SaveKey] = id
	}
	return saveStore, nil
}

// on a hash miss, the newest older generation still warms the build
// unusable entries degrade to a plain miss
func ResolveCaches(ctx context.Context, l *slog.Logger, index *db.DB, repoDID, engine, repoPath, rev string, entries []models.CacheEntry) []ResolvedCache {
	resolved := make([]ResolvedCache, 0, len(entries))
	for _, entry := range entries {
		hash := ""
		if len(entry.Hash) > 0 && repoPath != "" {
			sum, missing := hashKeyFiles(ctx, repoPath, rev, entry.Hash)
			for _, m := range missing {
				l.Warn("cache hash file not in repo", "key", entry.Key, "path", m)
			}
			hash = sum
		}
		rc := ResolvedCache{
			Paths:            entry.Paths,
			Key:              entry.Key,
			Hash:             hash,
			CompressionLevel: entry.CompressionLevel,
			When:             entry.When,
		}

		found, err := index.FindCacheEntry(ctx, repoDID, engine, entry.Key, hash)
		if errors.Is(err, sql.ErrNoRows) && hash != "" {
			found, err = index.FindFallbackCacheEntry(ctx, repoDID, engine, entry.Key, hash)
			if err == nil {
				rc.RestoreName = entry.Key + "-" + found.CacheHash
			}
		}
		if err != nil && !errors.Is(err, sql.ErrNoRows) {
			l.Warn("cache lookup failed; entry will save but not restore", "key", entry.Key, "err", err)
		} else if err == nil {
			rc.RestoreID = found.ID
			rc.RestoreKey = found.StorageKey
		}
		resolved = append(resolved, rc)
	}
	return resolved
}

func hashKeyFiles(ctx context.Context, repoPath, rev string, paths []string) (string, []string) {
	h := sha256.New()
	var missing []string
	hashed := 0
	for _, p := range paths {
		blob, err := gitBlobId(ctx, repoPath, rev, p)
		if err != nil {
			missing = append(missing, p)
			continue
		}
		fmt.Fprintf(h, "%s=%s\n", p, blob)
		hashed++
	}
	if hashed == 0 {
		return "", missing
	}
	return hex.EncodeToString(h.Sum(nil))[:12], missing
}

// sparse checkouts might not have the file, but the object database does
func gitBlobId(ctx context.Context, repoPath, rev, path string) (string, error) {
	out, err := exec.CommandContext(ctx, "git", "-C", repoPath, "rev-parse", rev+":"+path).Output()
	if err != nil {
		return "", err
	}
	return strings.TrimSpace(string(out)), nil
}