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

package handlers

import (
	"context"
	"log"
	"net/http"
	"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 LoginPage(w http.ResponseWriter, r *http.Request) {
// 	utils.Render(w, "login.html", map[string]interface{}{
// 		"Title":           "Login",
// 		"IsAuthenticated": false,
// 	})
// }

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 == "" {
		renderLoginError(w, "Email and password are required")
		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)
		renderLoginError(w, "Invalid email or password")
		return
	}

	// Verify password
	err = bcrypt.CompareHashAndPassword([]byte(storedHash), []byte(password))
	if err != nil {
		log.Printf("Password verification failed for user ID %d", userID)
		renderLoginError(w, "Invalid email or password")
		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.Error(w, "Could not log in", http.StatusInternalServerError)
		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")

	// 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
	}

	// Insert user into database
	_, err = models.DB.Exec(`
		INSERT INTO "users" (first_name, last_name, email, phone, password, role_id)
		VALUES ($1, $2, $3, $4, $5, $6)
	`, firstName, lastName, email, phone, string(hashedPassword), role)

	if err != nil {
		log.Printf("User registration failed for email %s: %v", email, err)
		renderRegisterError(w, "Could not create account. Email might already be in use.")
		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)
}

// // Admin Dashboard Handler
// func AdminDashboardHandler(w http.ResponseWriter, r *http.Request) {
// 	role := r.Context().Value("user_role").(int)
// 	userID := r.Context().Value("user_id").(int)
	
// 	// TODO: Fetch real data from database
// 	dashboardData := map[string]interface{}{
// 		"UserID":           userID,
// 		"TotalUsers":       100,  // Example: get from database
// 		"TotalVolunteers":  50,   // Example: get from database
// 		"TotalAddresses":   200,  // Example: get from database
// 		"RecentActivity":   []string{"User logged in", "New volunteer registered"}, // Example
// 	}
	
// 	data := createTemplateData("Admin Dashboard", "dashboard", role, true, dashboardData)
// 	utils.Render(w, "dashboard/dashboard.html", data)
// }

// // Volunteer Management Handler
// func VolunteerHandler(w http.ResponseWriter, r *http.Request) {
// 	role := r.Context().Value("user_role").(int)
	
// 	// TODO: Fetch real volunteer data from database
// 	volunteerData := map[string]interface{}{
// 		"Volunteers": []map[string]interface{}{
// 			{"ID": 1, "Name": "John Doe", "Email": "john@example.com", "Status": "Active"},
// 			{"ID": 2, "Name": "Jane Smith", "Email": "jane@example.com", "Status": "Active"},
// 		}, // Example: get from database
// 	}
	
// 	data := createTemplateData("Volunteers", "volunteer", role, true, volunteerData)
// 	utils.Render(w, "volunteers/volunteers.html", data)
// }

// // Address Management Handler
// func AddressHandler(w http.ResponseWriter, r *http.Request) {
// 	role := r.Context().Value("user_role").(int)
	
// 	// TODO: Fetch real address data from database
// 	addressData := map[string]interface{}{
// 		"Addresses": []map[string]interface{}{
// 			{"ID": 1, "Street": "123 Main St", "City": "Calgary", "Status": "Validated"},
// 			{"ID": 2, "Street": "456 Oak Ave", "City": "Calgary", "Status": "Pending"},
// 		}, // Example: get from database
// 	}
	
// 	data := createTemplateData("Addresses", "address", role, true, addressData)
// 	utils.Render(w, "addresses/addresses.html", data)
// }

// // Reports Handler
// func ReportHandler(w http.ResponseWriter, r *http.Request) {
// 	role := r.Context().Value("user_role").(int)
	
// 	// TODO: Fetch real report data from database
// 	reportData := map[string]interface{}{
// 		"Reports": []map[string]interface{}{
// 			{"ID": 1, "Name": "Weekly Summary", "Date": "2025-08-25", "Status": "Complete"},
// 			{"ID": 2, "Name": "Monthly Analytics", "Date": "2025-08-01", "Status": "Pending"},
// 		}, // Example: get from database
// 	}
	
