package main

import (
	"bufio"
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"

	"filippo.io/age"
	"golang.org/x/crypto/argon2"
	"golang.org/x/term"
)

const (
	defaultStoreDir = ".passage-store"
	passageExt      = ".passage"
	masterPassFile  = ".master-pass" // Just a hash, not the actual password (obviously)
)

// getStoreDir returns the password store directory path
// Tries to be smart about where to put things, but falls back gracefully
func getStoreDir() string {
	if dir := os.Getenv("PASSAGE_DIR"); dir != "" {
		// Paranoid check for path traversal - you never know what users will try
		if strings.Contains(dir, "..") {
			// Nope, not today. Fall back to default.
			return filepath.Join(os.Getenv("HOME"), defaultStoreDir)
		}
		// Make sure we have an absolute path
		if !filepath.IsAbs(dir) {
			// Relative paths are fine, just anchor them to HOME
			return filepath.Join(os.Getenv("HOME"), dir)
		}
		// Double-check that even absolute paths don't escape HOME
		// (because why would you want your passwords outside HOME anyway?)
		home := os.Getenv("HOME")
		if home != "" {
			rel, err := filepath.Rel(home, dir)
			if err != nil || strings.HasPrefix(rel, "..") {
				// Nice try, but no. Default it is.
				return filepath.Join(home, defaultStoreDir)
			}
		}
		return dir
	}
	// No custom dir set, use the default location
	home, err := os.UserHomeDir()
	if err != nil {
		// Can't get home dir? Just use current directory I guess
		return defaultStoreDir
	}
	return filepath.Join(home, defaultStoreDir)
}

// getMasterPassPath returns the path to the master password verification file
func getMasterPassPath(dir string) string {
	return filepath.Join(dir, masterPassFile)
}

// getMasterPassword prompts for the master password and returns SecureBytes
// Tries to hide input if possible, but falls back gracefully for pipes/scripts
func getMasterPassword(prompt string) (*SecureBytes, error) {
	if prompt == "" {
		prompt = "Enter master password: "
	}
	
	// Try to hide the password if we're in a real terminal
	if term.IsTerminal(int(os.Stdin.Fd())) {
		fmt.Fprint(os.Stderr, prompt)
		password, err := term.ReadPassword(int(os.Stdin.Fd()))
		fmt.Fprint(os.Stderr, "\n")
		if err != nil {
			return nil, fmt.Errorf("failed to read password: %w", err)
		}
		return NewSecureBytes(password), nil
	}
	
	// Not a terminal (probably a pipe or script) - can't hide it, but at least it works
	reader := bufio.NewReader(os.Stdin)
	fmt.Fprint(os.Stderr, prompt)
	password, err := reader.ReadString('\n')
	if err != nil {
		return nil, fmt.Errorf("failed to read password: %w", err)
	}
	trimmed := strings.TrimSpace(password)
	return NewSecureBytes([]byte(trimmed)), nil
}

// checkStoreInitialized checks if the store is initialized
func checkStoreInitialized(dir string) bool {
	masterPassPath := getMasterPassPath(dir)
	_, err := os.Stat(masterPassPath)
	return err == nil
}

// getRecipientForEncryption returns a ScryptRecipient for encrypting with master password
func getRecipientForEncryption(masterPassword []byte) (age.Recipient, error) {
	return age.NewScryptRecipient(string(masterPassword))
}

// getIdentityForDecryption returns a ScryptIdentity for decrypting with master password
func getIdentityForDecryption(masterPassword []byte) (age.Identity, error) {
	return age.NewScryptIdentity(string(masterPassword))
}



// hashPassword creates an Argon2id hash of the password for verification
// Returns base64-encoded salt:hash
func hashPassword(password []byte) (string, error) {
	// Generate random salt - makes each hash unique even for same password
	salt := make([]byte, 16)
	if _, err := rand.Read(salt); err != nil {
		return "", fmt.Errorf("failed to generate salt: %w", err)
	}

	// Argon2id with reasonable defaults:
	// - 3 iterations (fast enough, slow enough for attackers)
	// - 32 MB memory (makes GPU attacks expensive)
	// - 4 threads (parallelism)
	// - 32 bytes output (standard hash size)
	hash := argon2.IDKey(password, salt, 3, 32*1024, 4, 32)

	// Encode as base64: salt:hash
	saltB64 := base64.RawStdEncoding.EncodeToString(salt)
	hashB64 := base64.RawStdEncoding.EncodeToString(hash)
	return saltB64 + ":" + hashB64, nil
}

// verifyMasterPassword checks if the provided password matches the stored hash
func verifyMasterPassword(storeDir string, password []byte) (bool, error) {
	masterPassPath := getMasterPassPath(storeDir)
	data, err := os.ReadFile(masterPassPath)
	if err != nil {
		return false, err
	}
	stored := strings.TrimSpace(string(data))

	// Parse stored salt:hash
	parts := strings.Split(stored, ":")
	if len(parts) != 2 {
		// Legacy SHA256 format - try to verify and migrate
		return verifyLegacySHA256(password, stored)
	}

	saltB64 := parts[0]
	expectedHashB64 := parts[1]

	// Decode salt and expected hash
	salt, err := base64.RawStdEncoding.DecodeString(saltB64)
	if err != nil {
		return false, fmt.Errorf("failed to decode salt: %w", err)
	}

	expectedHash, err := base64.RawStdEncoding.DecodeString(expectedHashB64)
	if err != nil {
		return false, fmt.Errorf("failed to decode hash: %w", err)
	}

	// Compute hash with same parameters
	computedHash := argon2.IDKey(password, salt, 3, 32*1024, 4, 32)

	// Constant-time comparison
	return constantTimeCompare(computedHash, expectedHash), nil
}

