package aws

import (
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"time"
)

const stateFile = "aws-state.json"

// SessionState tracks AWS resources for a cp9 session.
type SessionState struct {
	SessionID          string            `json:"session_id"`
	Owner              string            `json:"owner"`
	Profile            string            `json:"profile"`
	SourceProfile      string            `json:"source_profile,omitempty"`
	Region             string            `json:"region"`
	Account            string            `json:"account"`
	RoleARN            string            `json:"role_arn,omitempty"`
	CredentialSource   string            `json:"credential_source,omitempty"`
	CredentialExpiry   string            `json:"credential_expiry,omitempty"`
	SessionTags        map[string]string `json:"session_tags,omitempty"`
	CredentialsSecret  string            `json:"credentials_secret,omitempty"`
	ProviderConfigName string            `json:"provider_config_name,omitempty"`
	CreatedAt          time.Time         `json:"created_at"`
}

// StatePath returns the path to the AWS state file.
func StatePath(stateDir string) string {
	return filepath.Join(stateDir, stateFile)
}

// SaveSessionState writes AWS session state to disk atomically.
func SaveSessionState(stateDir string, state *SessionState) error {
	if err := os.MkdirAll(stateDir, 0o755); err != nil {
		return fmt.Errorf("create state dir: %w", err)
	}

	if state.CreatedAt.IsZero() {
		state.CreatedAt = time.Now()
	}

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

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

	if _, err := tmpFile.Write(data); err != nil {
		_ = tmpFile.Close()
		return fmt.Errorf("write aws state temp: %w", err)
	}
	if err := tmpFile.Close(); err != nil {
		return fmt.Errorf("close aws state temp: %w", err)
	}

	target := StatePath(stateDir)
	if err := os.Rename(tmpPath, target); err != nil {
		return fmt.Errorf("rename aws state file: %w", err)
	}

	return nil
}

// LoadSessionState reads AWS session state from disk.
func LoadSessionState(stateDir string) (*SessionState, error) {
	data, err := os.ReadFile(StatePath(stateDir))
	if err != nil {
		return nil, fmt.Errorf("read aws state file: %w", err)
	}

	var state SessionState
	if err := json.Unmarshal(data, &state); err != nil {
		return nil, fmt.Errorf("parse aws state file: %w", err)
	}
	return &state, nil
}

// SessionStateExists checks if an AWS state file exists.
func SessionStateExists(stateDir string) bool {
	_, err := os.Stat(StatePath(stateDir))
	return err == nil
}

// RemoveSessionState removes the AWS state file.
func RemoveSessionState(stateDir string) error {
	path := StatePath(stateDir)
	if err := os.Remove(path); err != nil && !os.IsNotExist(err) {
		return fmt.Errorf("remove aws state file: %w", err)
	}
	return nil
}

// LoadAllSessionStates reads every session state file under the cp9 state root.
// Legacy flat state at the root is included when present.
func LoadAllSessionStates(stateRoot string) ([]SessionState, error) {
	var states []SessionState

	if SessionStateExists(stateRoot) {
		state, err := LoadSessionState(stateRoot)
		if err != nil {
			return nil, err
		}
		states = append(states, *state)
	}

	entries, err := os.ReadDir(stateRoot)
	if err != nil {
		if os.IsNotExist(err) {
			return states, nil
		}
		return nil, fmt.Errorf("read state root: %w", err)
	}

	for _, entry := range entries {
		if !entry.IsDir() {
			continue
		}
		stateDir := filepath.Join(stateRoot, entry.Name())
		if !SessionStateExists(stateDir) {
			continue
		}
		state, err := LoadSessionState(stateDir)
		if err != nil {
			return nil, err
		}
		states = append(states, *state)
	}

	return states, nil
}