package stress

import (
	"fmt"
	"strings"

	"auto-check/pkg/sshclient"

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

// ============================
// 国内镜像源
// ============================

// setupMirror 在安装前把对应包管理器的源切换为国内镜像（清华 TUNA）。
// useMirror=false 时直接跳过（先用系统自带源试装）。失败不致命，仅提示并继续用原源。
func setupMirror(client *ssh.Client, pm string, sudoPwd string, useMirror bool) {
	if !useMirror {
		return
	}
	const label = "清华 TUNA 镜像"
	switch pm {
	case "apt-get":
		// 备份并替换为清华镜像（Debian/Ubuntu 通用，按 lsb_release 选版本）
		script := `set -e
OS_ID=$(. /etc/os-release 2>/dev/null; echo "$ID")
OS_VER=$(. /etc/os-release 2>/dev/null; echo "$VERSION_CODENAME")
[ -z "$OS_VER" ] && OS_VER=$(. /etc/os-release 2>/dev/null; echo "$VERSION_ID")
MIRROR="https://mirrors.tuna.tsinghua.edu.cn"
if [ "$OS_ID" = "ubuntu" ]; then
  LIST="$MIRROR/ubuntu"
elif [ "$OS_ID" = "debian" ]; then
  LIST="$MIRROR/debian"
elif [ "$OS_ID" = "raspbian" ]; then
  LIST="$MIRROR/raspbian"
else
  LIST="$MIRROR/debian"
fi
BAK=/etc/apt/sources.list.d/auto-check-mirror.bak
if [ ! -f "$BAK" ]; then
  cp /etc/apt/sources.list "$BAK" 2>/dev/null || true
  cp -r /etc/apt/sources.list.d "$BAK.d" 2>/dev/null || true
fi
: > /etc/apt/sources.list
echo "deb $LIST $OS_VER main contrib non-free" > /etc/apt/sources.list
echo "deb $LIST $OS_VER-updates main contrib non-free" >> /etc/apt/sources.list
echo "deb $LIST $OS_VER-backports main contrib non-free" >> /etc/apt/sources.list
echo "deb $LIST-security $OS_VER-security main contrib non-free" >> /etc/apt/sources.list
echo OK`
		out, err := runPrivileged(client, "sh -c "+shellQuote(script), sudoPwd)
		if err != nil || !strings.Contains(out, "OK") {
			fmt.Printf("  [镜像] apt 切换国内源失败，将使用官方源: %v\n", err)
		} else {
			fmt.Printf("  [镜像] apt 已切换至%s\n", label)
		}
	case "apk":
		runPrivileged(client, "sed -i 's#https\\?://[^/]*alpinelinux.org#https://mirrors.tuna.tsinghua.edu.cn/alpine#g' /etc/apk/repositories", sudoPwd)
		fmt.Printf("  [镜像] apk 已切换至%s\n", label)
	default:
		// dnf/yum 等无通用一键切源命令，需手动配置 repo 文件，这里仅提示
		fmt.Printf("  [镜像] %s 无内置切源逻辑，建议使用系统默认源或手动配置%s repo\n", pm, label)
	}
}

// ============================
// Docker 容器化压测支持
// ============================

// DockerInfo 远程 Docker 环境信息
type DockerInfo struct {
	Available bool   // Docker 是否可用（已安装且有权限）
	Image     string // 镜像名（检测时传入，确认存在时回填）
	HasImage  bool   // 目标镜像是否已存在本地（避免每次 pull）
}

// DetectDocker 检测远程主机的 Docker 环境。
// 包括：docker 命令是否存在、是否需要 sudo、镜像是否已存在。
// sudoPwd 为非 root 用户执行时的密码（用于 sudo docker）。
func DetectDocker(client *ssh.Client, image string, sudoPwd string) DockerInfo {
	info := DockerInfo{Image: image}

	// 1. 检测 docker 命令是否存在（支持 root 和 sudo）
	if out, err := runPrivileged(client, "command -v docker", sudoPwd); err != nil || strings.TrimSpace(out) == "" {
		return info
	}
	info.Available = true

	// 2. 检测 docker 服务是否运行（docker info 成功即可）
	if _, err := runPrivileged(client, "docker info >/dev/null 2>&1", sudoPwd); err != nil {
		info.Available = false
		return info
	}

	// 3. 检测目标镜像是否已存在本地（避免每次拉取）
	if out, err := runPrivileged(client, fmt.Sprintf("docker image inspect %s >/dev/null 2>&1 && echo EXIST", image), sudoPwd); err == nil && strings.Contains(out, "EXIST") {
		info.HasImage = true
	}

	return info
}

// ============================
// 远程工具检测
// ============================

// requiredTools 压测核心依赖工具（缺则按需在线安装）。lm-sensors 依赖较多，
// 仅单文件二进制 stress-ng 与 fio 为必需；sensors 缺失仅影响温度曲线，不阻断压测。
// 这是工具检测的唯一来源，DetectTools/EnsureTools 都从这里取列表，避免两处不一致。
var requiredTools = []string{"stress-ng", "fio"}

// ToolInfo 远程工具信息
type ToolInfo struct {
	Name    string
	Path    string
	Version string
}

// ToolSet 远程可用工具集合
type ToolSet struct {
	Tools map[string]ToolInfo // name → info
}

// Has 工具是否可用
func (ts ToolSet) Has(name string) bool {
	_, ok := ts.Tools[name]
	return ok
}

// Get 获取工具信息
func (ts ToolSet) Get(name string) (ToolInfo, bool) {
	t, ok := ts.Tools[name]
	return t, ok
}

// DetectTools 检测远程所有相关工具是否可用
func DetectTools(client *ssh.Client) ToolSet {
	ts := ToolSet{Tools: make(map[string]ToolInfo)}
	for _, name := range requiredTools {
		if info, ok := detectOne(client, name); ok {
			ts.Tools[name] = info
		}
	}
	return ts
}

// EnsureTools 检测工具；缺失且 autoInstall=true 时按发行版自动安装。
// sudoPwd 为非 root 用户执行安装时提供（为空则尝试免密 sudo）。useMirror=true
// 时安装前把包管理器源切换为国内镜像（清华 TUNA）。返回安装后的工具集合，
// 以及最终仍缺失的工具名列表。
func EnsureTools(client *ssh.Client, names []string, autoInstall bool, sudoPwd string, useMirror bool) (ToolSet, []string) {
	ts := DetectTools(client)
	var missing []string
	for _, n := range names {
		if !ts.Has(n) {
			missing = append(missing, n)
		}
	}
	if len(missing) == 0 {
		return ts, nil
	}

	var failed []string
	if autoInstall {
		pm := detectPkgManager(client)
		if pm == "" {
			return ts, missing // 无法识别包管理器，跳过安装
		}
		// 安装前先把源切到国内镜像（使用官方源在网络差的环境会超时）
		setupMirror(client, pm, sudoPwd, useMirror)
		// 仅 apt 系需要先刷新元数据
		if pm == "apt-get" {
			if out, err := runPrivileged(client, "timeout 180 apt-get update 2>&1", sudoPwd); err != nil || strings.Contains(out, "Err") || strings.Contains(out, "Failed") {
				fmt.Printf("  [安装] apt-get update 异常(可能无可用源/无权限): %v\n%s\n", err, truncate(out, 400))
			}
		}
		for _, n := range missing {
			pkg := pkgName(pm, n)
			if pkg == "" {
				failed = append(failed, n)
				continue
			}
			installCmd := pkgInstallCmd(pm, pkg)
			// runPrivileged 统一处理 sudo 与引号，这里直接给裸安装命令（带 timeout 兜底）
			out, err := runPrivileged(client, fmt.Sprintf("timeout 300 %s 2>&1", installCmd), sudoPwd)
			if err != nil || !strings.Contains(out, "Setting up "+pkg) && !strings.Contains(out, "已安装") && !strings.Contains(out, "already newest") {
				fmt.Printf("  [安装] %s 安装失败: %v\n%s\n", pkg, err, truncate(out, 500))
			}
			// 安装后重新检测该工具是否到位
			if info, ok := detectOne(client, n); ok {
				ts.Tools[n] = info
			} else {
				failed = append(failed, n)
			}
		}
		if len(failed) == 0 {
			return ts, nil
		}
		return ts, failed
	}
	return ts, missing
}

// runPrivileged 以特权执行命令：root 直接执行；非 root 用 sudo，
// 若提供了 sudo 密码则通过 stdin 喂入（无需 NOPASSWD 配置）。
func runPrivileged(client *ssh.Client, cmd, sudoPwd string) (string, error) {
	// 已是 root 则直接执行
	if out, err := sshclient.RunCommand(client, "id -u"); err == nil && strings.TrimSpace(out) == "0" {
		return sshclient.RunCommand(client, cmd)
	}
	// 非 root：尝试 sudo
	if sudoPwd != "" {
		// echo 密码 | sudo -S 执行；整条命令用双引号包裹，避免命令内部单引号嵌套冲突。
		// 仅对命令内的双引号转义（安装命令均不含双引号，安全）。
		quoted := strings.ReplaceAll(singleLine(cmd), `"`, `\"`)
		wrapped := fmt.Sprintf("echo %s | sudo -S -- sh -c \"%s\"", shellQuote(sudoPwd), quoted)
		return sshclient.RunCommand(client, wrapped)
	}
	// 无密码：先尝试免密 sudo (-n)，失败则回退普通 sudo（会提示需要密码）
	quoted := strings.ReplaceAll(singleLine(cmd), `"`, `\"`)
	if out, err := sshclient.RunCommand(client, "sudo -n -- sh -c \""+quoted+"\""); err == nil {
		return out, nil
	}
	return sshclient.RunCommand(client, "sudo -- sh -c \""+quoted+"\"")
}

// shellQuote 转义单引号，避免密码破坏命令结构
func shellQuote(s string) string {
	return "'" + strings.ReplaceAll(s, "'", `'\''`) + "'"
}

// singleLine 去掉换行，保证能塞进 sh -c '...'
func singleLine(s string) string {
	return strings.ReplaceAll(strings.ReplaceAll(s, "\n", " "), "\r", "")
}

// truncate 截断过长的输出，便于日志展示
func truncate(s string, n int) string {
	// 去掉 ANSI/多余空行，只保留前 n 个字符
	s = strings.TrimSpace(s)
	if len([]rune(s)) > n {
		return string([]rune(s)[:n]) + "...(truncated)"
	}
	return s
}

// detectPkgManager 识别远程发行版的包管理器
func detectPkgManager(client *ssh.Client) string {
	for _, pm := range []string{"apt-get", "dnf", "yum", "apk", "zypper"} {
		if _, err := sshclient.RunCommand(client, fmt.Sprintf("command -v %s", pm)); err == nil {
			return pm
		}
	}
	return ""
}

// pkgName 工具名 → 对应发行版的包名
func pkgName(pm, tool string) string {
	m := map[string]string{
		"stress-ng":     "stress-ng",
		"fio":           "fio",
		"lm-sensors":    "lm-sensors",
		"stressapptest": "stressapptest",
		"iperf3":        "iperf3",
	}
	if pm == "dnf" || pm == "yum" {
		// RHEL 系 lm-sensors 包名为 lm_sensors
		if tool == "lm-sensors" {
			return "lm_sensors"
		}
	}
	if pm == "zypper" {
		if tool == "lm-sensors" {
			return "sensors"
		}
	}
	if v, ok := m[tool]; ok {
		return v
	}
	return tool
}

// pkgInstallCmd 生成安装命令
func pkgInstallCmd(pm, pkg string) string {
	switch pm {
	case "apt-get":
		return fmt.Sprintf("apt-get install -y %s", pkg)
	case "dnf":
		return fmt.Sprintf("dnf install -y %s", pkg)
	case "yum":
		return fmt.Sprintf("yum install -y %s", pkg)
	case "apk":
		return fmt.Sprintf("apk add %s", pkg)
	case "zypper":
		return fmt.Sprintf("zypper install -y %s", pkg)
	}
	return ""
}

// detectOne 检测单个工具（查系统 PATH；缺失由 EnsureTools 在线安装补全）
func detectOne(client *ssh.Client, name string) (ToolInfo, bool) {
	path := ""
	if out, err := sshclient.RunCommand(client, "which "+name+" 2>/dev/null"); err == nil {
		path = strings.TrimSpace(out)
	}
	if path == "" {
		return ToolInfo{}, false
	}

	// 取版本：lm-sensors 仅支持 -v，其余工具均支持 --version
	verFlag := "--version"
	if name == "lm-sensors" {
		verFlag = "-v"
	}
	v, _ := sshclient.RunCommand(client, path+" "+verFlag+" 2>&1 | head -1")
	version := strings.TrimSpace(v)

	return ToolInfo{Name: name, Path: path, Version: version}, true
}