package main

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"
	"syscall"
	"time"

	"github.com/spf13/cobra"
	"golang.org/x/term"
)

const (
	DefaultAlgorithm = "AES-256-GCM"
	AdminKeyDir      = ".pass_mgr"
	AdminKeyFile     = ".priv"
	SecretsDirName   = ".secrets"
)

// EncryptedData represents the structure of encrypted password file
type EncryptedData struct {
	Algorithm string `json:"algorithm"`
	Nonce     string `json:"nonce"`
	Data      string `json:"data"`
	User      string `json:"user,omitempty"`
}

// Config holds admin key configuration
type Config struct {
	KeyHash   string `json:"key_hash"`
	Algorithm string `json:"algorithm"`
}

var (
	workspaceDir string
	agentID      string
)

func main() {
	rootCmd := &cobra.Command{
		Use:   "pass_mgr",
		Short: "Password manager for OpenClaw agents",
		Long:  `A secure password management tool using AES-256-GCM encryption.`,
	}

	// Get environment variables
	workspaceDir = os.Getenv("AGENT_WORKSPACE")
	agentID = os.Getenv("AGENT_ID")

	// Commands
	rootCmd.AddCommand(getCmd())
	rootCmd.AddCommand(generateCmd())
	rootCmd.AddCommand(unsetCmd())
	rootCmd.AddCommand(rotateCmd())
	rootCmd.AddCommand(adminInitCmd())
	rootCmd.AddCommand(setCmd())

	if err := rootCmd.Execute(); err != nil {
		fmt.Fprintf(os.Stderr, "Error: %v\n", err)
		os.Exit(1)
	}
}

func getCmd() *cobra.Command {
	var showUsername bool
	cmd := &cobra.Command{
		Use:   "get [key]",
		Short: "Get password for a key",
		Args:  cobra.ExactArgs(1),
		Run: func(cmd *cobra.Command, args []string) {
			key := args[0]
			password, user, err := getPassword(key)
			if err != nil {
				fmt.Fprintf(os.Stderr, "Error: %v\n", err)
				os.Exit(1)
			}
			if showUsername {
				fmt.Println(user)
			} else {
				fmt.Println(password)
			}
		},
	}
	cmd.Flags().BoolVar(&showUsername, "username", false, "Show username instead of password")
	return cmd
}

func generateCmd() *cobra.Command {
	var user string
	cmd := &cobra.Command{
		Use:   "generate [key]",
		Short: "Generate a new password",
		Args:  cobra.ExactArgs(1),
		Run: func(cmd *cobra.Command, args []string) {
			key := args[0]
			// Check if agent is trying to set password
			if os.Getenv("AGENT") != "" || os.Getenv("AGENT_WORKSPACE") != "" {
				fmt.Fprintln(os.Stderr, "Error: Agents cannot set passwords. Use generate instead.")
				os.Exit(1)
			}
			password, err := generatePassword(32)
			if err != nil {
				fmt.Fprintf(os.Stderr, "Error: %v\n", err)
				os.Exit(1)
			}
			if err := setPassword(key, user, password); err != nil {
				fmt.Fprintf(os.Stderr, "Error: %v\n", err)
				os.Exit(1)
			}
			fmt.Println(password)
		},
	}
	cmd.Flags().StringVar(&user, "username", "", "Username associated with the password")
	return cmd
}

func unsetCmd() *cobra.Command {
	return &cobra.Command{
		Use:   "unset [key]",
		Short: "Remove a password",
		Args:  cobra.ExactArgs(1),
		Run: func(cmd *cobra.Command, args []string) {
			key := args[0]
			if err := removePassword(key); err != nil {
				fmt.Fprintf(os.Stderr, "Error: %v\n", err)
				os.Exit(1)
			}
		},
	}
}

func rotateCmd() *cobra.Command {
	return &cobra.Command{
		Use:   "rotate [key]",
		Short: "Rotate password for a key",
		Args:  cobra.ExactArgs(1),
		Run: func(cmd *cobra.Command, args []string) {
			key := args[0]
			// Check if initialized
			if !isInitialized() {
				fmt.Fprintln(os.Stderr, "Error: pass_mgr not initialized. Run 'pass_mgr admin init' first.")
				os.Exit(1)
			}
			
			// Get current user if exists
			_, user, err := getPassword(key)
			if err != nil {
				fmt.Fprintf(os.Stderr, "Error: %v\n", err)
				os.Exit(1)
			}
			
			// Generate new password
			newPassword, err := generatePassword(32)
			if err != nil {
				fmt.Fprintf(os.Stderr, "Error: %v\n", err)
				os.Exit(1)
			}
			
			if err := setPassword(key, user, newPassword); err != nil {
				fmt.Fprintf(os.Stderr, "Error: %v\n", err)
				os.Exit(1)
			}
			fmt.Println(newPassword)
		},
	}
}

