descrybe/apps/api/internal/httpapi/auth_handlers.go

package httpapi

import (
	"context"
	"errors"
	"net/http"
	"strings"

	"github.com/descrybe/descrybe-v2/apps/api/internal/auth"
	"github.com/google/uuid"
)

func (s *Server) handleRegister(w http.ResponseWriter, r *http.Request) {
	var body struct {
		Email       string `json:"email"`
		Password    string `json:"password"`
		Name        string `json:"name"`
		CompanyName string `json:"company_name"`
	}
	if err := DecodeJSON(r, &body); err != nil {
		Error(w, http.StatusBadRequest, "invalid json")
		return
	}
	res, err := s.Auth.Register(r.Context(), auth.RegisterInput{
		Email: body.Email, Password: body.Password, Name: body.Name, CompanyName: body.CompanyName,
	})
	if err != nil {
		if errors.Is(err, auth.ErrUserExists) {
			FieldError(w, http.StatusConflict, "user already exists", "user_already_exists", map[string]string{
				"email": "user already exists",
			})
			return
		}
		if errors.Is(err, auth.ErrPasswordTooShort) {
			FieldError(w, http.StatusBadRequest, "password must be at least 8 characters", "password_too_short", map[string]string{
				"password": "password must be at least 8 characters",
			})
			return
		}
		if errors.Is(err, auth.ErrRegisterFieldsRequired) {
			FieldError(w, http.StatusBadRequest, "email, password, and company name are required", "register_fields_required", map[string]string{
				"email":        "email, password, and company name are required",
				"password":     "email, password, and company name are required",
				"company_name": "email, password, and company name are required",
			})
			return
		}
		ClientOrLog(w, http.StatusBadRequest, "registration failed", err, auth.ClientError)
		return
	}
	_ = s.Billing.ProvisionFreePlan(r.Context(), res.CompanyID)
	if err := s.beginAuthenticatedSession(r.Context(), res.User.ID, res.CompanyID); err != nil {
		Error(w, http.StatusInternalServerError, "session start failed")
		return
	}
	JSON(w, http.StatusCreated, res)
}

func (s *Server) handleLogin(w http.ResponseWriter, r *http.Request) {
	var body struct {
		Email    string `json:"email"`
		Password string `json:"password"`
	}
	if err := DecodeJSON(r, &body); err != nil {
		Error(w, http.StatusBadRequest, "invalid json")
		return
	}
	lockout := s.loginAttempts()
	if locked, retryAfter := lockout.locked(body.Email); locked {
		writeRateLimited(w, loginLockoutMaxFails, retryAfter)
		return
	}
	res, err := s.Auth.Login(r.Context(), body.Email, body.Password)
	if errors.Is(err, auth.ErrMustSetPassword) {
		JSON(w, http.StatusForbidden, map[string]string{
			"error":   "password_not_set",
			"code":    "password_not_set",
			"message": PublicMessage(w, "This account still needs a password. Open your set-password invite link, or ask a company admin to re-issue one to this email."),
		})
		return
	}
	if errors.Is(err, auth.ErrInvalidCredentials) {
		lockout.recordFailure(body.Email)
		if locked, retryAfter := lockout.locked(body.Email); locked {
			writeRateLimited(w, loginLockoutMaxFails, retryAfter)
			return
		}
		FieldError(w, http.StatusUnauthorized, "invalid credentials", "invalid_credentials", map[string]string{
			"email":    "invalid credentials",
			"password": "invalid credentials",
		})
		return
	}
	if err != nil {
		Error(w, http.StatusInternalServerError, "login failed")
		return
	}
	lockout.clear(body.Email)
	if err := s.beginAuthenticatedSession(r.Context(), res.User.ID, res.CompanyID); err != nil {
		Error(w, http.StatusInternalServerError, "session start failed")
		return
	}
	JSON(w, http.StatusOK, res)
}

