package sshmanager

import (
	"bytes"
	"context"
	"fmt"
	"net"
	"os"
	"path/filepath"
	"strings"
	"time"

	"golang.org/x/crypto/ssh"
)

type ConnResult struct {
	Success    bool
	Output     string
	Error      string
	Fingerprint string
}

func TestSSHConnection(ctx context.Context, host string, port int, user, privKeyPath, knownHostsPath string, strictHostKeyChecking bool) (*ConnResult, error) {
	addr := fmt.Sprintf("%s:%d", host, port)

	auths := []ssh.AuthMethod{}
	if privKeyPath != "" {
		key, err := os.ReadFile(privKeyPath)
		if err != nil {
			return nil, fmt.Errorf("reading private key: %w", err)
		}
		signer, err := ssh.ParsePrivateKey(key)
		if err != nil {
			return nil, fmt.Errorf("parsing private key: %w", err)
		}
		auths = append(auths, ssh.PublicKeys(signer))
	}

	hostKeyPolicy := ssh.InsecureIgnoreHostKey()
	if strictHostKeyChecking && knownHostsPath != "" {
		hostKeyCallback, err := getHostKeyCallback(knownHostsPath, host, port)
		if err != nil {
			return &ConnResult{Success: false, Error: fmt.Sprintf("known_hosts: %v", err)}, nil
		}
		hostKeyPolicy = hostKeyCallback
	}

	cfg := &ssh.ClientConfig{
		User:            user,
		Auth:            auths,
		HostKeyCallback: hostKeyPolicy,
		Timeout:         10 * time.Second,
	}

	ctx, cancel := context.WithTimeout(ctx, 30*time.Second)
	defer cancel()

	conn, err := ssh.Dial("tcp", addr, cfg)
	if err != nil {
		if strings.Contains(err.Error(), "known_hosts") || strings.Contains(err.Error(), "host key") {
			return &ConnResult{
				Success: false,
				Error:   fmt.Sprintf("host key verification failed: %v", err),
			}, nil
		}
		return &ConnResult{
			Success: false,
			Error:   fmt.Sprintf("connection failed: %v", err),
		}, nil
	}
	defer conn.Close()

	session, err := conn.NewSession()
	if err != nil {
		return &ConnResult{Success: false, Error: fmt.Sprintf("session: %v", err)}, nil
	}
	defer session.Close()

	var stdout, stderr bytes.Buffer
	session.Stdout = &stdout
	session.Stderr = &stderr

	if err := session.Run("echo ok && uname -a"); err != nil {
		return &ConnResult{
			Success: false,
			Error:   fmt.Sprintf("exec: %v, stderr: %s", err, stderr.String()),
		}, nil
	}

	return &ConnResult{
		Success: true,
		Output:  stdout.String(),
	}, nil
}

func getHostKeyCallback(knownHostsPath, host string, port int) (ssh.HostKeyCallback, error) {
	return ssh.HostKeyCallback(func(hostname string, remote net.Addr, key ssh.PublicKey) error {
		kh, err := GetKnownHost(filepath.Dir(knownHostsPath), host, port)
		if err != nil {
			return fmt.Errorf("checking known_hosts: %w", err)
		}
		if kh == nil {
			return fmt.Errorf("host key not found in known_hosts: %s", hostname)
		}
		return nil
	}), nil
}