package main

import (
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"os"
	"path/filepath"
	"time"

	"github.com/spf13/cobra"
	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/agent"
	"gopkg.in/yaml.v2"
)

// Config structures
type Config struct {
	Types   map[string]DeviceType `yaml:"types"`
	Devices map[string]Device     `yaml:"devices"`
}

type DeviceType struct {
	Commands []string `yaml:"commands"`
}

type Device struct {
	User     string   `yaml:"user"`
	Type     string   `yaml:"type,omitempty"`
	Commands []string `yaml:"commands,omitempty"`
}

// RouterBackup handles SSH connections and command execution
type RouterBackup struct {
	hostname string
	username string
	password string
	keyFile  string
	port     int
	client   *ssh.Client
}

// NewRouterBackup creates a new RouterBackup instance
func NewRouterBackup(hostname, username, password, keyFile string, port int) *RouterBackup {
	return &RouterBackup{
		hostname: hostname,
		username: username,
		password: password,
		keyFile:  keyFile,
		port:     port,
	}
}

// Connect establishes SSH connection to the router
func (rb *RouterBackup) Connect() error {
	config := &ssh.ClientConfig{
		User:            rb.username,
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
		Timeout:         30 * time.Second,
	}

	// Try SSH agent first if available
	if sshAuthSock := os.Getenv("SSH_AUTH_SOCK"); sshAuthSock != "" {
		if conn, err := net.Dial("unix", sshAuthSock); err == nil {
			agentClient := agent.NewClient(conn)
			config.Auth = []ssh.AuthMethod{ssh.PublicKeysCallback(agentClient.Signers)}
		}
	}

	// If SSH agent didn't work, try key file
	if len(config.Auth) == 0 && rb.keyFile != "" {
		key, err := ioutil.ReadFile(rb.keyFile)
		if err != nil {
			return fmt.Errorf("unable to read private key: %v", err)
		}

		signer, err := ssh.ParsePrivateKey(key)
		if err != nil {
			return fmt.Errorf("unable to parse private key: %v", err)
		}

		config.Auth = []ssh.AuthMethod{ssh.PublicKeys(signer)}
	}

	// Fall back to password if available
	if len(config.Auth) == 0 && rb.password != "" {
		config.Auth = []ssh.AuthMethod{ssh.Password(rb.password)}
	}

	if len(config.Auth) == 0 {
		return fmt.Errorf("no authentication method available")
	}

	address := fmt.Sprintf("%s:%d", rb.hostname, rb.port)
	client, err := ssh.Dial("tcp", address, config)
	if err != nil {
		return fmt.Errorf("failed to connect to %s: %v", rb.hostname, err)
	}

	rb.client = client
	fmt.Printf("Successfully connected to %s\n", rb.hostname)
	return nil
}

// RunCommand executes a command on the router and returns the output
func (rb *RouterBackup) RunCommand(command string) (string, error) {
	if rb.client == nil {
		return "", fmt.Errorf("no active connection")
	}

	session, err := rb.client.NewSession()
	if err != nil {
		return "", fmt.Errorf("failed to create session: %v", err)
	}
	defer session.Close()

	output, err := session.CombinedOutput(command)
	if err != nil {
		return "", fmt.Errorf("failed to execute command '%s': %v", command, err)
	}

	return string(output), nil
}

// BackupCommands runs multiple commands and saves outputs to files
func (rb *RouterBackup) BackupCommands(commands []string, outputDir string) error {
	if err := os.MkdirAll(outputDir, 0755); err != nil {
		return fmt.Errorf("failed to create directory %s: %v", outputDir, err)
	}

	filename := rb.hostname
	filepath := filepath.Join(outputDir, filename)

	// Truncate file at start
	file, err := os.Create(filepath)
	if err != nil {
		return fmt.Errorf("failed to create file %s: %v", filepath, err)
	}
	file.Close()

	successCount := 0
	for i, command := range commands {
		fmt.Printf("Running command %d/%d: %s\n", i+1, len(commands), command)
		output, err := rb.RunCommand(command)

		if err != nil {
			fmt.Printf("Error executing '%s': %v\n", command, err)
			continue
		}

		// Append to file
		file, err := os.OpenFile(filepath, os.O_APPEND|os.O_WRONLY, 0644)
		if err != nil {
			fmt.Printf("Failed to open file for writing: %v\n", err)
			continue
		}

		fmt.Fprintf(file, "## COMMAND: %s\n", command)
		file.WriteString(output)
		file.Close()

		fmt.Printf("Output saved to %s\n", filepath)
		successCount++
	}

	fmt.Printf("Summary: %d/%d commands successful\n", successCount, len(commands))
	return nil
}

// Disconnect closes SSH connection
func (rb *RouterBackup) Disconnect() {
	if rb.client != nil {
		rb.client.Close()
		fmt.Printf("Disconnected from %s\n", rb.hostname)
	}
}

