Poll-system/app/internal/handlers/login.go

package handlers

import (
	"database/sql"
	"log"
	"net/http"
	"strconv"
	"time"

	"github.com/golang-jwt/jwt/v5"
	"golang.org/x/crypto/bcrypt"

	"github.com/patel-mann/poll-system/app/internal/models"
	"github.com/patel-mann/poll-system/app/internal/utils"
)

var jwtKey = []byte("your-secret-key") //TODO: Move to env/config

// Helper function to get redirect URL based on role
func getDefaultRedirectURL(role int) string {
	switch role {
	case 1: // Admin
		return "/dashboard"
	case 2: // Volunteer
		return "/volunteer/dashboard"
	case 3: // Volunteer
		return "/volunteer/dashboard"
	default:
		return "/" // Fallback to login page
	}
}

// Helper function to render error pages with consistent data
func renderLoginError(w http.ResponseWriter, errorMsg string) {
	utils.Render(w, "login.html", map[string]interface{}{
		"Error":           errorMsg,
		"Title":           "Login",
		"IsAuthenticated": false,
	})
}

func renderRegisterError(w http.ResponseWriter, errorMsg string) {
	utils.Render(w, "register.html", map[string]interface{}{
		"Error":           errorMsg,
		"Title":           "Register",
		"IsAuthenticated": false,
	})
}

// Helper function to create and sign JWT token
func createJWTToken(userID, role int) (string, time.Time, error) {
	expirationTime := time.Now().Add(12 * time.Hour)
	claims := &models.Claims{
		UserID: userID,
		Role:   role,
		RegisteredClaims: jwt.RegisteredClaims{
			ExpiresAt: jwt.NewNumericDate(expirationTime),
			IssuedAt:  jwt.NewNumericDate(time.Now()),
		},
	}

	token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
	tokenString, err := token.SignedString(jwtKey)
	return tokenString, expirationTime, err
}

// Helper function to set session cookie
func setSessionCookie(w http.ResponseWriter, tokenString string, expirationTime time.Time) {
	http.SetCookie(w, &http.Cookie{
		Name:     "session",
		Value:    tokenString,
		Path:     "/",
		HttpOnly: true,
		Secure:   false, // Set to true in production with HTTPS
		SameSite: http.SameSiteStrictMode,
		Expires:  expirationTime,
	})
}

// Helper function to clear session cookie
func clearSessionCookie(w http.ResponseWriter) {
	http.SetCookie(w, &http.Cookie{
		Name:     "session",
		Value:    "",
		Path:     "/",
		MaxAge:   -1,
		HttpOnly: true,
		Secure:   false, // Set to true in production with HTTPS
		SameSite: http.SameSiteStrictMode,
	})
}

func LoginHandler(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		http.Redirect(w, r, "/", http.StatusSeeOther)
		return
	}

	email := r.FormValue("email")
	password := r.FormValue("password")

	// Input validation
	if email == "" || password == "" {
		http.Redirect(w, r, "/?error=EmailAndPasswordRequired", http.StatusSeeOther)
		return
	}

	// Get user from database
	var storedHash string
	var userID int
	var role int

	err := models.DB.QueryRow(`
		SELECT user_id, password, role_id
		FROM "users"
		WHERE email = $1
	`, email).Scan(&userID, &storedHash, &role)

	if err != nil {
		log.Printf("Login failed for email %s: %v", email, err)
		http.Redirect(w, r, "/?error=InvalidCredentials", http.StatusSeeOther)
		return
	}

	// Verify password
	err = bcrypt.CompareHashAndPassword([]byte(storedHash), []byte(password))
	if err != nil {
		log.Printf("Password verification failed for user ID %d", userID)
		http.Redirect(w, r, "/?error=InvalidCredentials", http.StatusSeeOther)
		return
	}

	// Create JWT token
	tokenString, expirationTime, err := createJWTToken(userID, role)
	if err != nil {
		log.Printf("JWT token creation failed for user ID %d: %v", userID, err)
		http.Redirect(w, r, "/?error=InternalError", http.StatusSeeOther)
		return
	}

	// Set session cookie
	setSessionCookie(w, tokenString, expirationTime)

	// Redirect based on user role
	redirectURL := getDefaultRedirectURL(role)
	log.Printf("User %d (role %d) logged in successfully, redirecting to %s", userID, role, redirectURL)
	http.Redirect(w, r, redirectURL, http.StatusSeeOther)
}


func RegisterHandler(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		utils.Render(w, "register.html", map[string]interface{}{
			"Title":           "Register",
			"IsAuthenticated": false,
		})
		return
	}

	firstName := r.FormValue("first_name")
	lastName := r.FormValue("last_name")
	email := r.FormValue("email")
	phone := r.FormValue("phone")
	role := r.FormValue("role")
	password := r.FormValue("password")
	adminCode := r.FormValue("admin_code") // for volunteers

	// Input validation
	if firstName == "" || lastName == "" || email == "" || password == "" || role == "" {
		renderRegisterError(w, "All fields are required")
		return
	}

	// Hash password
	hashedPassword, err := bcrypt.GenerateFromPassword([]byte(password), bcrypt.DefaultCost)
	if err != nil {
		log.Printf("Password hashing failed: %v", err)
		http.Error(w, "Internal server error", http.StatusInternalServerError)
		return
	}

	// Convert role to int
	roleID, err := strconv.Atoi(role)
	if err != nil {
		renderRegisterError(w, "Invalid role")
		return
	}

	var adminID int
	if roleID == 3 { // volunteer
		if adminCode == "" {
			renderRegisterError(w, "Admin code is required for volunteers")
			return
		}

		// Check if admin exists
		err = models.DB.QueryRow(`SELECT user_id FROM users WHERE role_id = 1 AND admin_code = $1`, adminCode).Scan(&adminID)
		if err != nil {
			if err == sql.ErrNoRows {
				renderRegisterError(w, "Invalid admin code")
				return
			}
			log.Printf("DB error checking admin code: %v", err)
			http.Error(w, "Internal server error", http.StatusInternalServerError)
			return
		}
	}

	// Insert user and get ID
	var userID int
	err = models.DB.QueryRow(`
		INSERT INTO users (first_name, last_name, email, phone, password, role_id)
		VALUES ($1, $2, $3, $4, $5, $6)
		RETURNING user_id
	`, firstName, lastName, email, phone, string(hashedPassword), roleID).Scan(&userID)
	if err != nil {
		log.Printf("User registration failed: %v", err)
		renderRegisterError(w, "Could not create account. Email might already be in use.")
		return
	}

	// Link volunteer to admin if role is volunteer
	if roleID == 3 {
		_, err = models.DB.Exec(`
			INSERT INTO admin_volunteers (admin_id, volunteer_id)
			VALUES ($1, $2)
		`, adminID, userID)
		if err != nil {
			log.Printf("Failed to link volunteer to admin: %v", err)
			http.Error(w, "Internal server error", http.StatusInternalServerError)
			return
		}
	}

	log.Printf("User registered successfully: %s %s (%s)", firstName, lastName, email)
	http.Redirect(w, r, "/", http.StatusSeeOther)
}



func LogoutHandler(w http.ResponseWriter, r *http.Request) {
	clearSessionCookie(w)
	http.Redirect(w, r, "/", http.StatusSeeOther)
}