package tools

import (
	"crypto/rand"
	"encoding/hex"
	"fmt"
	"io/fs"
	"os"
	"sync"
	"time"

	"agentu/internal/safeio"
)

// RollbackSnapshot captures the state of a file before a mutation.
type RollbackSnapshot struct {
	ID           string
	Tool         string
	Path         string
	ResolvedPath string
	Existed      bool
	Data         []byte
	Mode         fs.FileMode
	CreatedAt    time.Time
}

// RollbackStore manages in-memory rollback snapshots for the current process.
type RollbackStore struct {
	mu        sync.Mutex
	snapshots []RollbackSnapshot
}

// NewRollbackStore creates an empty rollback store.
func NewRollbackStore() *RollbackStore {
	return &RollbackStore{}
}

// generateSnapshotID creates a rollback ID: "rb_" + 12 lowercase hex chars.
func generateSnapshotID() (string, error) {
	b := make([]byte, 6)
	if _, err := rand.Read(b); err != nil {
		return "", fmt.Errorf("generate snapshot id: %w", err)
	}
	return "rb_" + hex.EncodeToString(b), nil
}

// Capture stores the current state of the file at resolvedPath before a tool
// modifies it. Returns the snapshot ID.
func (s *RollbackStore) Capture(tool string, displayPath string, resolvedPath string) (string, error) {
	id, err := generateSnapshotID()
	if err != nil {
		return "", err
	}

	snap := RollbackSnapshot{
		ID:           id,
		Tool:         tool,
		Path:         displayPath,
		ResolvedPath: resolvedPath,
		CreatedAt:    time.Now().UTC(),
	}

	info, err := os.Stat(resolvedPath)
	if err != nil {
		if os.IsNotExist(err) {
			snap.Existed = false
		} else {
			return "", fmt.Errorf("stat file: %w", err)
		}
	} else {
		snap.Existed = true
		snap.Mode = info.Mode().Perm()
		data, err := os.ReadFile(resolvedPath)
		if err != nil {
			return "", fmt.Errorf("read file for snapshot: %w", err)
		}
		snap.Data = data
	}

	s.mu.Lock()
	s.snapshots = append(s.snapshots, snap)
	s.mu.Unlock()
	return id, nil
}

// Restore rolls back a file to the state captured in the given snapshot.
func (s *RollbackStore) Restore(id string) (RollbackSnapshot, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	for i := len(s.snapshots) - 1; i >= 0; i-- {
		if s.snapshots[i].ID == id {
			snap := s.snapshots[i]
			if snap.Existed {
				if err := safeio.AtomicWriteFile(snap.ResolvedPath, snap.Data, snap.Mode); err != nil {
					return RollbackSnapshot{}, fmt.Errorf("restore file: %w", err)
				}
			} else {
				if err := os.Remove(snap.ResolvedPath); err != nil && !os.IsNotExist(err) {
					return RollbackSnapshot{}, fmt.Errorf("remove created file: %w", err)
				}
			}
			s.snapshots = append(s.snapshots[:i], s.snapshots[i+1:]...)
			return snap, nil
		}
	}
	return RollbackSnapshot{}, fmt.Errorf("rollback snapshot %s not found", id)
}

// LastID returns the most recent snapshot ID, or "" if empty.
func (s *RollbackStore) LastID() string {
	s.mu.Lock()
	defer s.mu.Unlock()
	if len(s.snapshots) == 0 {
		return ""
	}
	return s.snapshots[len(s.snapshots)-1].ID
}