package main

import (
    "net/http"
    "time"

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

var (
    jwtSecret = []byte("your_strong_secret_key")
    users     = make(map[string]string)
)

type Credentials struct {
    Username string `json:"username" binding:"required"`
    Password string `json:"password" binding:"required"`
}

type Claims struct {
    Username string `json:"username"`
    jwt.RegisteredClaims
}

func main() {
    r := gin.Default()

    r.POST("/register", registerHandler)
    r.POST("/login", loginHandler)
    r.GET("/protected", authMiddleware, protectedHandler)

    r.Run(":8080")
}

func registerHandler(c *gin.Context) {
    var creds Credentials
    if err := c.ShouldBindJSON(&creds); err != nil {
        c.JSON(http.StatusBadRequest, gin.H{"error": "Invalid request"})
        return
    }

    if _, exists := users[creds.Username]; exists {
        c.JSON(http.StatusConflict, gin.H{"error": "User already exists"})
        return
    }

    hashedPassword, err := bcrypt.GenerateFromPassword([]byte(creds.Password), bcrypt.DefaultCost)
    if err != nil {
        c.JSON(http.StatusInternalServerError, gin.H{"error": "Could not hash password"})
        return
    }

    users[creds.Username] = string(hashedPassword)
    c.JSON(http.StatusCreated, gin.H{"message": "User created"})
}

func loginHandler(c *gin.Context) {
    var creds Credentials
    if err := c.ShouldBindJSON(&creds); err != nil {
        c.JSON(http.StatusBadRequest, gin.H{"error": "Invalid request"})
        return
    }

    storedHash, exists := users[creds.Username]
    if !exists {
        c.JSON(http.StatusUnauthorized, gin.H{"error": "Invalid credentials"})
        return
    }

    if err := bcrypt.CompareHashAndPassword([]byte(storedHash), []byte(creds.Password)); err != nil {
        c.JSON(http.StatusUnauthorized, gin.H{"error": "Invalid credentials"})
        return
    }

    expirationTime := time.Now().Add(15 * time.Minute)
    claims := &Claims{
        Username: creds.Username,
        RegisteredClaims: jwt.RegisteredClaims{
            ExpiresAt: jwt.NewNumericDate(expirationTime),
        },
    }

    token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
    tokenString, err := token.SignedString(jwtSecret)
    if err != nil {
        c.JSON(http.StatusInternalServerError, gin.H{"error": "Could not generate token"})
        return
    }

    c.JSON(http.StatusOK, gin.H{"token": tokenString})
}

func authMiddleware(c *gin.Context) {
    tokenString := c.GetHeader("Authorization")
    if tokenString == "" {
        c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{"error": "Missing authorization header"})
        return
    }

    token, err := jwt.ParseWithClaims(tokenString, &Claims{}, func(token *jwt.Token) (interface{}, error) {
        return jwtSecret, nil
    })

    if err != nil || !token.Valid {
        c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{"error": "Invalid token"})
        return
    }

    claims, ok := token.Claims.(*Claims)
    if !ok {
        c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{"error": "Invalid token claims"})
        return
    }

    c.Set("username", claims.Username)
    c.Next()
}

func protectedHandler(c *gin.Context) {
    username, exists := c.Get("username")
    if !exists {
        c.JSON(http.StatusInternalServerError, gin.H{"error": "User not found in context"})
        return
    }

    c.JSON(http.StatusOK, gin.H{"message": "Hello " + username.(string) + "! This is protected content"})
}