package main

import (
	"crypto/md5"
	_ "embed"
	"encoding/base64"
	"encoding/binary"
	"errors"
	"fmt"
	"log"
	"math/rand/v2"
	"os"
	"strings"
	"syscall"

	"github.com/pkg/sftp"
	"github.com/spf13/pflag"
	"golang.org/x/crypto/ssh"
	"golang.org/x/term"
)

//go:embed hytop
var hytop_bin []byte

var (
	flagHost = pflag.String("host", "", "host to run hytop")
	flagPort = pflag.Uint16P("port", "p", 22, "ssh port")
)

func main() {
	pflag.Parse()
	if flagHost != nil {
		// 远程执行
		run_remote()
	} else {
		// 本地执行
		run_local()
	}
}

func run_local() {
	dir := fmt.Sprintf("/tmp/hytop-%s", get_rand_str())
	err := os.Mkdir(dir, 0777)
	if err != nil {
		log.Fatal(err)
	}
	fp := fmt.Sprintf("%s/%s", dir, "hytop")
	os.WriteFile(fp, hytop_bin, 0777)
	syscall.Exec(fp, []string{"hytop"}, os.Environ())
}

func run_remote() {

	home := os.Getenv("HOME")
	pk, err := os.ReadFile(home + "/.ssh/id_rsa")
	if err != nil {
		log.Fatalf("read file ${HOME}/.ssh/id_rsa error")
	}
	signer, err := ssh.ParsePrivateKey(pk)
	if err != nil {
		log.Fatalf("unable to parse private key: %v", err)
	}
	// 1. SSH 客户端配置
	config := &ssh.ClientConfig{
		User: os.Getenv("USER"),
		Auth: []ssh.AuthMethod{
			ssh.PublicKeys(signer),
		},
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
	}

	// 2. 连接服务器
	client, err := ssh.Dial("tcp", fmt.Sprintf("%s:%d", *flagHost, *flagPort), config)
	if err != nil {
		log.Fatal("连接失败: ", err)
	}
	defer client.Close()

	binpath, err := send_file(client)
	if err != nil {
		log.Fatal("error send file: ", err)
	}

	// 3. 创建会话
	session, err := client.NewSession()
	if err != nil {
		log.Fatal("创建会话失败: ", err)
	}
	defer session.Close()

	// 4. 关键：请求 PTY (伪终端)
	// 这让远程进程认为它运行在一个真实的终端窗口中
	fd := int(os.Stdin.Fd())
	state, err := term.MakeRaw(fd) // 将本地终端设为 Raw 模式以透传按键
	if err != nil {
		log.Fatal(err)
	}
	defer term.Restore(fd, state)

	w, h, _ := term.GetSize(fd)
	if err := session.RequestPty("xterm-256color", h, w, ssh.TerminalModes{
		ssh.ECHO:          1,
		ssh.TTY_OP_ISPEED: 14400,
		ssh.TTY_OP_OSPEED: 14400,
	}); err != nil {
		log.Fatal("请求PTY失败: ", err)
	}

	// 5. 绑定输入输出
	session.Stdout = os.Stdout
	session.Stderr = os.Stderr
	session.Stdin = os.Stdin

	// 6. 启动持久命令（如 bash）
	if err := session.Start(fmt.Sprintf(`PATH="${PATH}:/opt/hyhal/bin"; %s`, binpath)); err != nil {
		log.Fatal("启动Shell失败: ", err)
	}

	// 等待命令结束
	session.Wait()

	del_file(client, binpath)
}

func get_rand_str() string {
	i := rand.Uint64()
	b := make([]byte, 8)
	binary.BigEndian.PutUint64(b, i)
	src := md5.Sum(b)
	return strings.Trim(strings.ReplaceAll(base64.StdEncoding.EncodeToString(src[:]), "/", "a"), "=")
}

func send_file(sshCli *ssh.Client) (string, error) {
	sftp_client, err := sftp.NewClient(sshCli)
	if err != nil {
		return "", err
	}
	tmp_path := get_rand_str()
	err = sftp_client.Mkdir(fmt.Sprintf("/tmp/hytop-%s", tmp_path))
	if err != nil {
		sftp_client.Close()
		return "", err
	}
	res := fmt.Sprintf("/tmp/hytop-%s/hytop", tmp_path)
	sftp_file, err := sftp_client.Create(res)
	if err != nil {
		sftp_client.Close()
		return "", err
	}
	s, err := sftp_file.Write(hytop_bin)
	if s != len(hytop_bin) || err != nil {
		sftp_file.Close()
		sftp_client.Close()
		if err != nil {
			return "", err
		} else {
			return "", errors.New("sftp send file size not equal with raw file")
		}
	}
	sftp_file.Chmod(0777)
	sftp_file.Close()
	sftp_client.Close()
	return res, nil
}

func del_file(sshCli *ssh.Client, filePath string) error {
	sftp_client, err := sftp.NewClient(sshCli)
	if err != nil {
		return err
	}
	defer sftp_client.Close()
	return sftp_client.Remove(filePath)
}