// loadConfig loads the YAML configuration file
func loadConfig(configPath string) (*Config, error) {
	data, err := ioutil.ReadFile(configPath)
	if err != nil {
		return nil, fmt.Errorf("failed to read config file %s: %v", configPath, err)
	}

	var config Config
	err = yaml.Unmarshal(data, &config)
	if err != nil {
		return nil, fmt.Errorf("failed to parse YAML: %v", err)
	}

	return &config, nil
}

// findDefaultSSHKey looks for default SSH keys
func findDefaultSSHKey() string {
	homeDir, err := os.UserHomeDir()
	if err != nil {
		return ""
	}

	defaultKeys := []string{
		filepath.Join(homeDir, ".ssh", "id_rsa"),
		filepath.Join(homeDir, ".ssh", "id_ed25519"),
		filepath.Join(homeDir, ".ssh", "id_ecdsa"),
	}

	for _, keyPath := range defaultKeys {
		if _, err := os.Stat(keyPath); err == nil {
			fmt.Printf("Using SSH key: %s\n", keyPath)
			return keyPath
		}
	}

	return ""
}

func main() {
	var configPath string
	var password string
	var keyFile string
	var port int
	var outputDir string
	var hostFilter []string

	var rootCmd = &cobra.Command{
		Use:   "ipng-router-backup",
		Short: "SSH Router Backup Tool",
		Long:  "Connects to routers via SSH and runs commands, saving output to local files.",
		Run: func(cmd *cobra.Command, args []string) {
			// Load configuration
			config, err := loadConfig(configPath)
			if err != nil {
				log.Fatalf("Failed to load config: %v", err)
			}

			// Check authentication setup
			if password == "" && keyFile == "" {
				if os.Getenv("SSH_AUTH_SOCK") != "" {
					fmt.Println("Using SSH agent for authentication")
				} else {
					keyFile = findDefaultSSHKey()
					if keyFile == "" {
						log.Fatal("No SSH key found and no password provided")
					}
				}
			}

			// Process devices
			if len(config.Devices) == 0 {
				log.Fatal("No devices found in config file")
			}

			// Filter devices if --host flags are provided
			devicesToProcess := config.Devices
			if len(hostFilter) > 0 {
				devicesToProcess = make(map[string]Device)
				for _, hostname := range hostFilter {
					if deviceConfig, exists := config.Devices[hostname]; exists {
						devicesToProcess[hostname] = deviceConfig
					} else {
						fmt.Printf("Warning: Host '%s' not found in config file\n", hostname)
					}
				}
			}

			successCount := 0
			totalCount := len(devicesToProcess)

			for hostname, deviceConfig := range devicesToProcess {
				fmt.Printf("\nProcessing device: %s (type: %s)\n", hostname, deviceConfig.Type)

				user := deviceConfig.User
				commands := deviceConfig.Commands
				deviceType := deviceConfig.Type

				// If device has a type, get commands from types section
				if deviceType != "" {
					if typeConfig, exists := config.Types[deviceType]; exists {
						commands = typeConfig.Commands
					}
				}

				if user == "" {
					fmt.Printf("No user specified for %s, skipping\n", hostname)
					continue
				}

				if len(commands) == 0 {
					fmt.Printf("No commands specified for %s, skipping\n", hostname)
					continue
				}

				// Create backup instance
				backup := NewRouterBackup(hostname, user, password, keyFile, port)

				// Connect and backup
				if err := backup.Connect(); err != nil {
					fmt.Printf("Failed to connect to %s: %v\n", hostname, err)
					continue
				}

				err = backup.BackupCommands(commands, outputDir)
				backup.Disconnect()

				if err != nil {
					fmt.Printf("Backup failed for %s: %v\n", hostname, err)
				} else {
					fmt.Printf("Backup completed for %s\n", hostname)
					successCount++
				}
			}

			fmt.Printf("\nOverall summary: %d/%d devices processed successfully\n", successCount, totalCount)
		},
	}

	rootCmd.Flags().StringVar(&configPath, "config", "", "YAML configuration file path (required)")
	rootCmd.Flags().StringVar(&password, "password", "", "SSH password")
	rootCmd.Flags().StringVar(&keyFile, "key-file", "", "SSH private key file path")
	rootCmd.Flags().IntVar(&port, "port", 22, "SSH port")
	rootCmd.Flags().StringVar(&outputDir, "output-dir", "/tmp", "Output directory for command output files")
	rootCmd.Flags().StringSliceVar(&hostFilter, "host", []string{}, "Specific host(s) to process (can be repeated, processes all if not specified)")

	rootCmd.MarkFlagRequired("config")

	if err := rootCmd.Execute(); err != nil {
		log.Fatal(err)
	}
}