package main

import (
	"encoding/base64"
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"net/http"
	"os"
	"os/exec"
	"sync"
	"time"

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

	"wxch.org/wxch/internal/protocol"
)

// Msg is the websocket envelope shared by the UI, the server and the agents.
type Msg = protocol.Msg

type term struct {
	f   *os.File
	cmd *exec.Cmd
}

func run(server, token, shell string) error {
	hdr := http.Header{"Authorization": {"Bearer " + token}}
	ws, _, err := websocket.DefaultDialer.Dial(server, hdr)
	if err != nil {
		return err
	}
	defer ws.Close()
	log.Println("connected to", server)

	var wmu sync.Mutex
	send := func(m Msg) {
		wmu.Lock()
		defer wmu.Unlock()
		ws.SetWriteDeadline(time.Now().Add(10 * time.Second))
		if err := ws.WriteJSON(m); err != nil {
			ws.Close()
		}
	}
	stopTel := make(chan struct{})
	defer close(stopTel)
	go telemetryLoop(send, stopTel)
	go moduleIdleLoop(stopTel)
	go hostProvisionAtStart()
	go limitsLoop(stopTel)
	setJobHost(server, token)
	ws.SetReadDeadline(time.Now().Add(60 * time.Second))
	ws.SetPingHandler(func(s string) error {
		ws.SetReadDeadline(time.Now().Add(60 * time.Second))
		wmu.Lock()
		defer wmu.Unlock()
		return ws.WriteControl(websocket.PongMessage, []byte(s), time.Now().Add(5*time.Second))
	})

	var tmu sync.Mutex
	terms := map[string]*term{}
	defer func() {
		tmu.Lock()
		for _, t := range terms {
			t.f.Close()
			t.cmd.Process.Kill()
		}
		tmu.Unlock()
	}()

	for {
		var m Msg
		if err := ws.ReadJSON(&m); err != nil {
			return err
		}
		ws.SetReadDeadline(time.Now().Add(60 * time.Second))
		switch m.T {
		case "open":
			cmd := exec.Command(shell, "-l")
			if m.Script != "" {
				cmd = exec.Command(shell, "-lc", m.Script)
			}
			cmd.Env = append(os.Environ(), "TERM=xterm-256color")
			f, err := pty.StartWithSize(cmd, &pty.Winsize{Cols: uint16(m.Cols), Rows: uint16(m.Rows)})
			if err != nil {
				log.Println("pty start:", err)
				send(Msg{T: "exit", Sid: m.Sid})
				continue
			}
			t := &term{f: f, cmd: cmd}
			tmu.Lock()
			terms[m.Sid] = t
			tmu.Unlock()
			sid := m.Sid
			go func() {
				buf := make([]byte, 16384)
				for {
					n, err := f.Read(buf)
					if n > 0 {
						send(Msg{T: "out", Sid: sid, D: base64.StdEncoding.EncodeToString(buf[:n])})
					}
					if err != nil {
						break
					}
				}
				cmd.Wait()
				tmu.Lock()
				delete(terms, sid)
				tmu.Unlock()
				send(Msg{T: "exit", Sid: sid})
			}()
		case "in":
			tmu.Lock()
			t := terms[m.Sid]
			tmu.Unlock()
			if t != nil {
				if b, err := base64.StdEncoding.DecodeString(m.D); err == nil {
					t.f.Write(b)
				}
			}
		case "resize":
			tmu.Lock()
			t := terms[m.Sid]
			tmu.Unlock()
			if t != nil {
				pty.Setsize(t.f, &pty.Winsize{Cols: uint16(m.Cols), Rows: uint16(m.Rows)})
			}
		case "rpc":
			go handleRPC(m, send)
		case "infer":
			go handleInfer(m, send)
		case "infer_cancel":
			cancelInfer(m.ID)
		case "job":
			go handleJob(m, send)
		case "job_cancel":
			cancelJob(m.ID)
		case "warm":
			go handleWarm(m)
		case "tun_req", "tun_ws_open", "tun_body", "tun_ws_msg", "tun_ws_close", "tun_cancel":
			handleTunMsg(m, send)
		case "update":
			go func(m Msg) {
				newVer, err := selfUpdate(server)
				if err != nil {
					log.Println("update failed:", err)
					send(Msg{T: "rpc_result", ID: m.ID, Err: "update failed: " + err.Error()})
					return
				}
				// tell the page what happened before the connection goes away with the restart
				b, _ := json.Marshal(map[string]string{"from": version, "version": newVer})
				send(Msg{T: "rpc_result", ID: m.ID, OK: true, Data: b})
				time.Sleep(700 * time.Millisecond)
				log.Println("update: restarting")
				os.Exit(0)
			}(m)
		case "close":
			tmu.Lock()
			t := terms[m.Sid]
			tmu.Unlock()
			if t != nil {
				t.f.Close()
				t.cmd.Process.Kill()
			}
		}
	}
}

func main() {
	if len(os.Args) > 1 && isCLICommand(os.Args[1]) {
		os.Exit(runCLI(os.Args[1], os.Args[2:]))
	}
	server := flag.String("server", os.Getenv("WXCH_SERVER"), "server websocket url, e.g. wss://control.example.org/ws/agent (or WXCH_SERVER)")
	token := flag.String("token", os.Getenv("WXCH_TOKEN"), "agent token (or WXCH_TOKEN)")
	shell := flag.String("shell", "/bin/bash", "shell to run")
	showVersion := flag.Bool("version", false, "print version and exit")
	doPause := flag.Bool("pause", false, "pause sharing on this machine (owner mode) and exit")
	doResume := flag.Bool("resume", false, "resume sharing on this machine and exit")
	flag.Parse()
	if *showVersion {
		fmt.Println(version)
		return
	}
	if *doPause || *doResume {
		if err := setPaused(*doPause); err != nil {
			log.Fatal(err)
		}
		if *doPause {
			fmt.Println("sharing paused; models will be unloaded within seconds. Use -resume to share again.")
		} else {
			fmt.Println("sharing resumed (it still stays off while you are actively using the machine).")
		}
		return
	}
	if err := hardenProcess(); err != nil {
		log.Println("could not make the agent non-dumpable:", err)
	}
	initGuard()
	go watchExternal()
	if *token == "" {
		log.Fatal("token required")
	}
	backoff := time.Second
	for {
		start := time.Now()
		err := run(*server, *token, *shell)
		log.Println("disconnected:", err)
		if time.Since(start) > time.Minute {
			backoff = time.Second
		}
		time.Sleep(backoff)
		if backoff < 30*time.Second {
			backoff *= 2
		}
	}
}