package user

import (
    "context"
    "database/sql"
    encore "encore.dev"
    "encore.dev/beta/auth"
    "encore.dev/beta/errs"
    "encore.dev/rlog"
    "encore.dev/storage/cache"
    "errors"
    "github.com/go-playground/validator/v10"
    "github.com/golang-jwt/jwt/v5"
    "github.com/google/uuid"
    "github.com/matthewhartstonge/argon2"
    "strings"
    "time"
)

var authCacheCluster = cache.NewCluster("auth-cache-cluster", cache.ClusterConfig{
    EvictionPolicy: cache.VolatileTTL,
})

type AuthTokenInfo struct {
    TokenId   string
    UserAgent string
    IpAddress string
    Kind      AuthTokenKind
}

type AuthTokenKind int

const (
    AccessToken AuthTokenKind = iota
    RefreshToken
)

type AuthToken struct {
    UserId  string
    TokenId string
    Kind    AuthTokenKind
    jwt.RegisteredClaims
}

type AuthInfo struct {
    User        User
    AuthTokenId string
}

var AccessTokenValidity = cache.NewStructKeyspace[string, AuthTokenInfo](authCacheCluster, cache.KeyspaceConfig{
    KeyPattern:    "auth/accessToken/:key",
    DefaultExpiry: cache.ExpireIn(10 * time.Minute),
})

var RefreshTokenValidity = cache.NewStructKeyspace[string, AuthTokenInfo](authCacheCluster, cache.KeyspaceConfig{
    KeyPattern:    "auth/refreshToken/:key",
    DefaultExpiry: cache.ExpireIn(30 * 24 * time.Hour),
})

var validate = validator.New()
var secrets struct {
    JWTSecret string
}

type LoginParams struct {
    Email    string `json:"email" validate:"required,email"`
    Password string `json:"password" validate:"required,min=8,max=100"`
}

type LoginResult struct {
    AccessToken  string `json:"access_token"`
    RefreshToken string `json:"refresh_token"`
}

//encore:authhandler
func ValidateLogin(ctx context.Context, token string) (auth.UID, *AuthInfo, error) {
    if strings.HasPrefix(token, "Bearer") {
        token = strings.Split(token, " ")[1]
    }
    parsedToken, _, err := getTokenWithoutValidation(token)
    if err != nil {
        rlog.Warn("Failed to decode token", "err", err)
        return "", nil, &errs.Error{
            Code: errs.Unauthenticated,
        }
    }

    tokenKind := AuthTokenKind(parsedToken.Claims.(jwt.MapClaims)["Kind"].(float64))
    if tokenKind != AccessToken {
        rlog.Warn("Got something other than access token", "kind", tokenKind)
        return "", nil, &errs.Error{
            Code: errs.Unauthenticated,
        }
    }

    userId := parsedToken.Claims.(jwt.MapClaims)["UserId"].(string)
    var userPassword string
    err = userdb.QueryRow(ctx, `select password from users where id = $1`, userId).Scan(&userPassword)
    if err != nil {
        rlog.Warn("Failed to fetch user", "err", err)
        return "", nil, &errs.Error{
            Code: errs.Unauthenticated,
        }
    }

    validatedToken, err := verifyToken(token, secrets.JWTSecret+userPassword)
    if err != nil {
        rlog.Warn("Failed to validate login", "err", err)
        return "", nil, &errs.Error{
            Code: errs.Unauthenticated,
        }
    }

    _, err = AccessTokenValidity.Get(ctx, validatedToken.Claims.(jwt.MapClaims)["TokenId"].(string))
    if err != nil {
        rlog.Warn("Failed to get access token from cache", "err", err)
        return "", nil, &errs.Error{
            Code: errs.Unauthenticated,
        }
    }

    user := User{}
    err = userdb.QueryRow(ctx, `select * from users where id = $1`, userId).Scan(
        &user.Id,
        &user.FirstName,
        &user.LastName,
        &user.Email,
        &user.Password,
        &user.Joined,
    )
    if err != nil {
        rlog.Error("Failed to fetch user after validation", "err", err)
        return "", nil, &errs.Error{
            Code: errs.Internal,
        }
    }

    return auth.UID(userId), &AuthInfo{
        User:        user,
        AuthTokenId: validatedToken.Claims.(jwt.MapClaims)["TokenId"].(string),
    }, nil
}

