package clientapi

import (
	"encoding/json"
	"errors"
	"net/http"
	"time"

	"github.com/dariuszjaniak/warp/server/internal/httputil"
	"github.com/dariuszjaniak/warp/server/internal/keyserver"
	"github.com/dariuszjaniak/warp/server/internal/userapi"
)

// The end-to-end encryption endpoints: device keys, one-time keys,
// cross-signing, signatures and the server-side key backup. Every one of
// them is the keyserver's; the handlers decode and shape.

// registerKeyRoutes adds the keyserver endpoints, if there is one.
func (c *ClientAPI) registerKeyRoutes(r *httputil.Router) {
	if c.keys == nil {
		return
	}

	c.authed(r, "POST "+v3+"/keys/upload", c.limitDefault, c.handleKeysUpload)
	c.authed(r, "POST "+v3+"/keys/query", c.limitDefault, c.handleKeysQuery)
	c.authed(r, "POST "+v3+"/keys/claim", c.limitDefault, c.handleKeysClaim)
	c.authed(r, "POST "+v3+"/keys/device_signing/upload", c.limitDefault, c.handleCrossSigningUpload)
	c.authed(r, "POST "+v3+"/keys/signatures/upload", c.limitDefault, c.handleSignaturesUpload)
	if c.sync != nil {
		c.authed(r, "GET "+v3+"/keys/changes", c.limitDefault, c.handleKeysChanges)
	}

	c.authed(r, "POST "+v3+"/room_keys/version", c.limitDefault, c.handleCreateBackup)
	c.authed(r, "GET "+v3+"/room_keys/version", c.limitDefault, c.handleGetBackup)
	c.authed(r, "GET "+v3+"/room_keys/version/{version}", c.limitDefault, c.handleGetBackup)
	c.authed(r, "PUT "+v3+"/room_keys/version/{version}", c.limitDefault, c.handlePutBackup)
	c.authed(r, "DELETE "+v3+"/room_keys/version/{version}", c.limitDefault, c.handleDeleteBackup)

	// The keys themselves, at the three granularities the spec gives:
	// the whole backup, one room, one session.
	for _, path := range []string{
		v3 + "/room_keys/keys",
		v3 + "/room_keys/keys/{roomID}",
		v3 + "/room_keys/keys/{roomID}/{sessionID}",
	} {
		c.authed(r, "PUT "+path, c.limitDefault, c.handlePutRoomKeys)
		c.authed(r, "GET "+path, c.limitDefault, c.handleGetRoomKeys)
		c.authed(r, "DELETE "+path, c.limitDefault, c.handleDeleteRoomKeys)
	}
}

// keyError maps a keyserver error onto the Matrix error a client is sent.
func keyError(err error) error {
	var (
		matrixErr   *httputil.MatrixError
		wrongBackup *keyserver.WrongBackupVersionError
	)
	switch {
	case err == nil:
		return nil
	case errors.As(err, &matrixErr):
		return err
	case errors.As(err, &wrongBackup):
		e := httputil.NewError(http.StatusForbidden, httputil.ErrCodeWrongRoomKeysVersion, "%s", err.Error())
		e.CurrentVersion = wrongBackup.Current
		return e
	case errors.Is(err, keyserver.ErrNoSuchBackup):
		return httputil.NotFound("Unknown backup version")
	case errors.Is(err, keyserver.ErrNoSuchBackupKey):
		return httputil.NotFound("No room keys found")
	case errors.Is(err, keyserver.ErrInvalidSignature):
		return httputil.NewError(http.StatusBadRequest, httputil.ErrCodeInvalidSignature, "%s", err.Error())
	case errors.Is(err, keyserver.ErrMissingParam):
		return httputil.MissingParam("%s", err.Error())
	case errors.Is(err, keyserver.ErrInvalidParam), errors.Is(err, keyserver.ErrKeyConflict):
		return httputil.InvalidParam("%s", err.Error())
	}
	return err
}

// --- device and one-time keys ---------------------------------------------------

