package signing

import (
	"context"
	"crypto/rsa"
	"crypto/x509"
	"encoding/base64"
	"encoding/pem"
	"errors"
	"fmt"
	"sort"
	"strconv"
	"strings"

	"github.com/go-jose/go-jose/v4"
	"github.com/go-viper/mapstructure/v2"
	baoapi "github.com/openbao/openbao/api/v2"
)

type TransitSigner struct {
	client *baoapi.Client
	mount  string
	key    string
}

type transitKeyData struct {
	LatestVersion int                          `mapstructure:"latest_version"`
	Keys          map[string]transitKeyVersion `mapstructure:"keys"`
}

type transitKeyVersion struct {
	PublicKey string `mapstructure:"public_key"`
}

func NewTransitSigner(client *baoapi.Client, mount, key string) (*TransitSigner, error) {
	if client == nil {
		return nil, errors.New("OpenBao client is required")
	}
	mount, key = strings.Trim(mount, "/"), strings.TrimSpace(key)
	if mount == "" || key == "" || strings.Contains(key, "/") {
		return nil, errors.New("valid Transit mount and key are required")
	}
	return &TransitSigner{client: client, mount: mount, key: key}, nil
}

func (s *TransitSigner) ActiveKey(ctx context.Context) (SigningKey, error) {
	data, err := s.readKeyData(ctx)
	if err != nil {
		return SigningKey{}, err
	}
	if data.LatestVersion < 1 {
		return SigningKey{}, errors.New("read Transit key: invalid latest_version")
	}
	return SigningKey{ID: s.keyID(data.LatestVersion), Version: data.LatestVersion}, nil
}

func (s *TransitSigner) SignRS256(ctx context.Context, key SigningKey, input []byte) ([]byte, error) {
	if key.Version < 1 || key.ID != s.keyID(key.Version) {
		return nil, errors.New("sign with Transit: invalid signing key")
	}
	secret, err := s.client.Logical().WriteWithContext(ctx, s.signPath(), map[string]any{
		"input":               base64.StdEncoding.EncodeToString(input),
		"key_version":         key.Version,
		"signature_algorithm": "pkcs1v15",
	})
	if err != nil {
		return nil, fmt.Errorf("sign with Transit: %w", err)
	}
	if secret == nil || secret.Data == nil {
		return nil, errors.New("sign with Transit: empty response")
	}
	encoded, ok := secret.Data["signature"].(string)
	if !ok {
		return nil, errors.New("sign with Transit: missing signature")
	}
	return decodeTransitSignature(encoded, key.Version)
}

func (s *TransitSigner) JWKS(ctx context.Context) (jose.JSONWebKeySet, error) {
	data, err := s.readKeyData(ctx)
	if err != nil {
		return jose.JSONWebKeySet{}, err
	}
	versions := make([]int, 0, len(data.Keys))
	for text := range data.Keys {
		version, err := strconv.Atoi(text)
		if err != nil || version < 1 {
			return jose.JSONWebKeySet{}, errors.New("read Transit keys: invalid key version")
		}
		versions = append(versions, version)
	}
	sort.Ints(versions)
	set := jose.JSONWebKeySet{Keys: make([]jose.JSONWebKey, 0, len(versions))}
	for _, version := range versions {
		publicKey, err := parseRSAPublicKey(data.Keys[strconv.Itoa(version)].PublicKey)
		if err != nil {
			return jose.JSONWebKeySet{}, fmt.Errorf("read Transit keys version %d: %w", version, err)
		}
		set.Keys = append(set.Keys, jose.JSONWebKey{
			Key:       publicKey,
			KeyID:     s.keyID(version),
			Algorithm: string(jose.RS256),
			Use:       "sig",
		})
	}
	return set, nil
}

func (s *TransitSigner) readKeyData(ctx context.Context) (transitKeyData, error) {
	secret, err := s.client.Logical().ReadWithContext(ctx, s.keyPath())
	if err != nil {
		return transitKeyData{}, fmt.Errorf("read Transit key: %w", err)
	}
	if secret == nil || secret.Data == nil {
		return transitKeyData{}, errors.New("read Transit key: empty response")
	}
	var data transitKeyData
	if err := mapstructure.Decode(secret.Data, &data); err != nil {
		return transitKeyData{}, fmt.Errorf("decode Transit key metadata: %w", err)
	}
	return data, nil
}

func decodeTransitSignature(encoded string, version int) ([]byte, error) {
	parts := strings.SplitN(encoded, ":", 3)
	if len(parts) != 3 || parts[0] != "vault" || parts[1] != "v"+strconv.Itoa(version) {
		return nil, errors.New("sign with Transit: unexpected signature version")
	}
	signature, err := base64.StdEncoding.DecodeString(parts[2])
	if err != nil {
		return nil, errors.New("sign with Transit: invalid signature encoding")
	}
	return signature, nil
}

func (s *TransitSigner) keyPath() string          { return s.mount + "/keys/" + s.key }
func (s *TransitSigner) signPath() string         { return s.mount + "/sign/" + s.key + "/sha2-256" }
func (s *TransitSigner) keyID(version int) string { return s.key + "-v" + strconv.Itoa(version) }

func parseRSAPublicKey(value string) (*rsa.PublicKey, error) {
	block, _ := pem.Decode([]byte(value))
	if block == nil {
		return nil, errors.New("invalid PEM public key")
	}
	parsed, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err != nil {
		return nil, fmt.Errorf("parse public key: %w", err)
	}
	publicKey, ok := parsed.(*rsa.PublicKey)
	if !ok {
		return nil, errors.New("Transit key is not RSA")
	}
	return publicKey, nil
}