// 	data := createTemplateData("Reports", "report", role, true, reportData)
// 	utils.Render(w, "reports/reports.html", data)
// }

// // Profile Handler (works for both admin and volunteer)
// func ProfileHandler(w http.ResponseWriter, r *http.Request) {
// 	role := r.Context().Value("user_role").(int)
// 	userID := r.Context().Value("user_id").(int)
	
// 	// Fetch real user data from database
// 	var firstName, lastName, email, phone string
// 	err := models.DB.QueryRow(`
// 		SELECT first_name, last_name, email, phone
// 		FROM "users"
// 		WHERE user_id = $1
// 	`, userID).Scan(&firstName, &lastName, &email, &phone)
	
// 	profileData := map[string]interface{}{
// 		"UserID": userID,
// 	}
	
// 	if err != nil {
// 		log.Printf("Error fetching user profile for ID %d: %v", userID, err)
// 		profileData["Error"] = "Could not load profile data"
// 	} else {
// 		profileData["FirstName"] = firstName
// 		profileData["LastName"] = lastName
// 		profileData["Email"] = email
// 		profileData["Phone"] = phone
// 	}
	
// 	data := createTemplateData("Profile", "profile", role, true, profileData)
// 	utils.Render(w, "profile/profile.html", data)
// }

// // Volunteer Dashboard Handler
// func VolunteerDashboardHandler(w http.ResponseWriter, r *http.Request) {
// 	role := r.Context().Value("user_role").(int)
// 	userID := r.Context().Value("user_id").(int)
	
// 	// TODO: Fetch volunteer-specific data from database
// 	dashboardData := map[string]interface{}{
// 		"UserID":           userID,
// 		"AssignedTasks":    5,    // Example: get from database
// 		"CompletedTasks":   12,   // Example: get from database
// 		"UpcomingEvents":   []string{"Community Meeting - Aug 30", "Training Session - Sep 5"}, // Example
// 	}
	
// 	data := createTemplateData("Volunteer Dashboard", "dashboard", role, true, dashboardData)
// 	utils.Render(w, "volunteer/dashboard.html", data)
// }

// // Schedule Handler for Volunteers
// func ScheduleHandler(w http.ResponseWriter, r *http.Request) {
// 	role := r.Context().Value("user_role").(int)
// 	userID := r.Context().Value("user_id").(int)
	
// 	// TODO: Fetch schedule data from database
// 	scheduleData := map[string]interface{}{
// 		"UserID": userID,
// 		"Schedule": []map[string]interface{}{
// 			{"Date": "2025-08-26", "Time": "10:00 AM", "Task": "Door-to-door survey", "Location": "Downtown"},
// 			{"Date": "2025-08-28", "Time": "2:00 PM", "Task": "Data entry", "Location": "Office"},
// 		}, // Example: get from database
// 	}
	
// 	data := createTemplateData("My Schedule", "schedual", role, true, scheduleData)
// 	utils.Render(w, "volunteer/schedule.html", data)
// }

// Enhanced middleware to check JWT auth and add user context
func AuthMiddleware(next http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		cookie, err := r.Cookie("session")
		if err != nil {
			http.Redirect(w, r, "/login", http.StatusSeeOther)
			return
		}

		claims := &models.Claims{}
		token, err := jwt.ParseWithClaims(cookie.Value, claims, func(token *jwt.Token) (interface{}, error) {
			return jwtKey, nil
		})

		if err != nil || !token.Valid {
			log.Printf("Invalid token: %v", err)
			clearSessionCookie(w) // Clear invalid cookie
			http.Redirect(w, r, "/login", http.StatusSeeOther)
			return
		}

		// Add user info to context
		ctx := context.WithValue(r.Context(), "user_id", claims.UserID)
		ctx = context.WithValue(ctx, "user_role", claims.Role)
		r = r.WithContext(ctx)

		next.ServeHTTP(w, r)
	}
}