// sessionUserEmail returns the signed-in user's email when a session cookie is present.
func (s *Server) sessionUserEmail(ctx context.Context) (string, bool) {
	uidStr := s.Sessions.GetString(ctx, auth.SessionUserIDKey)
	if uidStr == "" {
		return "", false
	}
	uid, err := uuid.Parse(uidStr)
	if err != nil {
		return "", false
	}
	user, err := s.Auth.GetUser(ctx, uid)
	if err != nil {
		return "", false
	}
	email := strings.TrimSpace(user.Email)
	if email == "" {
		return "", false
	}
	return email, true
}

func writeEmailMismatch(w http.ResponseWriter, sessionEmail, inviteEmail string) {
	JSON(w, http.StatusConflict, map[string]string{
		"error":         "email_mismatch",
		"code":          "email_mismatch",
		"message":       PublicMessage(w, "You're signed in as a different email than this invite. Sign out to continue with the invited account, or ask an admin to re-issue the invite to your signed-in email."),
		"session_email": sessionEmail,
		"invite_email":  inviteEmail,
	})
}

func (s *Server) beginAuthenticatedSession(ctx context.Context, userID, companyID uuid.UUID) error {
	if err := s.Sessions.RenewToken(ctx); err != nil {
		return err
	}
	s.Sessions.Put(ctx, auth.SessionUserIDKey, userID.String())
	s.putSessionVersion(ctx, userID)
	// Fresh login/register clears any prior impersonation chain / staff tenant act-as.
	s.Sessions.Remove(ctx, auth.SessionImpersonatorIDKey)
	s.Sessions.Remove(ctx, auth.SessionStaffHomeCompanyKey)
	if companyID == uuid.Nil {
		s.Sessions.Put(ctx, auth.SessionCompanyIDKey, "")
		return nil
	}
	s.Sessions.Put(ctx, auth.SessionCompanyIDKey, companyID.String())
	return nil
}

// beginImpersonatedSession swaps the signed-in user while preserving the original actor.
func (s *Server) beginImpersonatedSession(ctx context.Context, targetUserID, companyID, actorID uuid.UUID) error {
	if err := s.Sessions.RenewToken(ctx); err != nil {
		return err
	}
	s.Sessions.Put(ctx, auth.SessionUserIDKey, targetUserID.String())
	s.putSessionVersion(ctx, targetUserID)
	// Keep the original impersonator across chained switches.
	if existing := strings.TrimSpace(s.Sessions.GetString(ctx, auth.SessionImpersonatorIDKey)); existing == "" {
		s.Sessions.Put(ctx, auth.SessionImpersonatorIDKey, actorID.String())
	}
	if companyID == uuid.Nil {
		s.Sessions.Put(ctx, auth.SessionCompanyIDKey, "")
		return nil
	}
	s.Sessions.Put(ctx, auth.SessionCompanyIDKey, companyID.String())
	return nil
}

// putSessionVersion stamps users.session_version into the cookie session (0 when DB unavailable).
func (s *Server) putSessionVersion(ctx context.Context, userID uuid.UUID) {
	version := 0
	if s != nil && s.Auth != nil && s.Auth.Pool != nil {
		if st, err := s.Auth.UserSessionState(ctx, userID); err == nil {
			version = st.Version
		}
	}
	s.Sessions.Put(ctx, auth.SessionVersionKey, version)
}

func (s *Server) handleLogout(w http.ResponseWriter, r *http.Request) {
	if err := s.Sessions.Destroy(r.Context()); err != nil {
		Error(w, http.StatusInternalServerError, "logout failed")
		return
	}
	JSON(w, http.StatusOK, map[string]string{"status": "ok"})
}

