// Copyright 2025 OpenPubkey
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0

package commands

import (
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"

	"github.com/spf13/afero"
	"golang.org/x/crypto/ssh"
)

// LogoutCmd represents the logout command that removes opkssh-generated SSH keys and certificates.
type LogoutCmd struct {
	Fs         afero.Fs
	KeyPathArg string // Optional: specific key path to remove
	Verbosity  int    // Default verbosity is 0, 1 is verbose
	OutWriter  io.Writer
	ErrWriter  io.Writer
}

// NewLogoutCmd creates a new LogoutCmd instance.
func NewLogoutCmd(keyPathArg string) *LogoutCmd {
	return &LogoutCmd{
		Fs:         afero.NewOsFs(),
		KeyPathArg: keyPathArg,
	}
}

func (l *LogoutCmd) out() io.Writer {
	if l.OutWriter != nil {
		return l.OutWriter
	}
	return os.Stdout
}

func (l *LogoutCmd) errOut() io.Writer {
	if l.ErrWriter != nil {
		return l.ErrWriter
	}
	return os.Stderr
}

// Run executes the logout command, removing opkssh-generated SSH keys.
func (l *LogoutCmd) Run() error {
	if l.KeyPathArg != "" {
		return l.removeSpecificKey(l.KeyPathArg)
	}

	removedCount := 0

	// Remove keys from default SSH directory (~/.ssh/)
	n, err := l.removeDefaultSSHDirKeys()
	if err != nil {
		return err
	}
	removedCount += n

	// Remove keys from opkssh directory (~/.ssh/opkssh/)
	n, err = l.removeOpkSSHDirKeys()
	if err != nil {
		return err
	}
	removedCount += n

	if removedCount == 0 {
		fmt.Fprintln(l.out(), "No opkssh keys found to remove")
	} else {
		fmt.Fprintf(l.out(), "Successfully removed %d opkssh key pair(s)\n", removedCount)
	}
	return nil
}

// removeSpecificKey removes a specific key pair given the private key path.
func (l *LogoutCmd) removeSpecificKey(seckeyPath string) error {
	pubkeyPath := seckeyPath + "-cert.pub"

	afs := &afero.Afero{Fs: l.Fs}

	// Check if the cert file exists and was generated by openpubkey
	pubKeyBytes, err := afs.ReadFile(pubkeyPath)
	if err != nil {
		return fmt.Errorf("could not read certificate file %s: %w", pubkeyPath, err)
	}

	if !isOpenpubkeyComment(pubKeyBytes) {
		return fmt.Errorf("key at %s was not generated by opkssh", seckeyPath)
	}

	// Verify that the certificate matches the secret key before deleting
	secKeyBytes, err := afs.ReadFile(seckeyPath)
	if err != nil {
		return fmt.Errorf("could not read secret key file %s: %w", seckeyPath, err)
	}
	if err := verifyKeyPairMatch(secKeyBytes, pubKeyBytes); err != nil {
		return fmt.Errorf("key pair mismatch for %s: %w", seckeyPath, err)
	}

	if err := l.removeKeyPair(seckeyPath, pubkeyPath); err != nil {
		return err
	}

	fmt.Fprintf(l.out(), "Successfully removed opkssh key pair: %s\n", seckeyPath)
	return nil
}

// allDefaultSSHKeyFileNames returns all key file names from DefaultSSHKeyFileNames.
func allDefaultSSHKeyFileNames() []string {
	var all []string
	for _, names := range DefaultSSHKeyFileNames {
		all = append(all, names...)
	}
	return all
}

