package sshprox

import (
	"errors"
	"fmt"
	"net/http"
	"net/url"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strconv"
	"strings"

	"github.com/google/uuid"
	"imuslab.com/zoraxy/mod/reverseproxy"
	"imuslab.com/zoraxy/mod/utils"
	"imuslab.com/zoraxy/mod/websocketproxy"
)

/*
	SSH Proxy

	This is a tool to bind gotty into Zoraxy
	so that you can do something similar to
	online ssh terminal
*/

type Manager struct {
	StartingPort  int
	ReservedPorts map[string]int
	Instances     []*Instance
}

type Instance struct {
	UUID         string
	ExecPath     string
	RemoteAddr   string
	RemotePort   int
	AssignedPort int
	conn         *reverseproxy.ReverseProxy //HTTP proxy
	tty          *exec.Cmd                  //SSH connection ported to web interface
}

func NewSSHProxyManager() *Manager {
	return &Manager{
		StartingPort:  14810,
		ReservedPorts: map[string]int{},
		Instances:     []*Instance{},
	}
}

//Get the next free port in the list
func (m *Manager) GetNextPort() int {
	nextPort := m.StartingPort
	for {
		if _, exists := m.ReservedPorts[strconv.Itoa(nextPort)]; !exists {
			if !isPortInUse(nextPort) {
				return nextPort
			}
		}

		if nextPort == 65534 {
			return -1
		}
		nextPort++
	}
}

func (m *Manager) HandleHttpByInstanceId(instanceId string, w http.ResponseWriter, r *http.Request) {
	targetInstance, err := m.GetInstanceById(instanceId)
	if err != nil {
		http.Error(w, err.Error(), http.StatusNotFound)
		return
	}

	if targetInstance.tty == nil {
		//Server side already closed
		http.Error(w, "Connection already closed", http.StatusInternalServerError)
		return
	}

	r.Header.Set("X-Forwarded-Host", r.Host)
	requestURL := r.URL.String()
	if r.Header["Upgrade"] != nil && strings.ToLower(r.Header["Upgrade"][0]) == "websocket" {
		//Handle WebSocket request. Forward the custom Upgrade header and rewrite origin
		r.Header.Set("A-Upgrade", "websocket")
		requestURL = strings.TrimPrefix(requestURL, "/")
		u, _ := url.Parse("ws://127.0.0.1:" + strconv.Itoa(targetInstance.AssignedPort) + "/" + requestURL)
		wspHandler := websocketproxy.NewProxy(u)
		wspHandler.ServeHTTP(w, r)
		return
	}

	targetInstance.conn.ProxyHTTP(w, r)
}

func (m *Manager) GetInstanceById(instanceId string) (*Instance, error) {
	for _, instance := range m.Instances {
		if instance.UUID == instanceId {
			return instance, nil
		}
	}
	return nil, fmt.Errorf("instance not found: %s", instanceId)
}
func (m *Manager) NewSSHProxy(binaryRoot string) (*Instance, error) {
	//Check if the binary exists in system/gotty/
	binary := "gotty_" + runtime.GOOS + "_" + runtime.GOARCH

	if runtime.GOOS == "windows" {
		binary = binary + ".exe"
	}
	execPath := filepath.Join(binaryRoot, binary)

	if !utils.FileExists(execPath) {
		//Binary not found
		return nil, errors.New("binary not found at " + execPath)
	}

	//Convert the binary path to realpath
	realpath, err := filepath.Abs(execPath)
	if err != nil {
		return nil, err
	}

	thisInstance := Instance{
		UUID:         uuid.New().String(),
		ExecPath:     realpath,
		AssignedPort: -1,
	}

	m.Instances = append(m.Instances, &thisInstance)

	return &thisInstance, nil
}

//Create a new Connection to target address
func (i *Instance) CreateNewConnection(listenPort int, remoteIpAddr string, remotePort int) error {
	//Create a gotty instance
	cmd := exec.Command(i.ExecPath, "-w", "-p", strconv.Itoa(listenPort), "--once", "ssh", remoteIpAddr, "-p", strconv.Itoa(remotePort))
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	go func() {
		cmd.Run()
	}()
	i.tty = cmd
	i.AssignedPort = listenPort
	i.RemoteAddr = remoteIpAddr
	i.RemotePort = remotePort

	//Create a new proxy agent for this root
	path, err := url.Parse("http://127.0.0.1:" + strconv.Itoa(listenPort))
	if err != nil {
		return err
	}

	//Create new proxy objects to the proxy
	proxy := reverseproxy.NewReverseProxy(path)

	i.conn = proxy
	return nil
}