package repository

import (
	service "authservice/api/authservice"
	"authservice/internal/model"
	"authservice/internal/tools"
	"crypto"
	"crypto/rsa"
	"crypto/sha512"
	"crypto/x509"
	"encoding/base64"
	"encoding/pem"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
	"gopkg.in/square/go-jose.v2/jwt"
	"time"
)

type CryptoRepository struct {
	JwtCrypto *model.JwtCrypto
}

const (
	UserIdKey                  = "userId"
	AccessTokenKey             = "accessToken"
	RefreshTokenKey            = "refreshToken"
	PublicKeyKey               = "key"
	ExpirationTimeKey          = "exp"
	IssuerValue                = "SecNote"
	ExpirationTimeAccessValue  = time.Minute * 15
	ExpirationTimeRefreshValue = time.Hour * 24 * 7
)

func (repository *CryptoRepository) ParseKey(key string) (*rsa.PublicKey, error) {
	pubPEM, err := base64.StdEncoding.DecodeString(key)
	if err != nil {
		err = tools.ConvertError(codes.InvalidArgument, "Failed to decode provided public key", err)
		return nil, err
	}
	block, _ := pem.Decode(pubPEM)
	if block == nil {
		err = tools.CreateError(codes.InvalidArgument, "Failed to parse PEM block containing the public key", "Invalid PEM format")
		return nil, err
	}

	pubKey, err := x509.ParsePKIXPublicKey(block.Bytes)

	if err != nil {
		err = tools.ConvertError(codes.InvalidArgument, "Failed to parse DER encoded public key", err)
		return nil, err
	}
	return pubKey.(*rsa.PublicKey), nil
}

func (repository *CryptoRepository) VerifyJWT(token *jwt.JSONWebToken, isRefreshToken bool) error {
	claims, err := repository.GetJWTClaims(token) // Returns formatted error
	if err != nil {
		return err
	}
	if claims[RefreshTokenKey] != isRefreshToken {
		return tools.CreateError(codes.InvalidArgument, "Token is in invalid type, check type of token", "invalid token type")
	}

	expiration := int64(claims[ExpirationTimeKey].(float64))
	if time.Now().After(time.Unix(expiration, 0)) {
		return status.Error(codes.Unauthenticated, "token is expired")
	}

	return nil
}

func (repository *CryptoRepository) VerifyRequestSignature(token *jwt.JSONWebToken, encodedMessage string, signature string, encodedToken string) (bool, error) {
	headers, err := repository.GetJWTClaims(token) // Returns formatted error
	if err != nil {
		return false, err
	}
	pubKeyRaw := headers[PublicKeyKey].(string)
	pubKey, err := repository.ParseKey(pubKeyRaw) // Returns formatted error
	if err != nil {
		return false, err
	}
	sign, err := base64.StdEncoding.DecodeString(signature)
	if err != nil {
		err = tools.ConvertError(codes.InvalidArgument, "Invalid signature base64 encoding", err)
		return false, err
	}
	messageByte := []byte(encodedMessage)
	tokenByte := []byte(encodedToken)
	sha512Hash := sha512.New()
	sha512Hash.Write(append(messageByte, tokenByte...))
	hash := sha512Hash.Sum(nil)
	err = rsa.VerifyPKCS1v15(pubKey, crypto.SHA512, hash, sign)
	if err != nil {
		err = tools.ConvertError(codes.Unauthenticated, "Invalid signature, signature do not coresponde to encodedMessage", err)
		return false, err
	}
	return true, nil
}

func (repository *CryptoRepository) GetJWTClaims(token *jwt.JSONWebToken) (map[string]interface{}, error) {
	claims := make(map[string]interface{})
	err := token.Claims(&repository.JwtCrypto.Key.PublicKey, &claims)
	if err != nil {
		err = tools.ConvertError(codes.Aborted, "Can`t deserialize claims from JWT token", err)
		return nil, err
	}
	return claims, nil
}

func (repository *CryptoRepository) ParseJWTToken(token string) (*jwt.JSONWebToken, error) {
	parsedJWT, err := jwt.ParseSigned(token)
	if err != nil {
		err = tools.ConvertError(codes.Unauthenticated, "Can`t parse JWT token", err)
		return nil, err
	}
	return parsedJWT, nil
}

func (repository *CryptoRepository) CreateJWT(id string, key string) (*service.JWT, error) {

	currentTime := time.Now()

	accessPublicClaims := jwt.Claims{
		Issuer:   IssuerValue,
		IssuedAt: jwt.NewNumericDate(currentTime),
		Expiry:   jwt.NewNumericDate(currentTime.Add(ExpirationTimeAccessValue)),
	}

	accessPrivateClaims := map[string]interface{}{
		UserIdKey:       id,
		AccessTokenKey:  true,
		RefreshTokenKey: false,
		PublicKeyKey:    key,
	}

	builder := jwt.Signed(repository.JwtCrypto.Signer).Claims(accessPublicClaims).Claims(accessPrivateClaims)
	accessToken, err := builder.CompactSerialize()

	if err != nil {
		err = tools.ConvertError(codes.FailedPrecondition, "Failed to create access JWT", err)
		return nil, err
	}

	refreshPublicClaims := jwt.Claims{
		Issuer:    IssuerValue,
		IssuedAt:  jwt.NewNumericDate(currentTime),
		Expiry:    jwt.NewNumericDate(currentTime.Add(ExpirationTimeRefreshValue)),
		NotBefore: jwt.NewNumericDate(currentTime.Add(ExpirationTimeAccessValue)),
	}

	refreshPrivateClaims := map[string]interface{}{
		UserIdKey:       id,
		AccessTokenKey:  false,
		RefreshTokenKey: true,
		PublicKeyKey:    key,
	}

	builder = jwt.Signed(repository.JwtCrypto.Signer).Claims(refreshPublicClaims).Claims(refreshPrivateClaims)
	refreshToken, err := builder.CompactSerialize()

	if err != nil {
		err = tools.ConvertError(codes.FailedPrecondition, "Failed to create refresh JWT", err)
		return nil, err
	}

	return &service.JWT{
		AccessToken:  accessToken,
		RefreshToken: refreshToken,
	}, nil
}