//encore:api method=POST public
func Login(ctx context.Context, params *LoginParams) (LoginResult, error) {
    if err := validate.Struct(params); err != nil {
        return LoginResult{}, &errs.Error{
            Code:    errs.InvalidArgument,
            Message: err.Error(),
        }
    }

    user := User{}
    err := userdb.QueryRow(ctx, `select id, password from users where email = $1`, params.Email).Scan(
        &user.Id,
        &user.Password,
    )
    if err != nil {
        if errors.Is(err, sql.ErrNoRows) {
            return LoginResult{}, &errs.Error{
                Code:    errs.InvalidArgument,
                Message: "Invalid email or password",
            }
        } else {
            rlog.Warn("Failed to query for user", "email", params.Email, "err", err)
            return LoginResult{}, &errs.Error{
                Code: errs.Internal,
            }
        }
    }

    tokenId := uuid.New().String()

    signedAccessToken, err := signToken(AuthToken{
        UserId:  user.Id,
        TokenId: tokenId,
        Kind:    AccessToken,
        RegisteredClaims: jwt.RegisteredClaims{
            ExpiresAt: jwt.NewNumericDate(time.Now().Add(10 * time.Minute)),
        },
    }, secrets.JWTSecret+user.Password)
    if err != nil {
        rlog.Error("Failed to sign access token", "err", err)
        return LoginResult{}, &errs.Error{
            Code: errs.Internal,
        }
    }

    signedRefreshToken, err := signToken(AuthToken{
        UserId:  user.Id,
        TokenId: tokenId,
        Kind:    RefreshToken,
        RegisteredClaims: jwt.RegisteredClaims{
            ExpiresAt: jwt.NewNumericDate(time.Now().Add(30 * 24 * time.Hour)),
        },
    }, secrets.JWTSecret+user.Password)
    if err != nil {
        rlog.Error("Failed to sign refresh token", "err", err)
        return LoginResult{}, &errs.Error{
            Code: errs.Internal,
        }
    }

    accessTokenInfo := AuthTokenInfo{
        TokenId:   tokenId,
        UserAgent: encore.CurrentRequest().Headers.Get("User-Agent"),
        IpAddress: encore.CurrentRequest().Headers.Get("X-Forwarded-For"),
        Kind:      AccessToken,
    }
    err = AccessTokenValidity.Set(ctx, tokenId, accessTokenInfo)
    if err != nil {
        rlog.Error("Failed to save access token id in cache", "err", err)
        return LoginResult{}, &errs.Error{
            Code: errs.Internal,
        }
    }

    refreshTokenInfo := AuthTokenInfo{
        TokenId:   tokenId,
        UserAgent: encore.CurrentRequest().Headers.Get("User-Agent"),
        IpAddress: encore.CurrentRequest().Headers.Get("X-Forwarded-For"),
        Kind:      RefreshToken,
    }
    err = RefreshTokenValidity.Set(ctx, tokenId, refreshTokenInfo)
    if err != nil {
        rlog.Error("Failed to save refresh token id in cache", "err", err)
        return LoginResult{}, &errs.Error{
            Code: errs.Internal,
        }
    }

    res := LoginResult{
        AccessToken:  signedAccessToken,
        RefreshToken: signedRefreshToken,
    }
    return res, nil
}

type RegistrationParams struct {
    FirstName string `json:"first_name" validate:"required,min=1,max=64"`
    LastName  string `json:"last_name" validate:"required,min=1,max=64"`
    Email     string `json:"email" validate:"required,email,max=50"`
    Password  string `json:"password" validate:"required,min=8,max=100"`
}

type RegistrationResult struct {
    Message string `json:"message"`
}

//encore:api method=POST public
func Register(ctx context.Context, params *RegistrationParams) (RegistrationResult, error) {
    if err := validate.Struct(params); err != nil {
        return RegistrationResult{}, &errs.Error{
            Code:    errs.InvalidArgument,
            Message: err.Error(),
        }
    }

    rows, err := userdb.Query(ctx, `select * from users where email=($1)`, params.Email)
    if err != nil {
        rlog.Error("Failed to query for user", "email", params.Email, "err", err)
        return RegistrationResult{}, &errs.Error{
            Code: errs.Internal,
        }
    }

    if rows.Next() {
        return RegistrationResult{}, &errs.Error{
            Code:    errs.InvalidArgument,
            Message: "User already exists",
        }
    }

    argon := argon2.DefaultConfig()

    hashedPassword, err := argon.HashEncoded([]byte(params.Password))
    if err != nil {
        rlog.Error("Failed to hash password", "err", err)
        return RegistrationResult{}, &errs.Error{
            Code: errs.Internal,
        }
    }

    _, err = userdb.Exec(ctx, `insert into users (first_name, last_name, email, password) values ($1, $2, $3, $4)`,
        params.FirstName, params.LastName, params.Email, hashedPassword)
    if err != nil {
        rlog.Error("Failed to insert user", "email", params.Email, "err", err)
        return RegistrationResult{}, &errs.Error{
            Code: errs.Internal,
        }
    }

    return RegistrationResult{
        Message: "Registration successful!",
    }, nil
}

