package platform

import (
	"context"
	"crypto"
	"crypto/x509"
	"errors"
	"fmt"
	"net/url"
	"strings"

	"go.step.sm/crypto/kms/apiv1"
	"go.step.sm/crypto/kms/uri"
)

const Scheme = "kms"

func init() {
	apiv1.Register(apiv1.PlatformKMS, func(ctx context.Context, opts apiv1.Options) (apiv1.KeyManager, error) {
		return New(ctx, opts)
	})
}

const (
	backendKey = "backend"
	nameKey    = "name"
	hwKey      = "hw"
)

type kmsURI struct {
	uri         *uri.URI
	name        string
	hw          bool
	extraValues url.Values
}

func isDefaultKey(k string) bool {
	return k == nameKey || k == hwKey || k == backendKey
}

func getBackend(opts apiv1.Options) (apiv1.Type, error) {
	if opts.URI == "" {
		return opts.Type, nil
	}

	typ := opts.Type
	if opts.URI != "" {
		u, err := uri.ParseWithScheme(Scheme, opts.URI)
		if err != nil {
			return apiv1.DefaultKMS, err
		}
		if backend := u.Get(backendKey); backend != "" {
			typ = apiv1.Type(strings.ToLower(backend))
			if opts.Type != apiv1.DefaultKMS && opts.Type != typ {
				return apiv1.DefaultKMS, fmt.Errorf("options type %q and URI backend %q do not match", opts, typ)
			}
		}
	}

	return typ, nil
}

func parseURI(rawuri string) (*kmsURI, error) {
	u, err := uri.ParseWithScheme(Scheme, rawuri)
	if err != nil {
		return nil, err
	}

	extraValues := make(url.Values)
	for k, v := range uri.Values(u) {
		if !isDefaultKey(k) {
			extraValues[k] = v
		}
	}

	return &kmsURI{
		uri:         u,
		name:        u.Get(nameKey),
		hw:          u.GetBool(hwKey),
		extraValues: extraValues,
	}, nil
}

type extendedKeyManager interface {
	apiv1.KeyManager
	apiv1.KeyDeleter
	apiv1.CertificateManager
	apiv1.CertificateChainManager
	apiv1.CertificateDeleter
}

var (
	_ apiv1.KeyManager                   = (*KMS)(nil)
	_ apiv1.CertificateManager           = (*KMS)(nil)
	_ apiv1.CertificateChainManager      = (*KMS)(nil)
	_ apiv1.SearchableKeyManager         = (*KMS)(nil)
	_ apiv1.CredentialsCleaner           = (*KMS)(nil)
	_ apiv1.SearchableCertificateManager = (*KMS)(nil)
)

type KMS struct {
	typ              apiv1.Type
	backend          extendedKeyManager
	transformToURI   func(string) (string, error)
	transformFromURI func(string) (string, error)
}

func New(ctx context.Context, opts apiv1.Options) (*KMS, error) {
	return newKMS(ctx, opts)
}

func (k *KMS) Type() apiv1.Type {
	return k.typ
}

func (k *KMS) Close() error {
	return k.backend.Close()
}

func (k *KMS) GetPublicKey(req *apiv1.GetPublicKeyRequest) (crypto.PublicKey, error) {
	name, err := k.transformToURI(req.Name)
	if err != nil {
		return nil, err
	}

	return k.backend.GetPublicKey(&apiv1.GetPublicKeyRequest{
		Name: name,
	})
}

func (k *KMS) CreateKey(req *apiv1.CreateKeyRequest) (*apiv1.CreateKeyResponse, error) {
	name, err := k.transformToURI(req.Name)
	if err != nil {
		return nil, err
	}

	r := clone(req)
	r.Name = name
	resp, err := k.backend.CreateKey(r)
	if err != nil {
		return nil, err
	}

	return k.patchCreateKeyResponse(resp)
}

func (k *KMS) CreateSigner(req *apiv1.CreateSignerRequest) (crypto.Signer, error) {
	if req.Signer != nil {
		return req.Signer, nil
	}

	signingKey, err := k.transformToURI(req.SigningKey)
	if err != nil {
		return nil, err
	}

	r := clone(req)
	r.SigningKey = signingKey
	return k.backend.CreateSigner(r)
}

func (k *KMS) DeleteKey(req *apiv1.DeleteKeyRequest) error {
	name, err := k.transformToURI(req.Name)
	if err != nil {
		return err
	}

	r := clone(req)
	r.Name = name
	return k.backend.DeleteKey(r)
}

func (k *KMS) LoadCertificate(req *apiv1.LoadCertificateRequest) (*x509.Certificate, error) {
	name, err := k.transformToURI(req.Name)
	if err != nil {
		return nil, err
	}

	r := clone(req)
	r.Name = name
	return k.backend.LoadCertificate(r)
}

func (k *KMS) StoreCertificate(req *apiv1.StoreCertificateRequest) error {
	name, err := k.transformToURI(req.Name)
	if err != nil {
		return err
	}

	r := clone(req)
	r.Name = name
	return k.backend.StoreCertificate(r)
}

func (k *KMS) LoadCertificateChain(req *apiv1.LoadCertificateChainRequest) ([]*x509.Certificate, error) {
	name, err := k.transformToURI(req.Name)
	if err != nil {
		return nil, err
	}

	r := clone(req)
	r.Name = name
	return k.backend.LoadCertificateChain(r)
}

func (k *KMS) StoreCertificateChain(req *apiv1.StoreCertificateChainRequest) error {
	name, err := k.transformToURI(req.Name)
	if err != nil {
		return err
	}

	r := clone(req)
	r.Name = name
	return k.backend.StoreCertificateChain(r)
}

