package jws

import (
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"math/big"
	"strings"
	"time"

	appcrypto "netisjwt/internal/crypto"
	"netisjwt/internal/statuslist"
)

type StatusClaim struct {
	EncodedList string `json:"encodedList"`
	Index       int    `json:"index"`
}

type Claims struct {
	Issuer    string      `json:"iss"`
	IssuedAt  int64       `json:"iat"`
	ExpiresAt int64       `json:"exp"`
	Status    StatusClaim `json:"status"`
}

type TokenHeader struct {
	Algorithm string   `json:"alg"`
	Type      string   `json:"typ"`
	JWK       JWKValue `json:"jwk"`
}

type JWKValue struct {
	Kty string `json:"kty"`
	Crv string `json:"crv"`
	X   string `json:"x"`
	Y   string `json:"y"`
}

// SestaviTrditve pripravi JWT payload za status endpoint.
func SestaviTrditve(issuer string, list *statuslist.StatusList, index int, now time.Time) (Claims, error) {
	encoded, err := list.KodiranSeznam()
	if err != nil {
		return Claims{}, err
	}
	return Claims{
		Issuer:    issuer,
		IssuedAt:  now.Unix(),
		ExpiresAt: now.Add(24 * time.Hour).Unix(),
		Status: StatusClaim{
			EncodedList: encoded,
			Index:       index,
		},
	}, nil
}

// UstvariZetonIzPEM ustvari JWS compact token iz PEM kljuca.
func UstvariZetonIzPEM(privatePEM []byte, claims Claims) (string, error) {
	privateKey, err := appcrypto.PreberiZasebniKljucPEM(privatePEM)
	if err != nil {
		return "", err
	}
	return UstvariZeton(privateKey, claims)
}

// UstvariZeton ustvari JWS compact token iz ECDSA kljuca.
func UstvariZeton(privateKey *ecdsa.PrivateKey, claims Claims) (string, error) {
	header := TokenHeader{
		Algorithm: "ES256",
		Type:      "JWT",
		JWK:       javniKljucVJWK(&privateKey.PublicKey),
	}

	headerPart, err := kodirajDel(header)
	if err != nil {
		return "", err
	}
	claimsPart, err := kodirajDel(claims)
	if err != nil {
		return "", err
	}

	signingInput := headerPart + "." + claimsPart
	hash := sha256.Sum256([]byte(signingInput))
	signatureRaw, err := ecdsa.SignASN1(rand.Reader, privateKey, hash[:])
	if err != nil {
		return "", err
	}
	signature := base64.RawURLEncoding.EncodeToString(signatureRaw)
	return signingInput + "." + signature, nil
}

// PreveriZeton preveri podpis in vrne dekodirane trditve.
func PreveriZeton(token string) (Claims, error) {
	parts := strings.Split(token, ".")
	if len(parts) != 3 {
		return Claims{}, errors.New("token must have 3 parts")
	}

	headerBytes, err := base64.RawURLEncoding.DecodeString(parts[0])
	if err != nil {
		return Claims{}, err
	}
	claimsBytes, err := base64.RawURLEncoding.DecodeString(parts[1])
	if err != nil {
		return Claims{}, err
	}

	var header TokenHeader
	if err := json.Unmarshal(headerBytes, &header); err != nil {
		return Claims{}, err
	}
	var claims Claims
	if err := json.Unmarshal(claimsBytes, &claims); err != nil {
		return Claims{}, err
	}

	pub, err := jwkVJavniKljuc(header.JWK)
	if err != nil {
		return Claims{}, err
	}
	signature, err := base64.RawURLEncoding.DecodeString(parts[2])
	if err != nil {
		return Claims{}, err
	}

	hash := sha256.Sum256([]byte(parts[0] + "." + parts[1]))
	if !ecdsa.VerifyASN1(pub, hash[:], signature) {
		return Claims{}, errors.New("invalid token signature")
	}
	return claims, nil
}

// PreveriTrditve preveri veljavnost iat, exp in iss vrednosti.
func PreveriTrditve(claims Claims, expectedIssuer string, now time.Time) error {
	if expectedIssuer != "" && claims.Issuer != expectedIssuer {
		return fmt.Errorf("invalid issuer")
	}
	nowUnix := now.Unix()
	if claims.IssuedAt > nowUnix {
		return errors.New("iat is in future")
	}
	if claims.ExpiresAt < nowUnix {
		return errors.New("token expired")
	}
	return nil
}

// kodirajDel serializira del tokena v Base64URL JSON.
func kodirajDel(v any) (string, error) {
	raw, err := json.Marshal(v)
	if err != nil {
		return "", err
	}
	return base64.RawURLEncoding.EncodeToString(raw), nil
}

// javniKljucVJWK pretvori ECDSA javni kljuc v JWK zapis.
func javniKljucVJWK(pub *ecdsa.PublicKey) JWKValue {
	return JWKValue{
		Kty: "EC",
		Crv: "P-256",
		X:   base64.RawURLEncoding.EncodeToString(pub.X.FillBytes(make([]byte, 32))),
		Y:   base64.RawURLEncoding.EncodeToString(pub.Y.FillBytes(make([]byte, 32))),
	}
}

// jwkVJavniKljuc pretvori JWK zapis v ECDSA javni kljuc.
func jwkVJavniKljuc(jwk JWKValue) (*ecdsa.PublicKey, error) {
	if jwk.Kty != "EC" || jwk.Crv != "P-256" {
		return nil, errors.New("unsupported key type")
	}
	xBytes, err := base64.RawURLEncoding.DecodeString(jwk.X)
	if err != nil {
		return nil, err
	}
	yBytes, err := base64.RawURLEncoding.DecodeString(jwk.Y)
	if err != nil {
		return nil, err
	}
	pub := &ecdsa.PublicKey{
		Curve: elliptic.P256(),
		X:     new(big.Int).SetBytes(xBytes),
		Y:     new(big.Int).SetBytes(yBytes),
	}
	if !pub.Curve.IsOnCurve(pub.X, pub.Y) {
		return nil, errors.New("jwk is not on P-256 curve")
	}
	return pub, nil
}