//encore:api method=POST auth
func Logout(ctx context.Context) error {
    authTokenId := auth.Data().(*AuthInfo).AuthTokenId
    _, err := AccessTokenValidity.Delete(ctx, authTokenId)
    if err != nil {
        rlog.Error("Failed to delete access token info", "err", err)
        return &errs.Error{
            Code: errs.Internal,
        }
    }
    _, err = RefreshTokenValidity.Delete(ctx, authTokenId)
    if err != nil {
        rlog.Error("Failed to delete refresh token info", "err", err)
        return &errs.Error{
            Code: errs.Internal,
        }
    }

    return nil
}

type RefreshAccessTokenParams struct {
    RefreshToken string `json:"refresh_token"`
}

type RefreshAccessTokenResult struct {
    AccessToken string `json:"access_token"`
}

//encore:api method=POST public
func RefreshAccessToken(ctx context.Context, params *RefreshAccessTokenParams) (RefreshAccessTokenResult, error) {
    token := params.RefreshToken
    parsedToken, _, err := getTokenWithoutValidation(token)
    if err != nil {
        rlog.Warn("Failed to decode token", "err", err)
        return RefreshAccessTokenResult{}, &errs.Error{
            Code: errs.Unauthenticated,
        }
    }

    tokenKind := AuthTokenKind(parsedToken.Claims.(jwt.MapClaims)["Kind"].(float64))
    if tokenKind != RefreshToken {
        rlog.Warn("Got something other than refresh token", "kind", tokenKind)
        return RefreshAccessTokenResult{}, &errs.Error{
            Code: errs.Unauthenticated,
        }
    }

    userId := parsedToken.Claims.(jwt.MapClaims)["UserId"].(string)
    var userPassword string
    err = userdb.QueryRow(ctx, `select password from users where id = $1`, userId).Scan(&userPassword)
    if err != nil {
        rlog.Warn("Failed to fetch user", "err", err)
        return RefreshAccessTokenResult{}, &errs.Error{
            Code: errs.Unauthenticated,
        }
    }

    validatedToken, err := verifyToken(token, secrets.JWTSecret+userPassword)
    if err != nil {
        rlog.Warn("Failed to validate login", "err", err)
        return RefreshAccessTokenResult{}, &errs.Error{
            Code: errs.Unauthenticated,
        }
    }

    refreshTokenInfo, err := AccessTokenValidity.Get(ctx, validatedToken.Claims.(jwt.MapClaims)["TokenId"].(string))
    if err != nil {
        rlog.Warn("Failed to get refresh token from cache", "err", err)
        return RefreshAccessTokenResult{}, &errs.Error{
            Code: errs.Unauthenticated,
        }
    }

    accessToken, err := signToken(AuthToken{
        UserId:  userId,
        TokenId: refreshTokenInfo.TokenId,
        Kind:    AccessToken,
        RegisteredClaims: jwt.RegisteredClaims{
            ExpiresAt: jwt.NewNumericDate(time.Now().Add(10 * time.Minute)),
        },
    }, secrets.JWTSecret+userPassword)

    if err != nil {
        rlog.Error("Failed to sign refresh token", "err", err)
        return RefreshAccessTokenResult{}, &errs.Error{
            Code: errs.Internal,
        }
    }

    return RefreshAccessTokenResult{
        accessToken,
    }, nil
}

type CurrentUserResponse struct {
    CurrentUser User `json:"current_user"`
}

//encore:api method=GET auth
func GetCurrentUser(ctx context.Context) (CurrentUserResponse, error) {
    user := auth.Data().(*AuthInfo).User
    user.Password = "<redacted>"
    return CurrentUserResponse{
        CurrentUser: user,
    }, nil
}