Poll-system/app/main.go

package main

import (
	"context"
	"log"
	"net/http"
	"os"

	"github.com/golang-jwt/jwt/v5"
	"github.com/joho/godotenv"

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

	_ "github.com/lib/pq" // use PostgreSQL
)

// Helper function to determine navigation visibility based on role
func getNavFlags(role int) (bool, bool, bool) {
	showAdminNav := role == 1     // Admin role
	showLeaderNav := role == 2    // Team Leader role
	showVolunteerNav := role == 3 // Volunteer role
	return showAdminNav, showVolunteerNav, showLeaderNav
}

// Helper function to create template data with proper nav flags
func createTemplateData(title, activeSection string, role int, isAuthenticated bool, additionalData map[string]interface{}) map[string]interface{} {
	showAdminNav, showVolunteerNav, _ := getNavFlags(role)
	
	data := map[string]interface{}{
		"Title":             title,
		"IsAuthenticated":   isAuthenticated,
		"Role":              role,
		"ShowAdminNav":      showAdminNav,
		"ShowVolunteerNav":  showVolunteerNav,
		"ActiveSection":     activeSection,
	}
	
	// Add any additional data
	for key, value := range additionalData {
		data[key] = value
	}
	
	return data
}

func authMiddleware(next http.HandlerFunc) http.HandlerFunc {
	err := godotenv.Load()
	if err != nil {
		log.Fatalf("Error loading .env file: %v", err)
	}

	jwtSecret := os.Getenv("JWT_SECRET")
	var jwtKey = []byte(jwtSecret)

	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 {
			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)
	}
}

// Admin middleware to check if user has admin role
func adminMiddleware(next http.HandlerFunc) http.HandlerFunc {
	return authMiddleware(func(w http.ResponseWriter, r *http.Request) {
		role, ok := r.Context().Value("user_role").(int)
		if !ok || role != 1 {
			http.Redirect(w, r, "/", http.StatusSeeOther)
			return
		}
		next.ServeHTTP(w, r)
	})
}

// Volunteer middleware to check if user has volunteer role
func volunteerMiddleware(next http.HandlerFunc) http.HandlerFunc {
	return authMiddleware(func(w http.ResponseWriter, r *http.Request) {
		role, ok := r.Context().Value("user_role").(int)
		if !ok || (role != 3 && role != 2) {
			http.Redirect(w, r, "/", http.StatusSeeOther)
			return
		}
		next.ServeHTTP(w, r)
	})
}

func schedualHandler(w http.ResponseWriter, r *http.Request) {
	role := r.Context().Value("user_role").(int)
	data := createTemplateData("My Schedule", "schedual", role, true, nil)
	utils.Render(w, "Schedual/schedual.html", data)
}

func HomeHandler(w http.ResponseWriter, r *http.Request) {
	utils.Render(w, "dashboard/dashboard.html", map[string]interface{}{
		"Title":           "Admin Dashboard",
		"IsAuthenticated": false,
		"ActiveSection":   "dashboard",
	})
}

func main() {
	models.InitDB()
	
	// Static file servers
	fs := http.FileServer(http.Dir("static"))
	http.Handle("/static/", http.StripPrefix("/static/", fs))
	
	uploadsFs := http.FileServer(http.Dir("uploads"))
	http.Handle("/uploads/", http.StripPrefix("/uploads/", uploadsFs))

	// Public HTML Routes
	http.HandleFunc("/", HomeHandler)
	http.HandleFunc("/login", handlers.LoginHandler)
	http.HandleFunc("/register", handlers.RegisterHandler)

	//--- Protected HTML Routes
	http.HandleFunc("/logout", authMiddleware(handlers.LogoutHandler))

	// Common routes (both admin and volunteer can access)
	http.HandleFunc("/profile", authMiddleware(handlers.ProfileHandler))
	http.HandleFunc("/profile/update", authMiddleware(handlers.ProfileUpdateHandler))

	//--- Admin-only routes
	http.HandleFunc("/dashboard", adminMiddleware(handlers.AdminDashboardHandler))
	http.HandleFunc("/volunteers", adminMiddleware(handlers.VolunteerHandler))
	http.HandleFunc("/volunteer/edit", adminMiddleware(handlers.EditVolunteerHandler))

	http.HandleFunc("/team_builder", adminMiddleware(handlers.TeamBuilderHandler))
	http.HandleFunc("/team_builder/remove_volunteer", adminMiddleware(handlers.RemoveVolunteerHandler))

	http.HandleFunc("/addresses", adminMiddleware(handlers.AddressHandler))
	http.HandleFunc("/assign_address", adminMiddleware(handlers.AssignAddressHandler))
	http.HandleFunc("/remove_assigned_address", adminMiddleware(handlers.RemoveAssignedAddressHandler))

	http.HandleFunc("/addresses/upload-csv", adminMiddleware(handlers.CSVUploadHandler))


	http.HandleFunc("/posts", adminMiddleware(handlers.PostsHandler))

	//--- Volunteer-only routes  
	http.HandleFunc("/volunteer/dashboard", volunteerMiddleware(handlers.VolunteerPostsHandler))
	http.HandleFunc("/volunteer/Addresses", volunteerMiddleware(handlers.VolunteerAppointmentHandler))
	http.HandleFunc("/schedual", volunteerMiddleware(schedualHandler))

	// Poll routes (volunteer only)
	http.HandleFunc("/poll", volunteerMiddleware(handlers.PollHandler))
	
	log.Println("Server started on localhost:8080")
	log.Fatal(http.ListenAndServe("0.0.0.0:8080", nil))
}