type keysUploadResponse struct {
	OneTimeKeyCounts map[string]int `json:"one_time_key_counts"`
}

// handleKeysUpload implements POST /_matrix/client/v3/keys/upload.
func (c *ClientAPI) handleKeysUpload(w http.ResponseWriter, r *http.Request, auth userapi.Authenticated) error {
	var req keyserver.UploadRequest
	if err := httputil.DecodeJSON(w, r, &req); err != nil {
		return err
	}
	counts, err := c.keys.UploadKeys(r.Context(), auth.UserID, auth.DeviceID, req)
	if err != nil {
		return keyError(err)
	}
	httputil.WriteJSON(w, http.StatusOK, keysUploadResponse{OneTimeKeyCounts: counts})
	return nil
}

// keysQueryRequest is the body of POST /keys/query.
type keysQueryRequest struct {
	DeviceKeys map[string][]string `json:"device_keys"`

	// Timeout is in milliseconds.
	Timeout int64 `json:"timeout"`
}

// handleKeysQuery implements POST /_matrix/client/v3/keys/query.
func (c *ClientAPI) handleKeysQuery(w http.ResponseWriter, r *http.Request, auth userapi.Authenticated) error {
	var req keysQueryRequest
	if err := httputil.DecodeJSON(w, r, &req); err != nil {
		return err
	}
	resp, err := c.keys.QueryKeys(r.Context(), auth.UserID, keyserver.QueryRequest{
		DeviceKeys: req.DeviceKeys,
		Timeout:    time.Duration(req.Timeout) * time.Millisecond,
	})
	if err != nil {
		return keyError(err)
	}
	httputil.WriteJSON(w, http.StatusOK, resp)
	return nil
}

// keysClaimRequest is the body of POST /keys/claim.
type keysClaimRequest struct {
	OneTimeKeys map[string]map[string]string `json:"one_time_keys"`
	Timeout     int64                        `json:"timeout"`
}

// handleKeysClaim implements POST /_matrix/client/v3/keys/claim.
func (c *ClientAPI) handleKeysClaim(w http.ResponseWriter, r *http.Request, _ userapi.Authenticated) error {
	var req keysClaimRequest
	if err := httputil.DecodeJSON(w, r, &req); err != nil {
		return err
	}
	resp, err := c.keys.ClaimKeys(r.Context(), keyserver.ClaimRequest{
		OneTimeKeys: req.OneTimeKeys,
		Timeout:     time.Duration(req.Timeout) * time.Millisecond,
	})
	if err != nil {
		return keyError(err)
	}
	httputil.WriteJSON(w, http.StatusOK, resp)
	return nil
}

// handleKeysChanges implements GET /_matrix/client/v3/keys/changes.
func (c *ClientAPI) handleKeysChanges(w http.ResponseWriter, r *http.Request, auth userapi.Authenticated) error {
	query := r.URL.Query()
	from := query.Get("from")
	if from == "" {
		return httputil.MissingParam("from is required")
	}
	lists, err := c.sync.KeyChanges(r.Context(), auth.UserID, from, query.Get("to"))
	if err != nil {
		return syncError(err)
	}
	httputil.WriteJSON(w, http.StatusOK, lists)
	return nil
}

// --- cross-signing ------------------------------------------------------------

// crossSigningUploadRequest is the body of POST /keys/device_signing/upload.
type crossSigningUploadRequest struct {
	keyserver.CrossSigningRequest
	Auth *uiaRequest `json:"auth"`
}

