package licensekit

import (
	"crypto/ed25519"
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"net/url"
	"sort"
	"strings"
	"time"
)

const (
	TokenTypeLicense = "LICENSE"
	TokenTypeLease   = "LEASE"
	AlgorithmEdDSA   = "EdDSA"
	SchemaVersion    = 1
)

type VerificationMode string

const (
	ModeOffline VerificationMode = "offline"
	ModeHybrid  VerificationMode = "hybrid"
	ModeOnline  VerificationMode = "online"
)

type Header struct {
	Algorithm string `json:"alg"`
	Type      string `json:"typ"`
	KeyID     string `json:"kid"`
	Version   int    `json:"v"`
}

type VerificationPolicy struct {
	Mode                VerificationMode `json:"mode"`
	LeaseTTLSeconds     int64            `json:"leaseTtlSeconds,omitempty"`
	OfflineGraceSeconds int64            `json:"offlineGraceSeconds,omitempty"`
	ServerURL           string           `json:"serverUrl,omitempty"`
}

type Claims struct {
	Version      int                `json:"version"`
	LicenseID    string             `json:"licenseId"`
	Issuer       string             `json:"issuer"`
	Customer     string             `json:"customer"`
	Product      string             `json:"product"`
	Edition      string             `json:"edition"`
	Features     []string           `json:"features,omitempty"`
	Limits       map[string]int64   `json:"limits,omitempty"`
	Domains      []string           `json:"domains,omitempty"`
	InstanceIDs  []string           `json:"instanceIds,omitempty"`
	IssuedAt     int64              `json:"issuedAt"`
	NotBefore    int64              `json:"notBefore,omitempty"`
	ExpiresAt    int64              `json:"expiresAt"`
	Verification VerificationPolicy `json:"verification"`
	Metadata     map[string]string  `json:"metadata,omitempty"`
}

type LeaseClaims struct {
	Version    int      `json:"version"`
	LeaseID    string   `json:"leaseId"`
	LicenseID  string   `json:"licenseId"`
	Product    string   `json:"product"`
	Customer   string   `json:"customer"`
	Edition    string   `json:"edition"`
	Features   []string `json:"features,omitempty"`
	Host       string   `json:"host,omitempty"`
	InstanceID string   `json:"instanceId,omitempty"`
	IssuedAt   int64    `json:"issuedAt"`
	ExpiresAt  int64    `json:"expiresAt"`
}

type TrustStore struct {
	LicenseKeys map[string]string `json:"licenseKeys"`
	LeaseKeys   map[string]string `json:"leaseKeys"`
}

type VerifiedLicense struct {
	Header Header
	Claims Claims
}

type VerifiedLease struct {
	Header Header
	Claims LeaseClaims
}

func NewTrustStore() TrustStore {
	return TrustStore{LicenseKeys: map[string]string{}, LeaseKeys: map[string]string{}}
}

func ParseTrustStore(data []byte) (TrustStore, error) {
	var store TrustStore
	dec := json.NewDecoder(strings.NewReader(string(data)))
	dec.DisallowUnknownFields()
	if err := dec.Decode(&store); err != nil {
		return TrustStore{}, fmt.Errorf("decode trust store: %w", err)
	}
	if store.LicenseKeys == nil {
		store.LicenseKeys = map[string]string{}
	}
	if store.LeaseKeys == nil {
		store.LeaseKeys = map[string]string{}
	}
	for kid, encoded := range store.LicenseKeys {
		if strings.TrimSpace(kid) == "" {
			return TrustStore{}, errors.New("license trust store contains an empty key id")
		}
		if _, err := DecodePublicKey(encoded); err != nil {
			return TrustStore{}, fmt.Errorf("license key %q: %w", kid, err)
		}
	}
	for kid, encoded := range store.LeaseKeys {
		if strings.TrimSpace(kid) == "" {
			return TrustStore{}, errors.New("lease trust store contains an empty key id")
		}
		if _, err := DecodePublicKey(encoded); err != nil {
			return TrustStore{}, fmt.Errorf("lease key %q: %w", kid, err)
		}
	}
	return store, nil
}

func MarshalTrustStore(store TrustStore) ([]byte, error) {
	if store.LicenseKeys == nil {
		store.LicenseKeys = map[string]string{}
	}
	if store.LeaseKeys == nil {
		store.LeaseKeys = map[string]string{}
	}
	return json.MarshalIndent(store, "", "  ")
}