// removeDefaultSSHDirKeys finds and removes opkssh-generated keys from ~/.ssh/.
func (l *LogoutCmd) removeDefaultSSHDirKeys() (int, error) {
	homePath, err := os.UserHomeDir()
	if err != nil {
		return 0, fmt.Errorf("failed to get home directory: %w", err)
	}
	sshPath := filepath.Join(homePath, ".ssh")

	removedCount := 0
	afs := &afero.Afero{Fs: l.Fs}

	for _, keyFilename := range allDefaultSSHKeyFileNames() {
		seckeyPath := filepath.Join(sshPath, keyFilename)
		pubkeyPath := seckeyPath + "-cert.pub"

		// Check if the cert file exists
		pubKeyBytes, err := afs.ReadFile(pubkeyPath)
		if err != nil {
			if l.Verbosity >= 1 {
				fmt.Fprintf(l.errOut(), "Skipping %s: could not read certificate file\n", pubkeyPath)
			}
			continue // File doesn't exist or can't be read, skip
		}

		if !isOpenpubkeyComment(pubKeyBytes) {
			if l.Verbosity >= 1 {
				fmt.Fprintf(l.errOut(), "Skipping %s: not generated by opkssh\n", seckeyPath)
			}
			continue // Not generated by openpubkey, skip
		}

		// Verify that the certificate matches the secret key before deleting
		secKeyBytes, err := afs.ReadFile(seckeyPath)
		if err != nil {
			fmt.Fprintf(l.errOut(), "Skipping %s: could not read secret key file\n", seckeyPath)
			continue
		}
		if err := verifyKeyPairMatch(secKeyBytes, pubKeyBytes); err != nil {
			fmt.Fprintf(l.errOut(), "Skipping %s: certificate does not match secret key\n", seckeyPath)
			continue
		}

		if err := l.removeKeyPair(seckeyPath, pubkeyPath); err != nil {
			return removedCount, err
		}

		fmt.Fprintf(l.out(), "Removed %s and %s\n", seckeyPath, pubkeyPath)
		removedCount++
	}

	return removedCount, nil
}

// removeOpkSSHDirKeys finds and removes opkssh-generated keys from ~/.ssh/opkssh/.
func (l *LogoutCmd) removeOpkSSHDirKeys() (int, error) {
	homePath, err := os.UserHomeDir()
	if err != nil {
		return 0, fmt.Errorf("failed to get home directory: %w", err)
	}

	opkSSHDir := filepath.Join(homePath, ".ssh", "opkssh")
	afs := &afero.Afero{Fs: l.Fs}

	exists, err := afs.DirExists(opkSSHDir)
	if err != nil || !exists {
		return 0, nil
	}

	entries, err := afs.ReadDir(opkSSHDir)
	if err != nil {
		return 0, nil
	}

	removedCount := 0
	configPath := filepath.Join(opkSSHDir, "config")

	for _, entry := range entries {
		if entry.IsDir() {
			continue
		}
		name := entry.Name()

		// Skip config file and cert files (we handle them with their private key)
		if name == "config" || strings.HasSuffix(name, "-cert.pub") {
			continue
		}

		seckeyPath := filepath.Join(opkSSHDir, name)
		pubkeyPath := seckeyPath + "-cert.pub"

		// Check if the cert file exists
		pubKeyBytes, err := afs.ReadFile(pubkeyPath)
		if err != nil {
			if l.Verbosity >= 1 {
				fmt.Fprintf(l.errOut(), "Skipping %s: could not read certificate file\n", pubkeyPath)
			}
			continue
		}

		if !isOpenpubkeyComment(pubKeyBytes) {
			if l.Verbosity >= 1 {
				fmt.Fprintf(l.errOut(), "Skipping %s: not generated by opkssh\n", seckeyPath)
			}
			continue
		}

		// Verify that the certificate matches the secret key before deleting
		secKeyBytes, err := afs.ReadFile(seckeyPath)
		if err != nil {
			fmt.Fprintf(l.errOut(), "Skipping %s: could not read secret key file\n", seckeyPath)
			continue
		}
		if err := verifyKeyPairMatch(secKeyBytes, pubKeyBytes); err != nil {
			fmt.Fprintf(l.errOut(), "Skipping %s: certificate does not match secret key\n", seckeyPath)
			continue
		}

		if err := l.removeKeyPair(seckeyPath, pubkeyPath); err != nil {
			return removedCount, err
		}

		// Remove the IdentityFile entry from the opkssh config
		if err := l.removeFromOpkSSHConfig(configPath, seckeyPath); err != nil {
			return removedCount, fmt.Errorf("failed to update opkssh config: %w", err)
		}

		fmt.Fprintf(l.out(), "Removed %s and %s\n", seckeyPath, pubkeyPath)
		removedCount++
	}

	return removedCount, nil
}

