package sshmanager

import (
	"bytes"
	"context"
	"fmt"
	"strings"
	"time"
)

type ShutdownResult struct {
	Success bool
	Output  string
	Error   string
}

func IsShutdownCommand(cmd string) bool {
	c := strings.TrimSpace(strings.ToLower(cmd))
	for _, p := range []string{
		"shutdown", "poweroff", "halt", "reboot",
		"sudo shutdown", "sudo poweroff", "sudo halt", "sudo reboot",
		"systemctl poweroff", "systemctl halt", "systemctl reboot",
	} {
		if strings.HasPrefix(c, p) {
			return true
		}
	}
	return false
}

func RunRemoteCommand(ctx context.Context, host string, port int, user, privKeyPath, knownHostsPath string, strictHostKeyChecking bool, command string) (*ShutdownResult, error) {
	conn, _, _, err := dialSSH(ctx, host, port, user, privKeyPath, knownHostsPath, strictHostKeyChecking)
	if err != nil {
		return &ShutdownResult{Success: false, Error: err.Error()}, nil
	}
	defer conn.Close()

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

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

	effectiveCmd := command
	if IsShutdownCommand(command) {
		effectiveCmd = fmt.Sprintf(
			"nohup %s >/dev/null 2>&1 </dev/null & sleep 1; echo shutdown_initiated",
			command,
		)
	}

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

	done := make(chan error, 1)
	go func() {
		done <- session.Run(effectiveCmd)
	}()

	select {
	case err := <-done:
		if err != nil {
			return &ShutdownResult{
				Success: false,
				Error:   fmt.Sprintf("exec: %v, stderr: %s", err, stderr.String()),
			}, nil
		}
		return &ShutdownResult{
			Success: true,
			Output:  stdout.String(),
		}, nil
	case <-ctx.Done():
		return &ShutdownResult{
			Success: false,
			Error:   "command timed out after 15 seconds",
		}, nil
	}
}