func (k *KMS) DeleteCertificate(req *apiv1.DeleteCertificateRequest) error {
	name, err := k.transformToURI(req.Name)
	if err != nil {
		return err
	}

	r := clone(req)
	r.Name = name
	return k.backend.DeleteCertificate(r)
}

func (k *KMS) CreateAttestation(req *apiv1.CreateAttestationRequest) (*apiv1.CreateAttestationResponse, error) {
	if req.Name == "" {
		return nil, errors.New("createAttestationRequest 'name' cannot be empty")
	}

	name, err := k.transformToURI(req.Name)
	if err != nil {
		return nil, err
	}

	// Attestation implemented by the backend
	if km, ok := k.backend.(apiv1.Attester); ok {
		r := clone(req)
		r.Name = name

		return km.CreateAttestation(r)
	}

	// Attestation using a custom attestation client.
	if req.AttestationClient != nil {
		signer, err := k.backend.CreateSigner(&apiv1.CreateSignerRequest{
			SigningKey: name,
		})
		if err != nil {
			return nil, err
		}

		ctx := apiv1.NewAttestSignerContext(context.Background(), signer)
		chain, err := req.AttestationClient.Attest(ctx)
		if err != nil {
			return nil, err
		}

		var permanentIdentifier string
		if len(chain[0].URIs) > 0 {
			permanentIdentifier = chain[0].URIs[0].String()
		}

		return &apiv1.CreateAttestationResponse{
			Certificate:         chain[0],
			CertificateChain:    chain,
			PublicKey:           signer.Public(),
			PermanentIdentifier: permanentIdentifier,
		}, nil
	}

	return nil, apiv1.NotImplementedError{}
}

func (k *KMS) SearchKeys(req *apiv1.SearchKeysRequest) (*apiv1.SearchKeysResponse, error) {
	if km, ok := k.backend.(apiv1.SearchableKeyManager); ok {
		query, err := k.transformToURI(req.Query)
		if err != nil {
			return nil, err
		}

		r := clone(req)
		r.Query = query
		resp, err := km.SearchKeys(r)

		// SearchKeys is best-effort: the backend may return usable results
		// alongside an error describing the keys it couldn't load. For the
		// tpmkms backend this happens when a key's signer can't be loaded — a
		// corrupt or unreadable TPM or PCP/CNG keyset (e.g. NTE_BAD_KEYSET on
		// Windows, or a CNG container removed out of band) — or when the TPM
		// can't be enumerated at all. Forward whatever results came back (after
		// patching their URIs) together with that error, so callers can use the
		// subset and still see what failed. Only a response we can't patch is
		// fatal.
		if resp == nil {
			return nil, err
		}

		patched, patchErr := k.patchSearchKeysResponse(resp)
		if patchErr != nil {
			return nil, patchErr
		}

		return patched, err
	}

	return nil, apiv1.NotImplementedError{}
}

func (k *KMS) CleanupCredentials(req *apiv1.CleanupCredentialsRequest) error {
	km, ok := k.backend.(apiv1.CredentialsCleaner)
	if !ok {
		return apiv1.NotImplementedError{}
	}

	name, err := k.transformToURI(req.Name)
	if err != nil {
		return err
	}

	r := clone(req)
	r.Name = name
	return km.CleanupCredentials(r)
}

// SearchCertificates delegates to the backend's SearchableCertificateManager
// implementation, if it has one. req.Name is translated from a "kms:" URI to
// the backend's URI scheme before delegating, the same way CleanupCredentials
// translates its request name. The response carries only certificates and a
// bare key name, neither of which encodes a KMS URI, so no response
// translation is needed. Per-certificate metadata failures reported on a
// result's Err field are passed through untouched as well.
//
// SearchCertificates is best-effort: on a mid-enumeration failure the backend
// may return a non-nil response holding the certificates enumerated before
// the failure, together with the error describing it. That response is
// returned unmodified alongside the error, so callers wanting the partial
// results must check the response before the error.
func (k *KMS) SearchCertificates(req *apiv1.SearchCertificatesRequest) (*apiv1.SearchCertificatesResponse, error) {
	if req == nil {
		return nil, errors.New("searchCertificatesRequest cannot be nil")
	}

	km, ok := k.backend.(apiv1.SearchableCertificateManager)
	if !ok {
		return nil, apiv1.NotImplementedError{}
	}

	name, err := k.transformToURI(req.Name)
	if err != nil {
		return nil, err
	}

	r := clone(req)
	r.Name = name
	return km.SearchCertificates(r)
}

func (k *KMS) patchCreateKeyResponse(resp *apiv1.CreateKeyResponse) (*apiv1.CreateKeyResponse, error) {
	name, err := k.transformFromURI(resp.Name)
	if err != nil {
		return nil, err
	}

	resp.Name = name
	if resp.CreateSignerRequest.SigningKey != "" {
		resp.CreateSignerRequest.SigningKey = name
	}

	return resp, nil
}

func (k *KMS) patchSearchKeysResponse(resp *apiv1.SearchKeysResponse) (*apiv1.SearchKeysResponse, error) {
	for i := range resp.Results {
		name, err := k.transformFromURI(resp.Results[i].Name)
		if err != nil {
			return nil, err
		}

		resp.Results[i].Name = name
		if resp.Results[i].CreateSignerRequest.SigningKey != "" {
			resp.Results[i].CreateSignerRequest.SigningKey = name
		}
	}

	return resp, nil
}

func clone[T any](v *T) *T {
	c := *v
	return &c
}