func (s *Server) handleInvitePreview(w http.ResponseWriter, r *http.Request) {
	var body struct {
		Token string `json:"token"`
		Mode  string `json:"mode"`
	}
	if err := DecodeJSON(r, &body); err != nil {
		Error(w, http.StatusBadRequest, "invalid json")
		return
	}
	mode := strings.TrimSpace(strings.ToLower(body.Mode))
	if mode == "" {
		mode = "invite"
	}
	var inviteEmail string
	switch mode {
	case "set-password":
		uid, err := auth.ParseSetPasswordToken(s.Config.TokenSigningSecret, body.Token)
		if err != nil {
			Error(w, http.StatusBadRequest, "invalid or expired token")
			return
		}
		user, err := s.Auth.GetUser(r.Context(), uid)
		if err != nil {
			Error(w, http.StatusBadRequest, "invalid or expired token")
			return
		}
		inviteEmail = user.Email
	default:
		mode = "invite"
		email, err := s.Auth.ResolveInviteEmail(r.Context(), body.Token)
		if errors.Is(err, auth.ErrInviteInvalid) {
			Error(w, http.StatusBadRequest, "invite invalid or expired")
			return
		}
		if err != nil {
			ClientOrLog(w, http.StatusBadRequest, "invite preview failed", err, auth.ClientError)
			return
		}
		inviteEmail = email
	}
	out := map[string]any{
		"mode":         mode,
		"invite_email": inviteEmail,
		"valid":        true,
		"mismatch":     false,
	}
	if sessionEmail, ok := s.sessionUserEmail(r.Context()); ok {
		out["session_email"] = sessionEmail
		if !auth.EmailsEqual(sessionEmail, inviteEmail) {
			out["mismatch"] = true
		}
	}
	JSON(w, http.StatusOK, out)
}

func (s *Server) handleAcceptInvite(w http.ResponseWriter, r *http.Request) {
	var body struct {
		Token    string `json:"token"`
		Password string `json:"password"`
		Name     string `json:"name"`
	}
	if err := DecodeJSON(r, &body); err != nil {
		Error(w, http.StatusBadRequest, "invalid json")
		return
	}
	if sessionEmail, ok := s.sessionUserEmail(r.Context()); ok {
		inviteEmail, err := s.Auth.ResolveInviteEmail(r.Context(), body.Token)
		if errors.Is(err, auth.ErrInviteInvalid) {
			Error(w, http.StatusBadRequest, "invite invalid or expired")
			return
		}
		if err != nil {
			ClientOrLog(w, http.StatusBadRequest, "invite accept failed", err, auth.ClientError)
			return
		}
		if !auth.EmailsEqual(sessionEmail, inviteEmail) {
			writeEmailMismatch(w, sessionEmail, inviteEmail)
			return
		}
	}
	res, err := s.Auth.AcceptInvite(r.Context(), body.Token, body.Password, body.Name)
	if errors.Is(err, auth.ErrInviteInvalid) {
		Error(w, http.StatusBadRequest, "invite invalid or expired")
		return
	}
	if errors.Is(err, auth.ErrInvalidCredentials) {
		FieldError(w, http.StatusUnauthorized, "invalid credentials", "invalid_credentials", map[string]string{
			"password": "invalid credentials",
		})
		return
	}
	if err != nil {
		ClientOrLog(w, http.StatusBadRequest, "invite accept failed", err, auth.ClientError)
		return
	}
	if err := s.beginAuthenticatedSession(r.Context(), res.User.ID, res.CompanyID); err != nil {
		Error(w, http.StatusInternalServerError, "session start failed")
		return
	}
	JSON(w, http.StatusOK, res)
}