// handleCrossSigningUpload implements
// POST /_matrix/client/v3/keys/device_signing/upload.
//
// Replacing cross-signing keys asks for the password first, for the same
// reason deleting a device does: a stolen access token must not be enough
// to become the user to everyone who has verified them. The keyserver says
// when the upload is one of the two the spec exempts.
func (c *ClientAPI) handleCrossSigningUpload(w http.ResponseWriter, r *http.Request, auth userapi.Authenticated) error {
	var req crossSigningUploadRequest
	if err := httputil.DecodeJSON(w, r, &req); err != nil {
		return err
	}

	needsAuth, err := c.keys.CrossSigningNeedsAuth(r.Context(), auth.UserID, req.CrossSigningRequest)
	if err != nil {
		return keyError(err)
	}
	sessionID := ""
	if needsAuth {
		flows := []uiaFlow{{Stages: []string{stagePassword}}}
		var done bool
		sessionID, done, err = c.beginOrAdvanceUIA(w, flows, req.Auth,
			func(_ string, stage *uiaRequest) (*httputil.MatrixError, error) {
				return c.passwordStage(r, auth, stage)
			})
		if err != nil || !done {
			return err
		}
	}

	if err := c.keys.UploadCrossSigningKeys(r.Context(), auth.UserID, req.CrossSigningRequest); err != nil {
		return keyError(err)
	}
	if sessionID != "" {
		c.uia.remove(sessionID)
	}
	httputil.WriteJSON(w, http.StatusOK, struct{}{})
	return nil
}

type signaturesUploadResponse struct {
	Failures map[string]map[string]keyserver.SignatureFailure `json:"failures"`
}

// handleSignaturesUpload implements
// POST /_matrix/client/v3/keys/signatures/upload.
func (c *ClientAPI) handleSignaturesUpload(w http.ResponseWriter, r *http.Request, auth userapi.Authenticated) error {
	var req map[string]map[string]json.RawMessage
	if err := httputil.DecodeJSON(w, r, &req); err != nil {
		return err
	}
	failures, err := c.keys.UploadSignatures(r.Context(), auth.UserID, req)
	if err != nil {
		return keyError(err)
	}
	httputil.WriteJSON(w, http.StatusOK, signaturesUploadResponse{Failures: failures})
	return nil
}

// --- key backup ---------------------------------------------------------------

// backupVersionRequest is the body of POST and PUT /room_keys/version.
type backupVersionRequest struct {
	Algorithm string          `json:"algorithm"`
	AuthData  json.RawMessage `json:"auth_data"`
	Version   string          `json:"version"`
}

// handleCreateBackup implements POST /_matrix/client/v3/room_keys/version.
func (c *ClientAPI) handleCreateBackup(w http.ResponseWriter, r *http.Request, auth userapi.Authenticated) error {
	var req backupVersionRequest
	if err := httputil.DecodeJSON(w, r, &req); err != nil {
		return err
	}
	version, err := c.keys.CreateBackup(r.Context(), auth.UserID, req.Algorithm, req.AuthData)
	if err != nil {
		return keyError(err)
	}
	httputil.WriteJSON(w, http.StatusOK, map[string]string{"version": version})
	return nil
}

// handleGetBackup implements GET /_matrix/client/v3/room_keys/version and
// GET /_matrix/client/v3/room_keys/version/{version}.
func (c *ClientAPI) handleGetBackup(w http.ResponseWriter, r *http.Request, auth userapi.Authenticated) error {
	version := r.PathValue("version")
	info, err := c.keys.Backup(r.Context(), auth.UserID, version)
	if errors.Is(err, keyserver.ErrNoSuchBackup) && version == "" {
		return httputil.NotFound("No current backup version")
	}
	if err != nil {
		return keyError(err)
	}
	httputil.WriteJSON(w, http.StatusOK, info)
	return nil
}

// handlePutBackup implements PUT /_matrix/client/v3/room_keys/version/{version}.
func (c *ClientAPI) handlePutBackup(w http.ResponseWriter, r *http.Request, auth userapi.Authenticated) error {
	var req backupVersionRequest
	if err := httputil.DecodeJSON(w, r, &req); err != nil {
		return err
	}
	if err := c.keys.UpdateBackup(r.Context(), auth.UserID, r.PathValue("version"),
		req.Algorithm, req.Version, req.AuthData); err != nil {
		return keyError(err)
	}
	httputil.WriteJSON(w, http.StatusOK, struct{}{})
	return nil
}

