package ssh

import (
	"fmt"
	"io"

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

// Client wraps the SSH connection
type Client struct {
	config  *ssh.ClientConfig
	address string
	conn    *ssh.Client
}

// Config contains SSH connection parameters
type Config struct {
	Host     string
	Port     int
	User     string
	Password string // TODO: Add key-based auth support
}

// NewClient creates a new SSH client
func NewClient(cfg Config) *Client {
	config := &ssh.ClientConfig{
		User: cfg.User,
		Auth: []ssh.AuthMethod{
			ssh.Password(cfg.Password),
		},
		HostKeyCallback: ssh.InsecureIgnoreHostKey(), // TODO: Implement proper host key verification
	}

	return &Client{
		config:  config,
		address: fmt.Sprintf("%s:%d", cfg.Host, cfg.Port),
	}
}

// Connect establishes the SSH connection
func (c *Client) Connect() error {
	conn, err := ssh.Dial("tcp", c.address, c.config)
	if err != nil {
		return fmt.Errorf("failed to connect: %w", err)
	}
	c.conn = conn
	return nil
}

// Close closes the SSH connection
func (c *Client) Close() error {
	if c.conn != nil {
		return c.conn.Close()
	}
	return nil
}

// RunCommand runs a command and returns stdout
func (c *Client) RunCommand(cmd string) (string, error) {
	session, err := c.conn.NewSession()
	if err != nil {
		return "", fmt.Errorf("failed to create session: %w", err)
	}
	defer session.Close()

	output, err := session.CombinedOutput(cmd)
	if err != nil {
		return string(output), fmt.Errorf("command failed: %w", err)
	}
	return string(output), nil
}

// StreamCommand runs a command and streams stdout to a writer
func (c *Client) StreamCommand(cmd string, stdout io.Writer, stderr io.Writer) error {
	session, err := c.conn.NewSession()
	if err != nil {
		return fmt.Errorf("failed to create session: %w", err)
	}
	defer session.Close()

	session.Stdout = stdout
	session.Stderr = stderr

	if err := session.Run(cmd); err != nil {
		return fmt.Errorf("command failed: %w", err)
	}
	return nil
}

// StartCommand starts a command and returns pipes for streaming
func (c *Client) StartCommand(cmd string) (io.Reader, io.Reader, func() error, error) {
	session, err := c.conn.NewSession()
	if err != nil {
		return nil, nil, nil, fmt.Errorf("failed to create session: %w", err)
	}

	stdout, err := session.StdoutPipe()
	if err != nil {
		session.Close()
		return nil, nil, nil, fmt.Errorf("failed to get stdout: %w", err)
	}

	stderr, err := session.StderrPipe()
	if err != nil {
		session.Close()
		return nil, nil, nil, fmt.Errorf("failed to get stderr: %w", err)
	}

	if err := session.Start(cmd); err != nil {
		session.Close()
		return nil, nil, nil, fmt.Errorf("failed to start command: %w", err)
	}

	cleanup := func() error {
		session.Signal(ssh.SIGTERM)
		return session.Close()
	}

	return stdout, stderr, cleanup, nil
}