noteplace-server/internal/api/auth.go

package api

import (
	"net/http"
	"time"

	"git.pengzhan.dev/noteplace-server/internal/auth"
	"git.pengzhan.dev/noteplace-server/internal/models"
	"git.pengzhan.dev/noteplace-server/internal/store"
	"github.com/gin-gonic/gin"
	"github.com/google/uuid"
	"go.uber.org/zap"
	"golang.org/x/oauth2"
)

// AuthHandler handles the authentication API endpoints.
type AuthHandler struct {
	Store         *store.Store
	Authenticator *auth.Authenticator
	FrontendURL   string
}

// HandleGoogleLogin redirects the user to the Google consent page for the web flow.
func (h *AuthHandler) HandleGoogleLogin(c *gin.Context) {
	url := h.Authenticator.GetWebAuthURL()
	zap.S().Infof("Redirecting user to Google for web authentication: %s", url)
	c.Redirect(http.StatusTemporaryRedirect, url)
}

// HandleGoogleCallback handles the redirect from Google for the web flow.
func (h *AuthHandler) HandleGoogleCallback(c *gin.Context) {
	code := c.Query("code")
	if code == "" {
		zap.S().Warn("Web auth callback received without a code.")
		c.JSON(http.StatusBadRequest, gin.H{"error": "code is missing"})
		return
	}

	zap.S().Info("Received web auth code, exchanging for token.")
	token, err := h.Authenticator.ExchangeWebCodeForToken(code)
	if err != nil {
		zap.S().Errorf("Failed to exchange web code for token: %v", err)
		c.JSON(http.StatusUnauthorized, gin.H{"error": "Failed to validate code"})
		return
	}

	h.handleToken(c, token)
}

// HandleCliLogin provides the user with the URL to start the auth process for the CLI flow.
func (h *AuthHandler) HandleCliLogin(c *gin.Context) {
	authURL := h.Authenticator.GetCliAuthURL()
	zap.S().Info("Providing CLI user with auth URL.")

	response := gin.H{
		"message": `Please open the following URL in your browser to authorize the application.\nAfter authorizing, Google will provide a code. Paste it here to complete the login.\nExample: curl -X POST -d '{\"code\":\"YOUR_CODE_HERE\"}' http://localhost:3001/api/auth/google/cli/callback`,
		"url":     authURL,
	}
	c.JSON(http.StatusOK, response)
}

// HandleCliCallback exchanges the code for a token and creates a session for the CLI flow.
func (h *AuthHandler) HandleCliCallback(c *gin.Context) {
	var reqBody struct {
		Code string `json:"code"`
	}
	if err := c.ShouldBindJSON(&reqBody); err != nil {
		zap.S().Warnf("Invalid request body for CLI callback: %v", err)
		c.JSON(http.StatusBadRequest, gin.H{"error": "Invalid request body"})
		return
	}

	zap.S().Info("Received CLI auth code, exchanging for token.")
	token, err := h.Authenticator.ExchangeCliCodeForToken(reqBody.Code)
	if err != nil {
		zap.S().Errorf("Failed to exchange cli code for token: %v", err)
		c.JSON(http.StatusUnauthorized, gin.H{"error": "Failed to validate code"})
		return
	}

	h.handleToken(c, token)
}

// handleToken is a helper function to handle user/session creation after getting a token.
func (h *AuthHandler) handleToken(c *gin.Context, token *oauth2.Token) {
	zap.S().Info("Token acquired, fetching user info.")
	userInfo, err := h.Authenticator.GetUserInfo(token)
	if err != nil {
		zap.S().Errorf("Failed to get user info: %v", err)
		c.JSON(http.StatusInternalServerError, gin.H{"error": "Failed to get user info from Google"})
		return
	}

	user, found := h.Store.GetUserByGoogleID(userInfo.GoogleID)
	if !found {
		zap.S().Infof("User with Google ID %s not found, creating new user.", userInfo.GoogleID)
		userInfo.ID = "user-" + uuid.New().String()
		userInfo.CreatedAt = time.Now().UTC().Format(time.RFC3339)
		if err := h.Store.CreateUser(userInfo); err != nil {
			zap.S().Errorf("Failed to create user: %v", err)
			c.JSON(http.StatusInternalServerError, gin.H{"error": "Failed to create user"})
			return
		}
		user = userInfo
	} else {
		zap.S().Infof("Found existing user %s for Google ID %s.", user.ID, user.GoogleID)
	}

	sessionID := "session-" + uuid.New().String()
	now := time.Now().UTC().Format(time.RFC3339)
	newSession := models.Session{
		ID:        sessionID,
		UserID:    user.ID,
		CreatedAt: now,
		UpdatedAt: now,
	}
	if err := h.Store.CreateSession(newSession); err != nil {
		zap.S().Errorf("Failed to create session: %v", err)
		c.JSON(http.StatusInternalServerError, gin.H{"error": "Failed to create session"})
		return
	}
	zap.S().Infof("Created new session %s for user %s.", newSession.ID, user.ID)

	// Redirect back to the frontend for the web flow
	if c.Request.Method == "GET" {
		// Simple way to distinguish web callback from CLI POST
		redirectURL := h.FrontendURL + "/auth/callback?token=" + newSession.ID
		zap.S().Infof("Redirecting to frontend: %s", redirectURL)
		c.Redirect(http.StatusTemporaryRedirect, redirectURL)
	} else {
		// Respond with JSON for the CLI flow
		zap.S().Info("Returning session token to CLI client.")
		c.JSON(http.StatusOK, gin.H{"message": "Login successful!", "session_token": newSession.ID})
	}
}

// HandleGetUser returns the current user based on the session token.
func (h *AuthHandler) HandleGetUser(c *gin.Context) {
	user := c.MustGet("user").(models.User)
	c.JSON(http.StatusOK, user)
}