package aws

import (
	"encoding/base64"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/tidbcloud/controlplane9/internal/cluster"
)

const reapStateDirName = "reaper"
const reapStateSegmentPrefix = "b64-"

// ReapState persists the authoritative Stage-2 cleanup boundary across reruns.
type ReapState struct {
	Scope        Scope      `json:"scope"`
	Profile      string     `json:"profile"`
	ConfigScoped bool       `json:"config_scoped,omitempty"`
	ConfigDigest string     `json:"config_digest,omitempty"`
	Survivors    []Resource `json:"survivors,omitempty"`
	Completed    bool       `json:"completed,omitempty"`
	UpdatedAt    time.Time  `json:"updated_at"`
}

// ReapStatePath returns the persisted Stage-2 resume path for an owner/session scope.
func ReapStatePath(stateRoot, owner, sessionID string) string {
	return filepath.Join(
		reapStateRoot(stateRoot),
		reapStateOwnerDir(owner),
		reapStateFileName(sessionID),
	)
}

func reapStateRoot(stateRoot string) string {
	return filepath.Join(stateRoot, cluster.InternalStateDirName, reapStateDirName)
}

// previousReapStateRoot is read only for compatibility with state written
// before reaper metadata moved out of the valid "reaper" session root.
func previousReapStateRoot(stateRoot string) string {
	return filepath.Join(stateRoot, reapStateDirName)
}

func previousReapStatePath(stateRoot, owner, sessionID string) string {
	return filepath.Join(
		previousReapStateRoot(stateRoot),
		reapStateOwnerDir(owner),
		reapStateFileName(sessionID),
	)
}

// SaveReapState persists a Stage-2 resume file atomically.
func SaveReapState(stateRoot string, state *ReapState) error {
	if state == nil {
		return fmt.Errorf("missing reap state")
	}
	scope := state.Scope.Normalize()
	if err := scope.Validate(); err != nil {
		return err
	}
	state.Scope = scope
	if state.UpdatedAt.IsZero() {
		state.UpdatedAt = time.Now()
	}

	path := ReapStatePath(stateRoot, scope.Owner, scope.SessionID)
	if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
		return fmt.Errorf("create reap state dir: %w", err)
	}

	data, err := json.MarshalIndent(state, "", "  ")
	if err != nil {
		return fmt.Errorf("marshal reap state: %w", err)
	}

	tmp, err := os.CreateTemp(filepath.Dir(path), "reap-state.*.tmp")
	if err != nil {
		return fmt.Errorf("create reap state temp file: %w", err)
	}
	tmpPath := tmp.Name()
	defer os.Remove(tmpPath)

	if _, err := tmp.Write(data); err != nil {
		_ = tmp.Close()
		return fmt.Errorf("write reap state temp file: %w", err)
	}
	if err := tmp.Close(); err != nil {
		return fmt.Errorf("close reap state temp file: %w", err)
	}
	if err := os.Rename(tmpPath, path); err != nil {
		return fmt.Errorf("rename reap state: %w", err)
	}
	removePreviousReapStates(stateRoot, scope.Owner, scope.SessionID)
	return nil
}

// LoadReapState loads the persisted Stage-2 resume file for the given scope.
func LoadReapState(stateRoot, owner, sessionID string) (*ReapState, error) {
	path, err := firstExistingReapStatePath(stateRoot, owner, sessionID)
	if err != nil {
		return nil, fmt.Errorf("read reap state: %w", err)
	}
	return readReapStateFile(path)
}

// ReapStateExists reports whether the Stage-2 resume file exists.
func ReapStateExists(stateRoot, owner, sessionID string) bool {
	_, err := firstExistingReapStatePath(stateRoot, owner, sessionID)
	return err == nil
}

// RemoveReapState removes the persisted Stage-2 resume file and prunes empty parent dirs.
func RemoveReapState(stateRoot, owner, sessionID string) error {
	var firstErr error
	primary := ReapStatePath(stateRoot, owner, sessionID)
	if err := removeReapStateFile(primary, true); err != nil && !os.IsNotExist(err) {
		firstErr = err
	}
	previous := previousReapStatePath(stateRoot, owner, sessionID)
	if err := removeReapStateFile(previous, false); err != nil && !os.IsNotExist(err) && firstErr == nil {
		firstErr = err
	}

	legacyPath, err := findLegacyReapStatePath(stateRoot, owner, sessionID)
	if err != nil {
		if !os.IsNotExist(err) && firstErr == nil {
			firstErr = err
		}
	} else if err := removeReapStateFile(legacyPath, false); err != nil && !os.IsNotExist(err) && firstErr == nil {
		firstErr = err
	}
	if firstErr != nil {
		return fmt.Errorf("remove reap state: %w", firstErr)
	}
	return nil
}

