package controllers

import (
	"fmt"
	"net"
	"net/http"
	"strconv"

	"code.crute.us/mcrute/ssh-proxy/app/middleware"
	"code.crute.us/mcrute/ssh-proxy/app/models"
	"code.crute.us/mcrute/ssh-proxy/proxy"

	"github.com/gorilla/websocket"
	"github.com/labstack/echo/v4"
)

type ProxyHandler struct {
	Logger   echo.Logger
	Upgrader websocket.Upgrader
	Users    models.UserStore
}

func getConnectAddr(c echo.Context) string {
	p, err := strconv.Atoi(c.Param("port"))
	if err != nil {
		p = 22
	}
	return fmt.Sprintf("%s:%d", c.Param("host"), p)
}

func (h *ProxyHandler) authorizeRequest(c echo.Context) error {
	session := middleware.GetAuthorizedSession(c)

	user, err := h.Users.Get(c.Request().Context(), session.UserId)
	if err != nil {
		return err
	}

	if !session.HasScope("ssh:proxy") {
		return fmt.Errorf("Authorized session does not have scope ssh:proxy")
	}

	host := c.Param("host")
	if user.AuthorizedForHost(host) {
		h.Logger.Infof("Allowing user %s to proxy to host %s", session.UserId, host)
		return nil
	}

	return fmt.Errorf("User %s not authorized for host %s", session.UserId, host)
}

func (h *ProxyHandler) Handle(c echo.Context) error {
	if err := h.authorizeRequest(c); err != nil {
		h.Logger.Error(err)
		return c.NoContent(http.StatusUnauthorized)
	}

	wsconn, err := h.Upgrader.Upgrade(c.Response(), c.Request(), nil)
	if err != nil {
		return err
	}
	defer wsconn.Close()

	proxyconn, err := net.Dial("tcp", getConnectAddr(c))
	if err != nil {
		return err
	}
	defer proxyconn.Close()

	errc := make(chan error)
	ws := &proxy.WebsocketReadWriter{W: wsconn}

	go proxy.CopyWithErrors(proxyconn, ws, errc)
	go proxy.CopyWithErrors(ws, proxyconn, errc)

	<-errc
	return nil
}