package main

import (
	"fmt"
	"io"
	"os"
	"path/filepath"

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

type UploadArgs struct {
	Command string `json:"command"`
	LocalFilepath string `json:"local_filepath"`
	DestinationFilepath string `json:"destination_filepath"`
	PrivateKeyFile string `json:"private_key_file"`
	Username string `json:"username"`
	Host string `json:"host"`
}

type RunCommandArgs struct {
	Command string `json:"command"`
	RemoteCommand string `json:"remote_command"`
	PrivateKeyFile string `json:"private_key_file"`
	Username string `json:"username"`
	Host string `json:"host"`
}

func upload_read_args() UploadArgs {

	if len(os.Args) != 7 {
		fmt.Fprintf(os.Stderr, "Usage: %s <command> <local-file> <remote-path> <private-key> <username> <host>\n", os.Args[0])
		os.Exit(1)
	}

	return UploadArgs{
		os.Args[1],
		os.Args[2],
		os.Args[3],
		os.Args[4],
		os.Args[5],
		os.Args[6],
	}
}

func run_command_args() RunCommandArgs {

	if len(os.Args) != 6 {
		fmt.Fprintf(os.Stderr, "Usage: %s <command> <remote-command> <private-key> <username> <host>\n", os.Args[0])
		os.Exit(1)
	}

	return RunCommandArgs{
		os.Args[1],
		os.Args[2],
		os.Args[3],
		os.Args[4],
		os.Args[5],
	}

}

func get_signer(private_key_file string) ssh.Signer {

	key_bytes, err := os.ReadFile(private_key_file)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error reading private key: %v\n", err)
		os.Exit(1)
	}

	signer, err := ssh.ParsePrivateKey(key_bytes)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error parsing private key: %v\n", err)
		os.Exit(1)
	}

	return signer

}

func resolve_port(host string) string {

	if len(host) > 0 {
		hasPort := false
		for i := len(host) - 1; i >= 0; i-- {
			if host[i] == ':' {
				hasPort = true
				break
			}
			if host[i] == ']' {
				break
			}
		}
		if !hasPort {
			host = host + ":22"
		}
	}

	return host

}

func get_client(private_key_file, username, host string) *ssh.Client {

	signer := get_signer(private_key_file)

	config := &ssh.ClientConfig{
		User: username,
		Auth: []ssh.AuthMethod{
			ssh.PublicKeys(signer),
		},
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
	}

	addr := resolve_port(host)

	client, err := ssh.Dial("tcp", addr, config)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error connecting to server: %v\n", err)
		os.Exit(1)
	}

	return client
}

func upload_file() {
	args := upload_read_args()

	client := get_client(args.PrivateKeyFile, args.Username, args.Host)
	defer client.Close()

	session, err := client.NewSession()
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error creating SSH session: %v\n", err)
		os.Exit(1)
	}
	defer session.Close()

	source_file, err := os.Open(args.LocalFilepath)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error opening local file: %v\n", err)
		os.Exit(1)
	}
	defer source_file.Close()

	info, err := source_file.Stat()
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error stating local file: %v\n", err)
		os.Exit(1)
	}

	w, err := session.StdinPipe()
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error creating stdin pipe: %v\n", err)
		os.Exit(1)
	}

	if err := session.Start("scp -t " + args.DestinationFilepath); err != nil {
		fmt.Fprintf(os.Stderr, "Error starting scp: %v\n", err)
		os.Exit(1)
	}

	fmt.Fprintf(w, "C0644 %d %s\n", info.Size(), filepath.Base(args.LocalFilepath))
	if _, err := io.Copy(w, source_file); err != nil {
		fmt.Fprintf(os.Stderr, "Error sending file: %v\n", err)
		os.Exit(1)
	}
	fmt.Fprint(w, "\x00")
	w.Close()

	if err := session.Wait(); err != nil {
		fmt.Fprintf(os.Stderr, "Error during scp transfer: %v\n", err)
		os.Exit(1)
	}

	fmt.Println("Success")
}

func run_command() {

	args := run_command_args()

	client := get_client(args.PrivateKeyFile, args.Username, args.Host)
	defer client.Close()

	session, err := client.NewSession()
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error creating SSH session: %v\n", err)
		os.Exit(1)
	}
	defer session.Close()

	session.Stdout = os.Stdout
	session.Stderr = os.Stderr
	if err := session.Run(args.RemoteCommand); err != nil {
		fmt.Fprintf(os.Stderr, "Error running command: %v\n", err)
		os.Exit(1)
	}
	return

}

func main() {
	switch os.Args[1] {
	case "upload_file":
		upload_file()
	case "run_command":
		run_command()
	default: 
		fmt.Println("HERE")
	}
}