func adminInitCmd() *cobra.Command {
	return &cobra.Command{
		Use:   "admin init",
		Short: "Initialize pass_mgr with admin key",
		Run: func(cmd *cobra.Command, args []string) {
			if err := initAdminInteractive(); err != nil {
				fmt.Fprintf(os.Stderr, "Error: %v\n", err)
				os.Exit(1)
			}
			fmt.Println("pass_mgr initialized successfully")
		},
	}
}

func setCmd() *cobra.Command {
	var user string
	cmd := &cobra.Command{
		Use:   "set [key] [password]",
		Short: "Set password (admin only)",
		Args:  cobra.ExactArgs(2),
		Run: func(cmd *cobra.Command, args []string) {
			// Check if agent is trying to set password
			if os.Getenv("AGENT") != "" || os.Getenv("AGENT_WORKSPACE") != "" {
				fmt.Fprintln(os.Stderr, "Error: Agents cannot set passwords. Only humans can use 'set'.")
				os.Exit(1)
			}
			
			key := args[0]
			password := args[1]
			
			if err := setPassword(key, user, password); err != nil {
				fmt.Fprintf(os.Stderr, "Error: %v\n", err)
				os.Exit(1)
			}
		},
	}
	cmd.Flags().StringVar(&user, "username", "", "Username associated with the password")
	return cmd
}

// Helper functions

func getHomeDir() string {
	home, err := os.UserHomeDir()
	if err != nil {
		return "."
	}
	return home
}

func getAdminKeyPath() string {
	return filepath.Join(getHomeDir(), AdminKeyDir, AdminKeyFile)
}

func getConfigPath() string {
	return filepath.Join(getHomeDir(), AdminKeyDir, "config.json")
}

func isInitialized() bool {
	_, err := os.Stat(getConfigPath())
	return err == nil
}

func loadAdminKey() ([]byte, error) {
	keyPath := getAdminKeyPath()
	key, err := os.ReadFile(keyPath)
	if err != nil {
		return nil, fmt.Errorf("failed to load admin key: %w", err)
	}
	// Hash the key to get 32 bytes for AES-256
	hash := sha256.Sum256(key)
	return hash[:], nil
}

func initAdminInteractive() error {
	fmt.Print("Enter admin password: ")
	password1, err := term.ReadPassword(int(syscall.Stdin))
	if err != nil {
		return fmt.Errorf("failed to read password: %w", err)
	}
	fmt.Println()

	// Trim whitespace/newlines
	password1 = []byte(strings.TrimSpace(string(password1)))

	// Validate password length
	if len(password1) < 6 {
		return fmt.Errorf("password must be at least 6 characters long (got %d)", len(password1))
	}

	fmt.Print("Confirm admin password: ")
	password2, err := term.ReadPassword(int(syscall.Stdin))
	if err != nil {
		return fmt.Errorf("failed to read password confirmation: %w", err)
	}
	fmt.Println()

	// Trim whitespace/newlines
	password2 = []byte(strings.TrimSpace(string(password2)))

	// Check passwords match
	if string(password1) != string(password2) {
		return fmt.Errorf("passwords do not match")
	}

	// Save the key
	return saveAdminKey(password1)
}

func saveAdminKey(key []byte) error {
	homeDir := getHomeDir()
	adminDir := filepath.Join(homeDir, AdminKeyDir)

	// Create admin directory
	if err := os.MkdirAll(adminDir, 0700); err != nil {
		return fmt.Errorf("failed to create admin directory: %w", err)
	}

	// Save key
	keyFile := filepath.Join(adminDir, AdminKeyFile)
	if err := os.WriteFile(keyFile, key, 0600); err != nil {
		return fmt.Errorf("failed to save key: %w", err)
	}

	// Save config
	config := Config{
		KeyHash:   fmt.Sprintf("%x", sha256.Sum256(key)),
		Algorithm: DefaultAlgorithm,
	}
	configData, _ := json.MarshalIndent(config, "", "  ")
	configPath := filepath.Join(adminDir, "config.json")
	if err := os.WriteFile(configPath, configData, 0600); err != nil {
		return fmt.Errorf("failed to save config: %w", err)
	}

	return nil
}

func initAdmin(keyPath string) error {
	// Read provided key
	key, err := os.ReadFile(keyPath)
	if err != nil {
		return fmt.Errorf("failed to read key file: %w", err)
	}

	// Validate password length (must be >= 6 characters)
	if len(key) < 6 {
		return fmt.Errorf("password must be at least 6 characters long (got %d)", len(key))
	}

	return saveAdminKey(key)
}

func getSecretsDir() string {
	if workspaceDir != "" && agentID != "" {
		return filepath.Join(workspaceDir, SecretsDirName, agentID)
	}
	// Fallback to home directory
	return filepath.Join(getHomeDir(), SecretsDirName, "default")
}

func getPasswordFilePath(key string) string {
	return filepath.Join(getSecretsDir(), key+".gpg")
}