func SignLicense(privateKey ed25519.PrivateKey, keyID string, claims Claims) (string, error) {
	if err := validateLicenseClaims(claims, time.Unix(claims.IssuedAt, 0), false); err != nil {
		return "", err
	}
	return sign(TokenTypeLicense, keyID, privateKey, claims)
}

func SignLease(privateKey ed25519.PrivateKey, keyID string, claims LeaseClaims) (string, error) {
	if err := validateLeaseClaims(claims, time.Unix(claims.IssuedAt, 0), false); err != nil {
		return "", err
	}
	return sign(TokenTypeLease, keyID, privateKey, claims)
}

func sign(tokenType, keyID string, privateKey ed25519.PrivateKey, claims any) (string, error) {
	if len(privateKey) != ed25519.PrivateKeySize {
		return "", errors.New("invalid Ed25519 private key")
	}
	keyID = strings.TrimSpace(keyID)
	if keyID == "" {
		return "", errors.New("key id is required")
	}
	header := Header{Algorithm: AlgorithmEdDSA, Type: tokenType, KeyID: keyID, Version: SchemaVersion}
	headerJSON, err := json.Marshal(header)
	if err != nil {
		return "", fmt.Errorf("marshal token header: %w", err)
	}
	payloadJSON, err := json.Marshal(claims)
	if err != nil {
		return "", fmt.Errorf("marshal token payload: %w", err)
	}
	headerPart := base64.RawURLEncoding.EncodeToString(headerJSON)
	payloadPart := base64.RawURLEncoding.EncodeToString(payloadJSON)
	signingInput := headerPart + "." + payloadPart
	signature := ed25519.Sign(privateKey, []byte(signingInput))
	return signingInput + "." + base64.RawURLEncoding.EncodeToString(signature), nil
}

func VerifyLicense(store TrustStore, token string, now time.Time) (VerifiedLicense, error) {
	header, payload, err := verifyToken(store.LicenseKeys, TokenTypeLicense, token)
	if err != nil {
		return VerifiedLicense{}, err
	}
	var claims Claims
	if err := decodeStrict(payload, &claims); err != nil {
		return VerifiedLicense{}, fmt.Errorf("decode license payload: %w", err)
	}
	if err := validateLicenseClaims(claims, now, true); err != nil {
		return VerifiedLicense{}, err
	}
	claims.Features = UniqueSorted(claims.Features)
	claims.Domains = UniqueSorted(claims.Domains)
	claims.InstanceIDs = UniqueSorted(claims.InstanceIDs)
	return VerifiedLicense{Header: header, Claims: claims}, nil
}

func VerifyLease(store TrustStore, token string, now time.Time, allowGrace time.Duration) (VerifiedLease, error) {
	header, payload, err := verifyToken(store.LeaseKeys, TokenTypeLease, token)
	if err != nil {
		return VerifiedLease{}, err
	}
	var claims LeaseClaims
	if err := decodeStrict(payload, &claims); err != nil {
		return VerifiedLease{}, fmt.Errorf("decode lease payload: %w", err)
	}
	if err := validateLeaseClaims(claims, now, false); err != nil {
		return VerifiedLease{}, err
	}
	if now.Unix() >= claims.ExpiresAt+int64(allowGrace.Seconds()) {
		return VerifiedLease{}, errors.New("lease has expired")
	}
	claims.Features = UniqueSorted(claims.Features)
	return VerifiedLease{Header: header, Claims: claims}, nil
}