// FindReapStateBySession finds the unique persisted Stage-2 resume file for a session ID.
func FindReapStateBySession(stateRoot, sessionID string) (*ReapState, error) {
	sessionID = strings.TrimSpace(sessionID)
	targets := map[string]struct{}{
		reapStateFileName(sessionID):       {},
		legacyReapStateFileName(sessionID): {},
	}
	matches := make(map[string]*ReapState)
	for _, root := range []string{reapStateRoot(stateRoot), previousReapStateRoot(stateRoot)} {
		entries, err := os.ReadDir(root)
		if err != nil {
			if os.IsNotExist(err) {
				continue
			}
			return nil, fmt.Errorf("list reap states: %w", err)
		}
		for _, entry := range entries {
			if !entry.IsDir() {
				continue
			}
			ownerDir := filepath.Join(root, entry.Name())
			files, err := os.ReadDir(ownerDir)
			if err != nil {
				return nil, fmt.Errorf("list owner reap states: %w", err)
			}
			for _, file := range files {
				if file.IsDir() {
					continue
				}
				if _, ok := targets[file.Name()]; !ok {
					continue
				}
				state, err := readReapStateFile(filepath.Join(ownerDir, file.Name()))
				if err != nil {
					return nil, err
				}
				if state.Scope.SessionID != sessionID {
					continue
				}
				key := state.Scope.Owner + "\x00" + state.Scope.SessionID
				if _, exists := matches[key]; !exists {
					matches[key] = state
				}
			}
		}
	}

	switch len(matches) {
	case 0:
		return nil, os.ErrNotExist
	case 1:
		for _, state := range matches {
			return state, nil
		}
		return nil, os.ErrNotExist
	default:
		return nil, fmt.Errorf("multiple reap states found for session %q", sessionID)
	}
}

func encodeReapStateSegment(segment string) string {
	segment = strings.TrimSpace(segment)
	if segment == "" {
		return reapStateSegmentPrefix
	}
	return reapStateSegmentPrefix + base64.RawURLEncoding.EncodeToString([]byte(segment))
}

func reapStateOwnerDir(owner string) string {
	return encodeReapStateSegment(owner)
}

func reapStateFileName(sessionID string) string {
	return encodeReapStateSegment(sessionID) + ".json"
}

func legacyReapStateFileName(sessionID string) string {
	return legacyReapStateSegment(sessionID) + ".json"
}

func legacyReapStateSegment(segment string) string {
	segment = strings.TrimSpace(segment)
	segment = strings.ReplaceAll(segment, "/", "_")
	if segment == "" {
		return "_"
	}
	return segment
}

func firstExistingReapStatePath(stateRoot, owner, sessionID string) (string, error) {
	for _, path := range []string{
		ReapStatePath(stateRoot, owner, sessionID),
		previousReapStatePath(stateRoot, owner, sessionID),
	} {
		if _, err := os.Stat(path); err == nil {
			return path, nil
		} else if !os.IsNotExist(err) {
			return "", err
		}
	}
	return findLegacyReapStatePath(stateRoot, owner, sessionID)
}

func readReapStateFile(path string) (*ReapState, error) {
	data, err := os.ReadFile(path)
	if err != nil {
		return nil, fmt.Errorf("read reap state: %w", err)
	}
	var state ReapState
	if err := json.Unmarshal(data, &state); err != nil {
		return nil, fmt.Errorf("parse reap state: %w", err)
	}
	state.Scope = state.Scope.Normalize()
	return &state, nil
}

func removePreviousReapStates(stateRoot, owner, sessionID string) {
	previous := previousReapStatePath(stateRoot, owner, sessionID)
	if err := removeReapStateFile(previous, false); err != nil && !os.IsNotExist(err) {
		return
	}
	legacy, err := findLegacyReapStatePath(stateRoot, owner, sessionID)
	if err != nil {
		return
	}
	_ = removeReapStateFile(legacy, false)
}

func findLegacyReapStatePath(stateRoot, owner, sessionID string) (string, error) {
	root := previousReapStateRoot(stateRoot)
	ownerEntries, err := os.ReadDir(root)
	if err != nil {
		if os.IsNotExist(err) {
			return "", os.ErrNotExist
		}
		return "", err
	}

	targetOwner := legacyReapStateSegment(owner)
	targetFile := legacyReapStateFileName(sessionID)
	for _, ownerEntry := range ownerEntries {
		if !ownerEntry.IsDir() || ownerEntry.Name() != targetOwner {
			continue
		}
		ownerDir := filepath.Join(root, ownerEntry.Name())
		files, err := os.ReadDir(ownerDir)
		if err != nil {
			return "", err
		}
		for _, file := range files {
			if file.IsDir() || file.Name() != targetFile {
				continue
			}
			return filepath.Join(ownerDir, file.Name()), nil
		}
	}
	return "", os.ErrNotExist
}

func removeReapStateFile(path string, pruneRoot bool) error {
	if err := os.Remove(path); err != nil {
		return err
	}
	pruneReapStateParents(path, pruneRoot)
	return nil
}

func pruneReapStateParents(path string, pruneRoot bool) {
	parent := filepath.Dir(path)
	_ = os.Remove(parent)
	if pruneRoot {
		_ = os.Remove(filepath.Dir(parent))
	}
}