func encrypt(plaintext []byte, key []byte) (*EncryptedData, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return nil, err
	}
	
	nonce := make([]byte, gcm.NonceSize())
	if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
		return nil, err
	}
	
	ciphertext := gcm.Seal(nonce, nonce, plaintext, nil)
	
	return &EncryptedData{
		Algorithm: DefaultAlgorithm,
		Nonce:     base64.StdEncoding.EncodeToString(nonce),
		Data:      base64.StdEncoding.EncodeToString(ciphertext[gcm.NonceSize():]),
	}, nil
}

func decrypt(data *EncryptedData, key []byte) ([]byte, error) {
	ciphertext, err := base64.StdEncoding.DecodeString(data.Data)
	if err != nil {
		return nil, err
	}
	
	nonce, err := base64.StdEncoding.DecodeString(data.Nonce)
	if err != nil {
		return nil, err
	}
	
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return nil, err
	}
	
	plaintext, err := gcm.Open(nil, nonce, ciphertext, nil)
	if err != nil {
		return nil, err
	}
	
	return plaintext, nil
}

func setPassword(key, user, password string) error {
	if !isInitialized() {
		return fmt.Errorf("pass_mgr not initialized. Run 'pass_mgr admin init' first")
	}
	
	adminKey, err := loadAdminKey()
	if err != nil {
		return err
	}
	
	// Create secrets directory
	secretsDir := getSecretsDir()
	if err := os.MkdirAll(secretsDir, 0700); err != nil {
		return fmt.Errorf("failed to create secrets directory: %w", err)
	}
	
	// Encrypt password
	data := map[string]string{
		"password": password,
		"user":     user,
	}
	plaintext, _ := json.Marshal(data)
	
	encrypted, err := encrypt(plaintext, adminKey)
	if err != nil {
		return fmt.Errorf("failed to encrypt: %w", err)
	}
	encrypted.User = user
	
	// Save to file
	filePath := getPasswordFilePath(key)
	fileData, _ := json.MarshalIndent(encrypted, "", "  ")
	if err := os.WriteFile(filePath, fileData, 0600); err != nil {
		return fmt.Errorf("failed to save password: %w", err)
	}
	
	return nil
}

func getPassword(key string) (string, string, error) {
	if !isInitialized() {
		return "", "", fmt.Errorf("pass_mgr not initialized. Run 'pass_mgr admin init' first")
	}
	
	adminKey, err := loadAdminKey()
	if err != nil {
		return "", "", err
	}
	
	filePath := getPasswordFilePath(key)
	fileData, err := os.ReadFile(filePath)
	if err != nil {
		return "", "", fmt.Errorf("password not found: %w", err)
	}
	
	var encrypted EncryptedData
	if err := json.Unmarshal(fileData, &encrypted); err != nil {
		return "", "", fmt.Errorf("failed to parse password file: %w", err)
	}
	
	plaintext, err := decrypt(&encrypted, adminKey)
	if err != nil {
		return "", "", fmt.Errorf("failed to decrypt: %w", err)
	}
	
	var data map[string]string
	if err := json.Unmarshal(plaintext, &data); err != nil {
		return "", "", fmt.Errorf("failed to parse decrypted data: %w", err)
	}
	
	return data["password"], data["user"], nil
}

func removePassword(key string) error {
	if !isInitialized() {
		return fmt.Errorf("pass_mgr not initialized. Run 'pass_mgr admin init' first")
	}
	
	filePath := getPasswordFilePath(key)
	if err := os.Remove(filePath); err != nil {
		return fmt.Errorf("failed to remove password: %w", err)
	}
	return nil
}

func generatePassword(length int) (string, error) {
	const charset = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789!@#$%^&*"
	password := make([]byte, length)
	for i := range password {
		randomByte := make([]byte, 1)
		if _, err := rand.Read(randomByte); err != nil {
			return "", err
		}
		password[i] = charset[int(randomByte[0])%len(charset)]
	}
	return string(password), nil
}

// CheckForAdminLeak checks if admin password appears in message/tool calling
func CheckForAdminLeak(content string) bool {
	// This is a placeholder - actual implementation should check against actual admin password
	// This function should be called by the plugin to monitor messages
	configPath := getConfigPath()
	if _, err := os.Stat(configPath); err != nil {
		return false
	}
	
	// TODO: Implement actual leak detection
	// For now, just check if content contains common patterns
	return strings.Contains(content, "admin") && strings.Contains(content, "password")
}

// ResetOnLeak resets pass_mgr to uninitialized state and logs security breach
func ResetOnLeak() error {
	configPath := getConfigPath()
	
	// Remove config (but keep key file for potential recovery)
	if err := os.Remove(configPath); err != nil {
		return err
	}
	
	// Log security breach
	logPath := filepath.Join(getHomeDir(), AdminKeyDir, "security_breach.log")
	logEntry := fmt.Sprintf("[%s] CRITICAL: Admin password leaked! pass_mgr reset to uninitialized state.\n", 
		time.Now().Format(time.RFC3339))
	
	f, err := os.OpenFile(logPath, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0600)
	if err != nil {
		return err
	}
	defer f.Close()
	
	if _, err := f.WriteString(logEntry); err != nil {
		return err
	}
	
	return nil
}