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"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promauto"
)

var (
	proxyError = promauto.NewCounterVec(prometheus.CounterOpts{
		Namespace: "ssh_proxy",
		Name:      "proxy_error",
		Help:      "Total number of errors during proxy setup operation",
	}, []string{"type"})
	proxySuccess = promauto.NewCounter(prometheus.CounterOpts{
		Namespace: "ssh_proxy",
		Name:      "proxy_success",
		Help:      "Total number of successful proxy sessions",
	})
)

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") {
		proxyError.With(prometheus.Labels{"type": "token_missing_scope"}).Inc()
		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
	}

	proxyError.With(prometheus.Labels{"type": "not_authorized"}).Inc()
	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)
	}

	// TODO: Set Access-Control-Max-Age header if policy requires time-limited
	// sessions then terminate the session once that timer expires

	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}

	proxySuccess.Inc()

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

	<-errc
	return nil
}