// verifyLegacySHA256 verifies against old SHA256 hashes (for migration)
// Kept for backwards compatibility with stores created before Argon2id migration
func verifyLegacySHA256(password []byte, storedHash string) (bool, error) {
	// Old stores used SHA256 - not ideal but we support it for migration
	hash := sha256Sum(password)
	providedHash := hex.EncodeToString(hash[:])
	
	// Constant-time comparison - can't let attackers time their way in
	return constantTimeStringCompare(providedHash, storedHash), nil
}

// sha256Sum computes SHA256 hash (for legacy verification only)
func sha256Sum(data []byte) [32]byte {
	return sha256.Sum256(data)
}

// constantTimeCompare performs constant-time comparison of two byte slices
func constantTimeCompare(a, b []byte) bool {
	if len(a) != len(b) {
		return false
	}
	var diff byte
	for i := 0; i < len(a); i++ {
		diff |= a[i] ^ b[i]
	}
	return diff == 0
}

// constantTimeStringCompare performs constant-time comparison of two strings
func constantTimeStringCompare(a, b string) bool {
	if len(a) != len(b) {
		return false
	}
	var diff byte
	for i := 0; i < len(a); i++ {
		diff |= a[i] ^ b[i]
	}
	return diff == 0
}

// storeMasterPasswordHash stores an Argon2id hash of the master password for verification
func storeMasterPasswordHash(storeDir string, password []byte) error {
	masterPassPath := getMasterPassPath(storeDir)
	hash, err := hashPassword(password)
	if err != nil {
		return fmt.Errorf("failed to hash password: %w", err)
	}
	return os.WriteFile(masterPassPath, []byte(hash), 0600)
}

// encryptFile encrypts a file using AGE with passphrase encryption
// Uses AGE's scrypt encryption - simple but effective
func encryptFile(inputPath, outputPath string, masterPassword []byte) error {
	recipient, err := getRecipientForEncryption(masterPassword)
	if err != nil {
		return fmt.Errorf("failed to create recipient: %w", err)
	}

	var input io.Reader
	if inputPath == "" || inputPath == "-" {
		input = os.Stdin
	} else {
		file, err := os.Open(inputPath)
		if err != nil {
			return err
		}
		defer file.Close()
		input = file
	}

	output, err := os.Create(outputPath)
	if err != nil {
		return err
	}
	defer output.Close()

	encrypted, err := age.Encrypt(output, recipient)
	if err != nil {
		return fmt.Errorf("failed to create encrypted writer: %w", err)
	}

	if _, err := io.Copy(encrypted, input); err != nil {
		return fmt.Errorf("failed to encrypt data: %w", err)
	}

	if err := encrypted.Close(); err != nil {
		return fmt.Errorf("failed to finalize encryption: %w", err)
	}

	return nil
}

// decryptFile decrypts an AGE-encrypted file using master password
// If outputPath is empty or "-", returns a reader for the decrypted content
// If outputPath is specified, writes decrypted content to file and returns nil reader
func decryptFile(inputPath, outputPath string, masterPassword []byte) (io.Reader, error) {
	identity, err := getIdentityForDecryption(masterPassword)
	if err != nil {
		return nil, fmt.Errorf("failed to create identity: %w", err)
	}

	var input io.Reader
	var file *os.File
	if inputPath == "" || inputPath == "-" {
		input = os.Stdin
	} else {
		file, err = os.Open(inputPath)
		if err != nil {
			return nil, err
		}
		input = file
	}

	decrypted, err := age.Decrypt(input, identity)
	if err != nil {
		if file != nil {
			file.Close()
		}
		return nil, fmt.Errorf("failed to decrypt: %w", err)
	}

	// If output path specified, write to file and close everything
	if outputPath != "" && outputPath != "-" {
		outputFile, err := os.Create(outputPath)
		if err != nil {
			if file != nil {
				file.Close()
			}
			return nil, err
		}
		defer outputFile.Close()
		if file != nil {
			defer file.Close()
		}
		if _, err := io.Copy(outputFile, decrypted); err != nil {
			return nil, fmt.Errorf("failed to write decrypted data: %w", err)
		}
		return nil, nil
	}

	// Return a reader that keeps the file open until reading is complete
	// This was a fun bug to track down - file was closing too early!
	// Now fileReader handles cleanup properly
	if file == nil {
		// Reading from stdin - no file to manage
		return decrypted, nil
	}
	return &fileReader{
		file:   file,
		reader: decrypted,
	}, nil
}

// fileReader wraps a reader and keeps the underlying file open until reading completes
type fileReader struct {
	file   *os.File
	reader io.Reader
}

func (fr *fileReader) Read(p []byte) (n int, err error) {
	if fr == nil || fr.reader == nil {
		return 0, io.EOF
	}
	n, err = fr.reader.Read(p)
	if err == io.EOF || err != nil {
		// Clean up when we're done - don't leak file handles
		if fr.file != nil {
			fr.file.Close()
			fr.file = nil
		}
		fr.reader = nil
	}
	return n, err
}