// handleDeleteBackup implements DELETE /_matrix/client/v3/room_keys/version/{version}.
func (c *ClientAPI) handleDeleteBackup(w http.ResponseWriter, r *http.Request, auth userapi.Authenticated) error {
	if err := c.keys.DeleteBackup(r.Context(), auth.UserID, r.PathValue("version")); err != nil {
		return keyError(err)
	}
	httputil.WriteJSON(w, http.StatusOK, struct{}{})
	return nil
}

// roomKeysRequest is the body of PUT /room_keys/keys at whatever
// granularity the path names; only the matching field is read.
type roomKeysRequest struct {
	Rooms    map[string]keyserver.RoomKeyBackup `json:"rooms"`
	Sessions map[string]keyserver.KeyBackupData `json:"sessions"`
}

// handlePutRoomKeys implements PUT /_matrix/client/v3/room_keys/keys,
// .../{roomId} and .../{roomId}/{sessionId}.
func (c *ClientAPI) handlePutRoomKeys(w http.ResponseWriter, r *http.Request, auth userapi.Authenticated) error {
	roomID, sessionID := r.PathValue("roomID"), r.PathValue("sessionID")

	var rooms map[string]keyserver.RoomKeyBackup
	switch {
	case sessionID != "":
		var data keyserver.KeyBackupData
		if err := httputil.DecodeJSON(w, r, &data); err != nil {
			return err
		}
		rooms = map[string]keyserver.RoomKeyBackup{
			roomID: {Sessions: map[string]keyserver.KeyBackupData{sessionID: data}},
		}
	case roomID != "":
		var req roomKeysRequest
		if err := httputil.DecodeJSON(w, r, &req); err != nil {
			return err
		}
		rooms = map[string]keyserver.RoomKeyBackup{roomID: {Sessions: req.Sessions}}
	default:
		var req roomKeysRequest
		if err := httputil.DecodeJSON(w, r, &req); err != nil {
			return err
		}
		rooms = req.Rooms
	}

	result, err := c.keys.PutBackupKeys(r.Context(), auth.UserID, r.URL.Query().Get("version"), rooms)
	if err != nil {
		return keyError(err)
	}
	httputil.WriteJSON(w, http.StatusOK, result)
	return nil
}

// handleGetRoomKeys implements GET /_matrix/client/v3/room_keys/keys,
// .../{roomId} and .../{roomId}/{sessionId}.
func (c *ClientAPI) handleGetRoomKeys(w http.ResponseWriter, r *http.Request, auth userapi.Authenticated) error {
	roomID, sessionID := r.PathValue("roomID"), r.PathValue("sessionID")
	rooms, err := c.keys.BackupKeys(r.Context(), auth.UserID, r.URL.Query().Get("version"), roomID, sessionID)
	if err != nil {
		return keyError(err)
	}

	switch {
	case sessionID != "":
		data, ok := rooms[roomID].Sessions[sessionID]
		if !ok {
			return httputil.NotFound("No room key found for that session")
		}
		httputil.WriteJSON(w, http.StatusOK, data)
	case roomID != "":
		room, ok := rooms[roomID]
		if !ok {
			room = keyserver.RoomKeyBackup{Sessions: map[string]keyserver.KeyBackupData{}}
		}
		httputil.WriteJSON(w, http.StatusOK, room)
	default:
		httputil.WriteJSON(w, http.StatusOK, map[string]any{"rooms": rooms})
	}
	return nil
}

// handleDeleteRoomKeys implements DELETE /_matrix/client/v3/room_keys/keys,
// .../{roomId} and .../{roomId}/{sessionId}.
func (c *ClientAPI) handleDeleteRoomKeys(w http.ResponseWriter, r *http.Request, auth userapi.Authenticated) error {
	result, err := c.keys.DeleteBackupKeys(r.Context(), auth.UserID, r.URL.Query().Get("version"),
		r.PathValue("roomID"), r.PathValue("sessionID"))
	if err != nil {
		return keyError(err)
	}
	httputil.WriteJSON(w, http.StatusOK, result)
	return nil
}