func (s *Server) handleMe(w http.ResponseWriter, r *http.Request) {
	uid, _ := UserIDFromContext(r.Context())
	user, err := s.Auth.GetUser(r.Context(), uid)
	if err != nil {
		Error(w, http.StatusUnauthorized, "unauthorized")
		return
	}
	companies, err := s.Auth.ListUserCompanies(r.Context(), uid)
	if err != nil {
		Error(w, http.StatusInternalServerError, "failed to load companies")
		return
	}
	staffTenantSwitch := false
	access, accessErr := s.Auth.GetStaffAccess(r.Context(), uid)
	if accessErr == nil && access.FullAdmin && access.Role == auth.StaffRoleAdmin {
		staffTenantSwitch = true
		all, listErr := s.Auth.ListCompanies(r.Context(), 500)
		if listErr != nil {
			Error(w, http.StatusInternalServerError, "failed to load companies")
			return
		}
		companies = all
	}
	cidStr := s.Sessions.GetString(r.Context(), auth.SessionCompanyIDKey)
	homeStr := strings.TrimSpace(s.Sessions.GetString(r.Context(), auth.SessionStaffHomeCompanyKey))
	out := map[string]any{
		"user":              user,
		"companies":         companies,
		"active_company_id": cidStr,
	}
	if accessErr == nil && (access.FullAdmin || access.SupportDesk) {
		out["staff_access"] = access
		out["staff_capabilities"] = auth.StaffCapabilities(access.Role)
	}
	if staffTenantSwitch {
		out["staff_tenant_switch"] = true
	}
	staffOverride := false
	if cid, err := uuid.Parse(cidStr); err == nil {
		for _, c := range companies {
			if c.ID == cid {
				out["company"] = c
				break
			}
		}
		// Staff may have selected a company not in the capped list — resolve by id.
		if _, ok := out["company"]; !ok && staffTenantSwitch {
			if c, cerr := s.Auth.CompanyByID(r.Context(), cid); cerr == nil {
				out["company"] = c
			}
		}
		if credits, err := s.Billing.CreditsOverview(r.Context(), cid, s.Config.LowCreditsThreshold); err == nil {
			out["credits"] = credits
		}
		if m, err := s.Auth.EnsureMembership(r.Context(), uid, cid); err == nil {
			out["membership"] = map[string]string{"role": m.Role, "status": m.Status}
			if isOwner, oerr := s.Auth.IsCompanyOwner(r.Context(), cid, uid); oerr == nil {
				out["is_owner"] = isOwner
			}
		} else if staffTenantSwitch && errors.Is(err, auth.ErrNotCompanyMember) {
			staffOverride = true
			out["membership"] = map[string]any{"role": "admin", "status": "active", "staff_override": true}
			out["is_owner"] = false
		}
	}
	if staffTenantSwitch && homeStr != "" {
		out["staff_home_company_id"] = homeStr
		if hid, err := uuid.Parse(homeStr); err == nil {
			if c, cerr := s.Auth.CompanyByID(r.Context(), hid); cerr == nil {
				out["staff_home_company"] = c
			}
		}
	}
	if staffOverride {
		out["staff_tenant_acting"] = true
	}
	impersonating := false
	if impStr := strings.TrimSpace(s.Sessions.GetString(r.Context(), auth.SessionImpersonatorIDKey)); impStr != "" {
		if impID, err := uuid.Parse(impStr); err == nil && impID != uuid.Nil {
			impersonating = true
			out["impersonating"] = true
			if impUser, err := s.Auth.GetUser(r.Context(), impID); err == nil {
				out["impersonator"] = map[string]any{
					"id":    impUser.ID,
					"email": impUser.Email,
					"name":  impUser.Name,
				}
			} else {
				out["impersonator"] = map[string]any{"id": impID}
			}
		}
	}
	canUserSwitch := impersonating
	if !canUserSwitch {
		if accessErr == nil && access.FullAdmin && access.Role == auth.StaffRoleAdmin {
			canUserSwitch = true
		} else if !s.Config.IsProduction() {
			if accessErr == nil && access.FullAdmin {
				canUserSwitch = true
			} else if isLocalDemoEmail(user.Email) {
				canUserSwitch = true
			}
		}
	}
	if canUserSwitch {
		out["dev_user_switch"] = true
	}
	JSON(w, http.StatusOK, out)
}

func (s *Server) handleSetPassword(w http.ResponseWriter, r *http.Request) {
	uid, _ := UserIDFromContext(r.Context())
	var body struct {
		Password string `json:"password"`
	}
	if err := DecodeJSON(r, &body); err != nil {
		Error(w, http.StatusBadRequest, "invalid json")
		return
	}
	if err := s.Auth.SetPassword(r.Context(), uid, body.Password); err != nil {
		if errors.Is(err, auth.ErrPasswordAlreadySet) {
			Error(w, http.StatusBadRequest, "password already set")
			return
		}
		ClientOrLog(w, http.StatusBadRequest, "could not set password", err, auth.ClientError)
		return
	}
	s.putSessionVersion(r.Context(), uid)
	JSON(w, http.StatusOK, map[string]string{"status": "ok"})
}