func verifyToken(keys map[string]string, expectedType, token string) (Header, []byte, error) {
	parts := strings.Split(strings.TrimSpace(token), ".")
	if len(parts) != 3 {
		return Header{}, nil, errors.New("token has invalid format")
	}
	headerBytes, err := base64.RawURLEncoding.DecodeString(parts[0])
	if err != nil {
		return Header{}, nil, errors.New("token header is not valid base64url")
	}
	var header Header
	if err := decodeStrict(headerBytes, &header); err != nil {
		return Header{}, nil, fmt.Errorf("decode token header: %w", err)
	}
	if header.Algorithm != AlgorithmEdDSA || header.Type != expectedType || header.Version != SchemaVersion {
		return Header{}, nil, errors.New("unsupported token header")
	}
	encodedKey, ok := keys[header.KeyID]
	if !ok {
		return Header{}, nil, fmt.Errorf("token is signed by unknown key %q", header.KeyID)
	}
	publicKey, err := DecodePublicKey(encodedKey)
	if err != nil {
		return Header{}, nil, fmt.Errorf("decode trusted key %q: %w", header.KeyID, err)
	}
	signature, err := base64.RawURLEncoding.DecodeString(parts[2])
	if err != nil {
		return Header{}, nil, errors.New("token signature is not valid base64url")
	}
	signingInput := parts[0] + "." + parts[1]
	if !ed25519.Verify(publicKey, []byte(signingInput), signature) {
		return Header{}, nil, errors.New("token signature verification failed")
	}
	payload, err := base64.RawURLEncoding.DecodeString(parts[1])
	if err != nil {
		return Header{}, nil, errors.New("token payload is not valid base64url")
	}
	return header, payload, nil
}

func ValidateLicenseContext(claims Claims, product, baseURL, instanceID string) error {
	if strings.TrimSpace(product) == "" {
		return errors.New("client product id is required")
	}
	if claims.Product != product {
		return fmt.Errorf("license is for product %q, not %q", claims.Product, product)
	}
	if err := ValidateDomain(claims.Domains, baseURL); err != nil {
		return err
	}
	if len(claims.InstanceIDs) > 0 {
		instanceID = strings.TrimSpace(instanceID)
		if instanceID == "" {
			return errors.New("license requires an instance id")
		}
		allowed := false
		for _, candidate := range claims.InstanceIDs {
			if candidate == "*" || candidate == instanceID {
				allowed = true
				break
			}
		}
		if !allowed {
			return fmt.Errorf("instance %q is not covered by the license", instanceID)
		}
	}
	return nil
}

func ValidateLeaseContext(claims LeaseClaims, license Claims, product, baseURL, instanceID string) error {
	if claims.LicenseID != license.LicenseID {
		return errors.New("lease does not belong to the configured license")
	}
	if claims.Product != product || claims.Product != license.Product {
		return errors.New("lease product does not match")
	}
	host, err := HostFromBaseURL(baseURL)
	if err != nil {
		return err
	}
	if claims.Host != "" && !strings.EqualFold(claims.Host, host) {
		return errors.New("lease host does not match")
	}
	if claims.InstanceID != "" && claims.InstanceID != instanceID {
		return errors.New("lease instance does not match")
	}
	return nil
}

func ValidateDomain(domains []string, baseURL string) error {
	if len(domains) == 0 {
		return nil
	}
	host, err := HostFromBaseURL(baseURL)
	if err != nil {
		return err
	}
	for _, allowed := range domains {
		allowed = strings.ToLower(strings.TrimSpace(allowed))
		if allowed == "*" || host == allowed {
			return nil
		}
		if strings.HasPrefix(allowed, "*.") {
			root := strings.TrimPrefix(allowed, "*.")
			if host != root && strings.HasSuffix(host, "."+root) {
				return nil
			}
		}
	}
	return fmt.Errorf("host %q is not covered by the license", host)
}

func HostFromBaseURL(baseURL string) (string, error) {
	u, err := url.Parse(strings.TrimSpace(baseURL))
	if err != nil || u.Hostname() == "" {
		return "", errors.New("base URL has no valid host")
	}
	return strings.ToLower(u.Hostname()), nil
}

func StricterMode(a, b VerificationMode) VerificationMode {
	rank := map[VerificationMode]int{ModeOffline: 0, ModeHybrid: 1, ModeOnline: 2}
	if rank[b] > rank[a] {
		return b
	}
	if _, ok := rank[a]; !ok {
		return ModeOffline
	}
	return a
}

func ParseMode(value string) (VerificationMode, error) {
	mode := VerificationMode(strings.ToLower(strings.TrimSpace(value)))
	switch mode {
	case "", ModeOffline:
		return ModeOffline, nil
	case ModeHybrid, ModeOnline:
		return mode, nil
	default:
		return "", fmt.Errorf("unknown verification mode %q", value)
	}
}

func TokenHash(token string) string {
	sum := sha256.Sum256([]byte(strings.TrimSpace(token)))
	return hex.EncodeToString(sum[:])
}

