package main

import (
	"net/http"
	"os"
	"os/exec"
	"syscall"

	"github.com/creack/pty"
	"github.com/gorilla/websocket"
)

var sshUpgrader = websocket.Upgrader{
	CheckOrigin: func(r *http.Request) bool { return true },
}

func sshWSHandler(w http.ResponseWriter, r *http.Request) {
	ws, err := sshUpgrader.Upgrade(w, r, nil)
	if err != nil {
		return
	}
	defer ws.Close()

	//cmd := exec.Command("bash", "-l", "-i")
	cmd := exec.Command("login")
	cmd.Env = append(os.Environ(),
		"TERM=xterm-256color",
		"LANG=C.UTF-8",
		"LC_ALL=C.UTF-8",
		"COLORTERM=truecolor",
		"TERM_PROGRAM=xterm",
	)
	cmd.SysProcAttr = &syscall.SysProcAttr{Setsid: true}

	ptmx, err := pty.Start(cmd)
	if err != nil {
		return
	}

	defer func() {
		ptmx.Close()
		if cmd.Process != nil {
			syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL)
		}
	}()

	pty.Setsize(ptmx, &pty.Winsize{Cols: 120, Rows: 40})

	go func() {
		buf := make([]byte, 4096)
		for {
			n, err := ptmx.Read(buf)
			if n > 0 {
				if ws.WriteMessage(websocket.TextMessage, buf[:n]) != nil {
					return
				}
			}
			if err != nil {
				ws.Close()
				return
			}
		}
	}()

	for {
		var msg struct {
			Type string `json:"type"`
			Data string `json:"data"`
			Cols uint16 `json:"cols"`
			Rows uint16 `json:"rows"`
		}

		if err := ws.ReadJSON(&msg); err != nil {
			return
		}

		switch msg.Type {
		case "input":
			if _, err := ptmx.Write([]byte(msg.Data)); err != nil {
				return
			}
		case "resize":
			if msg.Cols > 0 && msg.Rows > 0 {
				pty.Setsize(ptmx, &pty.Winsize{
					Cols: msg.Cols,
					Rows: msg.Rows,
				})
			}
		}
	}
}