func (s *Server) handleChangePassword(w http.ResponseWriter, r *http.Request) {
	uid, _ := UserIDFromContext(r.Context())
	var body struct {
		CurrentPassword string `json:"current_password"`
		Password        string `json:"password"`
	}
	if err := DecodeJSON(r, &body); err != nil {
		Error(w, http.StatusBadRequest, "invalid json")
		return
	}
	if err := s.Auth.ChangePassword(r.Context(), uid, body.CurrentPassword, body.Password); err != nil {
		if errors.Is(err, auth.ErrMustSetPassword) {
			Error(w, http.StatusBadRequest, "set password first")
			return
		}
		if errors.Is(err, auth.ErrInvalidCredentials) {
			Error(w, http.StatusBadRequest, "current password is incorrect")
			return
		}
		ClientOrLog(w, http.StatusBadRequest, "could not change password", err, auth.ClientError)
		return
	}
	s.putSessionVersion(r.Context(), uid)
	JSON(w, http.StatusOK, map[string]string{"status": "ok"})
}

func (s *Server) handleSelectCompany(w http.ResponseWriter, r *http.Request) {
	uid, _ := UserIDFromContext(r.Context())
	var body struct {
		CompanyID string `json:"company_id"`
	}
	if err := DecodeJSON(r, &body); err != nil {
		Error(w, http.StatusBadRequest, "invalid json")
		return
	}
	cid, err := uuid.Parse(body.CompanyID)
	if err != nil {
		Error(w, http.StatusBadRequest, "invalid company_id")
		return
	}
	_, memErr := s.Auth.EnsureMembership(r.Context(), uid, cid)
	staffOK := s.staffMayActAsCompany(r.Context(), uid)
	if memErr != nil {
		if !(errors.Is(memErr, auth.ErrNotCompanyMember) && staffOK) {
			Error(w, http.StatusForbidden, "forbidden")
			return
		}
		if _, cerr := s.Auth.CompanyByID(r.Context(), cid); cerr != nil {
			if errors.Is(cerr, auth.ErrCompanyNotFound) {
				Error(w, http.StatusNotFound, "company not found")
				return
			}
			Error(w, http.StatusInternalServerError, "company lookup failed")
			return
		}
		// First hop into a foreign tenant: remember home company for revert.
		if strings.TrimSpace(s.Sessions.GetString(r.Context(), auth.SessionStaffHomeCompanyKey)) == "" {
			if home := s.resolveStaffHomeCompany(r.Context(), uid); home != uuid.Nil {
				s.Sessions.Put(r.Context(), auth.SessionStaffHomeCompanyKey, home.String())
			}
		}
	} else {
		// Back on a membership company — clear act-as home.
		s.Sessions.Remove(r.Context(), auth.SessionStaffHomeCompanyKey)
	}
	s.Sessions.Put(r.Context(), auth.SessionCompanyIDKey, cid.String())
	JSON(w, http.StatusOK, map[string]string{"company_id": cid.String()})
}

// resolveStaffHomeCompany picks the company to restore on revert: current session
// company when the admin is a member, else their first membership company.
func (s *Server) resolveStaffHomeCompany(ctx context.Context, userID uuid.UUID) uuid.UUID {
	if cur := strings.TrimSpace(s.Sessions.GetString(ctx, auth.SessionCompanyIDKey)); cur != "" {
		if cid, err := uuid.Parse(cur); err == nil {
			if _, err := s.Auth.EnsureMembership(ctx, userID, cid); err == nil {
				return cid
			}
		}
	}
	companies, err := s.Auth.ListUserCompanies(ctx, userID)
	if err != nil || len(companies) == 0 {
		return uuid.Nil
	}
	return companies[0].ID
}