package dnssec

import (
	"context"
	"crypto"
	"crypto/ecdsa"
	"crypto/rsa"
	"encoding/json"
	"errors"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/coredns/coredns/request"

	"github.com/aws/aws-sdk-go-v2/config"
	"github.com/aws/aws-sdk-go-v2/service/secretsmanager"
	"github.com/miekg/dns"
	"golang.org/x/crypto/ed25519"
)

// DNSKEY holds a DNSSEC public and private key used for on-the-fly signing.
type DNSKEY struct {
	K   *dns.DNSKEY
	D   *dns.DS
	s   crypto.Signer
	tag uint16
}

// SecretKeyData represents the structure of the DNS keys stored in AWS Secrets Manager.
type SecretKeyData struct {
	Key     string `json:"key"`
	Private string `json:"private"`
}

// ParseKeyFile read a DNSSEC keyfile as generated by dnssec-keygen or other
// utilities. It adds ".key" for the public key and ".private" for the private key.
func ParseKeyFile(pubFile, privFile string) (*DNSKEY, error) {
	f, e := os.Open(filepath.Clean(pubFile))
	if e != nil {
		return nil, e
	}
	defer f.Close()
	k, e := dns.ReadRR(f, pubFile)
	if e != nil {
		return nil, e
	}

	f, e = os.Open(filepath.Clean(privFile))
	if e != nil {
		return nil, e
	}
	defer f.Close()

	dk, ok := k.(*dns.DNSKEY)
	if !ok {
		return nil, errors.New("no public key found")
	}
	p, e := dk.ReadPrivateKey(f, privFile)
	if e != nil {
		return nil, e
	}

	if s, ok := p.(*rsa.PrivateKey); ok {
		return &DNSKEY{K: dk, D: dk.ToDS(dns.SHA256), s: s, tag: dk.KeyTag()}, nil
	}
	if s, ok := p.(*ecdsa.PrivateKey); ok {
		return &DNSKEY{K: dk, D: dk.ToDS(dns.SHA256), s: s, tag: dk.KeyTag()}, nil
	}
	if s, ok := p.(ed25519.PrivateKey); ok {
		return &DNSKEY{K: dk, D: dk.ToDS(dns.SHA256), s: s, tag: dk.KeyTag()}, nil
	}
	return &DNSKEY{K: dk, D: dk.ToDS(dns.SHA256), s: nil, tag: 0}, errors.New("no private key found")
}

// ParseKeyFromAWSSecretsManager retrieves and parses a DNSSEC key pair from AWS Secrets Manager.
func ParseKeyFromAWSSecretsManager(secretID string) (*DNSKEY, error) {
	// Load the AWS SDK configuration
	cfg, err := config.LoadDefaultConfig(context.TODO())
	if err != nil {
		return nil, err
	}

	// Create a Secrets Manager client
	client := secretsmanager.NewFromConfig(cfg)

	// Retrieve the secret value
	input := &secretsmanager.GetSecretValueInput{
		SecretId: &secretID,
	}
	result, err := client.GetSecretValue(context.TODO(), input)
	if err != nil {
		return nil, err
	}

	// Parse the secret string into SecretKeyData
	var secretData SecretKeyData
	err = json.Unmarshal([]byte(*result.SecretString), &secretData)
	if err != nil {
		return nil, err
	}

	// Parse the public key
	rr, err := dns.NewRR(secretData.Key)
	if err != nil {
		return nil, err
	}
	dk, ok := rr.(*dns.DNSKEY)
	if !ok {
		return nil, errors.New("invalid public key format")
	}

	// Parse the private key
	p, err := dk.ReadPrivateKey(strings.NewReader(secretData.Private), secretID)
	if err != nil {
		return nil, err
	}

	// Create the DNSKEY structure
	var s crypto.Signer
	var tag uint16
	switch key := p.(type) {
	case *rsa.PrivateKey:
		s = key
		tag = dk.KeyTag()
	case *ecdsa.PrivateKey:
		s = key
		tag = dk.KeyTag()
	case ed25519.PrivateKey:
		s = key
		tag = dk.KeyTag()
	default:
		return nil, errors.New("unsupported key type")
	}

	return &DNSKEY{K: dk, D: dk.ToDS(dns.SHA256), s: s, tag: tag}, nil
}

// getDNSKEY returns the correct DNSKEY to the client. Signatures are added when do is true.
func (d Dnssec) getDNSKEY(state request.Request, zone string, do bool, server string) *dns.Msg {
	keys := make([]dns.RR, len(d.keys))
	for i, k := range d.keys {
		keys[i] = dns.Copy(k.K)
		keys[i].Header().Name = zone
	}
	m := new(dns.Msg)
	m.SetReply(state.Req)
	m.Answer = keys
	if !do {
		return m
	}

	incep, expir := incepExpir(time.Now().UTC())
	if sigs, err := d.sign(keys, zone, 3600, incep, expir, server); err == nil {
		m.Answer = append(m.Answer, sigs...)
	}
	return m
}

// Return true if, and only if, this is a zone key with the SEP bit unset. This implies a ZSK (rfc4034 2.1.1).
func (k DNSKEY) isZSK() bool {
	return k.K.Flags&(1<<8) == (1<<8) && k.K.Flags&1 == 0
}

// Return true if, and only if, this is a zone key with the SEP bit set. This implies a KSK (rfc4034 2.1.1).
func (k DNSKEY) isKSK() bool {
	return k.K.Flags&(1<<8) == (1<<8) && k.K.Flags&1 == 1
}