package web

import (
	"context"
	"crypto/rsa"
	"encoding/json"
	"errors"
	"net/http"
	"time"

	"github.com/labstack/echo/v5"
	"github.com/shellhub-io/shellhub/pkg/api/authorizer"
	"github.com/shellhub-io/shellhub/pkg/api/identity"
	"github.com/shellhub-io/shellhub/pkg/models"
	routesmiddleware "github.com/shellhub-io/shellhub/server/api/routes/middleware"
	"github.com/shellhub-io/shellhub/server/api/services"
	"github.com/shellhub-io/shellhub/server/ssh/pkg/magickey"
	"github.com/shellhub-io/shellhub/server/ssh/pkg/webhandoff"
	"github.com/shellhub-io/shellhub/server/ssh/web/pkg/token"
	log "github.com/sirupsen/logrus"
	"golang.org/x/net/websocket"
)

const (
	// WebsocketSSHBridgeRoute is the WebSocket upgrade that carries the terminal
	// stream. It is gated by a single-use token rather than authenticated.
	WebsocketSSHBridgeRoute = "/ws/ssh"
	// WebSessionRoute is the credential/token POST. It is authenticated, so the
	// bridge learns the logged-in ShellHub account (used in a new namespace). It
	// requires authorizer.DeviceConnect, and refuses a device outside the caller's
	// namespace, since the role it checks is the one held there. In a legacy
	// namespace it is the only permission check a browser terminal meets.
	WebSessionRoute = "/ws/ssh/session"
)

func exitLogLevel(err error) log.Level {
	switch {
	case errors.Is(err, ErrConnect),
		errors.Is(err, ErrAccessDenied),
		errors.Is(err, ErrInvalidSSHID),
		errors.Is(err, ErrAuthentication),
		errors.Is(err, ErrGetAuth),
		errors.Is(err, ErrFindDevice),
		errors.Is(err, ErrForbiddenPublicKey),
		errors.Is(err, ErrBridgeCredentialsNotFound),
		errors.Is(err, ErrWebSocketGetToken),
		errors.Is(err, ErrWebSocketGetDimensions),
		errors.Is(err, ErrWebSocketGetIP):
		return log.WarnLevel
	default:
		return log.ErrorLevel
	}
}

// Config is what the bridge needs: how to reach the SSH server it dials over the loopback, and
// how to name this instance in the tokens it signs.
type Config struct {
	// HostKeyFile is the path to that server's host key. The bridge pins its connection to
	// the key's public half, so nothing else answering on the loopback port is accepted.
	HostKeyFile string
	// Issuer names the instance signing web session tokens; it is written to their iss claim.
	Issuer string
}

// NewSSHServerBridge creates routes into a [echo.Router] to connect a webscoket to SSH using Shell session.
//
// authn is the API's authenticator, used to declare the WebSocket upgrade as
// reachable without a credential. It may be nil in tests.
//
// A web session token opens one WebSocket. The upgrade spends it before checking the terminal size
// and the client IP, so an open that fails on either cannot be retried with the same token.
//
// It fails when the SSH server's host key cannot be read, since without it the bridge has
// nothing to pin its connection to.
func NewSSHServerBridge(router *echo.Echo, authn *routesmiddleware.Authenticator, service services.Service, handoff *webhandoff.Store, config *Config) error {
	hostKey, err := readHostKey(config.HostKeyFile)
	if err != nil {
		return err
	}

	manager := newManager(30 * time.Second)

	if authn != nil {
		authn.AllowAnonymous(http.MethodGet, WebsocketSSHBridgeRoute)
	}

	router.Add(http.MethodPost, WebSessionRoute, echo.WrapHandler(
		http.HandlerFunc(func(res http.ResponseWriter, req *http.Request) {
			type Success struct {
				Token string `json:"token"`
			}

			type Fail struct {
				Error string `json:"error"`
			}

			decoder := json.NewDecoder(req.Body)
			encoder := json.NewEncoder(res)

			response := func(res http.ResponseWriter, status int, data any) {
				res.Header().Set("Content-Type", "application/json")
				res.WriteHeader(status)

				encoder.Encode(data) //nolint:errcheck
			}

			var request Credentials
			if err := decoder.Decode(&request); err != nil {
				response(res, http.StatusBadRequest, Fail{Error: err.Error()})

				return
			}

			caller, _ := identity.FromContext(req.Context())

			request.UserID = caller.ID

			sc, err := services.BoundTo(caller.TenantID)
			if err != nil {
				response(res, http.StatusForbidden, Fail{Error: err.Error()})

				return
			}

			if _, err := service.GetDevice(req.Context(), sc, models.UID(request.Device)); err != nil {
				response(res, http.StatusForbidden, Fail{Error: err.Error()})

				return
			}

			key := magickey.GetReference()

			token, err := token.NewToken(config.Issuer)
			if err != nil {
				response(res, http.StatusBadRequest, Fail{Error: err.Error()})

				return
			}

			if err := request.encryptPassword(key); err != nil {
				status := http.StatusInternalServerError
				if errors.Is(err, rsa.ErrMessageTooLong) {
					status = http.StatusBadRequest
				}

				response(res, status, Fail{Error: err.Error()})

				return
			}

			manager.save(token.ID, &request)

			response(res, http.StatusOK, Success{Token: token.ID})
		})),
		routesmiddleware.RequiresPermission(authorizer.DeviceConnect),
	)

	router.Add(http.MethodGet, WebsocketSSHBridgeRoute, echo.WrapHandler(websocket.Handler(func(wsconn *websocket.Conn) {
		conn := NewConn(wsconn)
		defer conn.Close() //nolint:errcheck // the handler is returning and has no one left to report a close failure to

		exit := func(err error) {
			log.WithError(err).Log(exitLogLevel(err), "web terminal error")

			buffer, marshalErr := json.Marshal(Message{
				Kind: messageKindError,
				Data: err.Error(),
			})
			if marshalErr != nil {
				log.WithError(marshalErr).Error("failed to marshal error message")

				return
			}

			conn.Write(buffer) //nolint:errcheck // a failed write means the browser is gone, so there is no one left to tell
		}

		token, err := getToken(wsconn.Request())
		if err != nil {
			exit(ErrWebSocketGetToken)

			return
		}

		creds, ok := manager.get(token)
		if !ok {
			exit(ErrBridgeCredentialsNotFound)

			return
		}

		cols, rows, err := getDimensions(wsconn.Request())
		if err != nil {
			exit(ErrWebSocketGetDimensions)

			return
		}

		ip, err := getIP(wsconn.Request())
		if err != nil {
			exit(ErrWebSocketGetIP)

			return
		}

		ctx, cancel := whileBrowserAnswers(wsconn.Request().Context(), conn)
		defer cancel()

		if err := creds.decryptPassword(magickey.GetReference()); err != nil {
			exit(errors.Join(ErrBridgeDecryptPassword, err))

			return
		}

		if err := newSession(
			ctx,
			service,
			handoff,
			conn,
			creds,
			Dimensions{cols, rows},
			Info{IP: ip},
			hostKey,
		); err != nil {
			exit(err)

			return
		}
	})))

	return nil
}

func whileBrowserAnswers(parent context.Context, conn *Conn) (context.Context, context.CancelFunc) {
	ctx, cancel := context.WithCancel(parent)

	go func() {
		conn.KeepAlive()
		cancel()
	}()

	return ctx, cancel
}