func DecodePrivateKey(encoded string) (ed25519.PrivateKey, error) {
	b, err := decodeKey(encoded)
	if err != nil {
		return nil, err
	}
	if len(b) == ed25519.SeedSize {
		return ed25519.NewKeyFromSeed(b), nil
	}
	if len(b) != ed25519.PrivateKeySize {
		return nil, errors.New("private key must contain an Ed25519 seed or private key")
	}
	return ed25519.PrivateKey(b), nil
}

func DecodePublicKey(encoded string) (ed25519.PublicKey, error) {
	b, err := decodeKey(encoded)
	if err != nil {
		return nil, err
	}
	if len(b) != ed25519.PublicKeySize {
		return nil, errors.New("public key must contain an Ed25519 public key")
	}
	return ed25519.PublicKey(b), nil
}

func EncodeKey(key []byte) string { return base64.RawURLEncoding.EncodeToString(key) }

func decodeKey(value string) ([]byte, error) {
	value = strings.TrimSpace(value)
	if b, err := base64.RawURLEncoding.DecodeString(value); err == nil {
		return b, nil
	}
	if b, err := base64.StdEncoding.DecodeString(value); err == nil {
		return b, nil
	}
	return nil, errors.New("key is not valid base64")
}

func UniqueSorted(values []string) []string {
	seen := map[string]bool{}
	out := make([]string, 0, len(values))
	for _, value := range values {
		value = strings.TrimSpace(value)
		if value != "" && !seen[value] {
			seen[value] = true
			out = append(out, value)
		}
	}
	sort.Strings(out)
	return out
}

func validateLicenseClaims(c Claims, now time.Time, checkTime bool) error {
	if c.Version != SchemaVersion {
		return errors.New("unsupported license version")
	}
	if strings.TrimSpace(c.LicenseID) == "" || strings.TrimSpace(c.Issuer) == "" {
		return errors.New("license id and issuer are required")
	}
	if strings.TrimSpace(c.Customer) == "" || strings.TrimSpace(c.Product) == "" || strings.TrimSpace(c.Edition) == "" {
		return errors.New("customer, product and edition are required")
	}
	if c.IssuedAt <= 0 || c.ExpiresAt <= c.IssuedAt {
		return errors.New("license timestamps are invalid")
	}
	notBefore := c.NotBefore
	if notBefore == 0 {
		notBefore = c.IssuedAt
	}
	if checkTime {
		if now.Unix() < notBefore-300 {
			return errors.New("license is not active yet")
		}
		if now.Unix() >= c.ExpiresAt {
			return errors.New("license has expired")
		}
	}
	if _, err := ParseMode(string(c.Verification.Mode)); err != nil {
		return err
	}
	if c.Verification.LeaseTTLSeconds < 0 || c.Verification.OfflineGraceSeconds < 0 {
		return errors.New("verification durations cannot be negative")
	}
	if raw := strings.TrimSpace(c.Verification.ServerURL); raw != "" {
		parsed, err := url.Parse(raw)
		if err != nil || parsed.Host == "" || (parsed.Scheme != "https" && parsed.Scheme != "http") {
			return errors.New("verification server URL must be an absolute HTTP(S) URL")
		}
	}
	return nil
}

func validateLeaseClaims(c LeaseClaims, now time.Time, checkExpiration bool) error {
	if c.Version != SchemaVersion {
		return errors.New("unsupported lease version")
	}
	if strings.TrimSpace(c.LeaseID) == "" || strings.TrimSpace(c.LicenseID) == "" || strings.TrimSpace(c.Product) == "" {
		return errors.New("lease id, license id and product are required")
	}
	if c.IssuedAt <= 0 || c.ExpiresAt <= c.IssuedAt {
		return errors.New("lease timestamps are invalid")
	}
	if now.Unix() < c.IssuedAt-300 {
		return errors.New("lease is not active yet")
	}
	if checkExpiration && now.Unix() >= c.ExpiresAt {
		return errors.New("lease has expired")
	}
	return nil
}

func decodeStrict(data []byte, target any) error {
	dec := json.NewDecoder(strings.NewReader(string(data)))
	dec.DisallowUnknownFields()
	if err := dec.Decode(target); err != nil {
		return err
	}
	return nil
}