// removeKeyPair removes both the private key and certificate files.
// The secret key is removed first so that if an error occurs removing the
// certificate, the secret key will not be left orphaned on disk.
func (l *LogoutCmd) removeKeyPair(seckeyPath string, pubkeyPath string) error {
	// Remove secret key first to avoid leaving it orphaned if cert removal fails
	if err := l.Fs.Remove(seckeyPath); err != nil && !os.IsNotExist(err) {
		return fmt.Errorf("failed to remove key %s: %w", seckeyPath, err)
	}

	// Remove certificate file
	if err := l.Fs.Remove(pubkeyPath); err != nil && !os.IsNotExist(err) {
		return fmt.Errorf("failed to remove certificate %s: %w", pubkeyPath, err)
	}

	return nil
}

// removeFromOpkSSHConfig removes the IdentityFile line for the given key from the opkssh config file.
func (l *LogoutCmd) removeFromOpkSSHConfig(configPath string, seckeyPath string) error {
	afs := &afero.Afero{Fs: l.Fs}

	content, err := afs.ReadFile(configPath)
	if err != nil {
		return nil // Config file doesn't exist, nothing to clean up
	}

	identityLine := "IdentityFile " + seckeyPath

	// Split handling both \r\n (Windows) and \n (Unix) line endings
	normalized := strings.ReplaceAll(string(content), "\r\n", "\n")
	lines := strings.Split(normalized, "\n")
	var newLines []string
	for _, line := range lines {
		if strings.TrimSpace(line) != identityLine {
			newLines = append(newLines, line)
		}
	}

	newContent := strings.Join(newLines, "\n")
	return afs.WriteFile(configPath, []byte(newContent), 0o600)
}

// verifyKeyPairMatch checks that the public key in the certificate corresponds
// to the given secret key. This prevents accidentally deleting a secret key
// when someone has overwritten the public key file with a different cert.
func verifyKeyPairMatch(secKeyBytes []byte, pubKeyBytes []byte) error {
	secKey, err := ssh.ParsePrivateKey(secKeyBytes)
	if err != nil {
		return fmt.Errorf("failed to parse secret key: %w", err)
	}

	pubKey, _, _, _, err := ssh.ParseAuthorizedKey(pubKeyBytes)
	if err != nil {
		return fmt.Errorf("failed to parse certificate: %w", err)
	}

	// If the pubKey is a certificate, extract the underlying key
	cert, ok := pubKey.(*ssh.Certificate)
	if !ok {
		return fmt.Errorf("public key file does not contain a certificate")
	}

	// Compare the public key from the certificate with the public key derived from the secret key
	secPubKey := secKey.PublicKey()
	certPubKey := cert.Key

	if secPubKey.Type() != certPubKey.Type() {
		return fmt.Errorf("key type mismatch: secret key is %s, certificate key is %s", secPubKey.Type(), certPubKey.Type())
	}

	if string(secPubKey.Marshal()) != string(certPubKey.Marshal()) {
		return fmt.Errorf("public key in certificate does not match secret key")
	}

	return nil
}

// isOpenpubkeyComment checks if an SSH public key file has an openpubkey-generated comment.
func isOpenpubkeyComment(pubKeyBytes []byte) bool {
	_, comment, _, _, err := ssh.ParseAuthorizedKey(pubKeyBytes)
	if err != nil {
		return false
	}
	return strings.HasPrefix(comment, "openpubkey")
}
