package jws

import (
	"bytes"
	"fmt"

	"github.com/lestrrat-go/jwx/v3/internal/base64"
	"github.com/lestrrat-go/jwx/v3/internal/json"
	"github.com/lestrrat-go/jwx/v3/internal/pool"
	"github.com/lestrrat-go/jwx/v3/internal/tokens"
	"github.com/lestrrat-go/jwx/v3/jwa"
	"github.com/lestrrat-go/jwx/v3/jwk"
)

func NewSignature() *Signature {
	return &Signature{}
}

func (s *Signature) DecodeCtx() DecodeCtx {
	return s.dc
}

func (s *Signature) SetDecodeCtx(dc DecodeCtx) {
	s.dc = dc
}

func (s Signature) PublicHeaders() Headers {
	return s.headers
}

func (s *Signature) SetPublicHeaders(v Headers) *Signature {
	s.headers = v
	return s
}

func (s Signature) ProtectedHeaders() Headers {
	return s.protected
}

func (s *Signature) SetProtectedHeaders(v Headers) *Signature {
	s.protected = v
	return s
}

func (s Signature) Signature() []byte {
	return s.signature
}

func (s *Signature) SetSignature(v []byte) *Signature {
	s.signature = v
	return s
}

type signatureUnmarshalProbe struct {
	Header    Headers `json:"header,omitempty"`
	Protected *string `json:"protected,omitempty"`
	Signature *string `json:"signature,omitempty"`
}

func (s *Signature) UnmarshalJSON(data []byte) error {
	var sup signatureUnmarshalProbe
	sup.Header = NewHeaders()
	if err := json.Unmarshal(data, &sup); err != nil {
		return fmt.Errorf(`failed to unmarshal signature into temporary struct: %w`, err)
	}

	s.headers = sup.Header
	if buf := sup.Protected; buf != nil {
		// RFC 7515 §3 mandates that "protected" be base64url-encoded.
		// Earlier code carried a relaxed probe that accepted a literal-
		// JSON form (a JSON string whose content begins with "{") and
		// skipped base64 decoding — that was asymmetric with the
		// flattened branch (which only base64-decodes) and gave callers
		// a non-conforming wire form useful for evading byte-exact JWS
		// dedup / replay caches.
		decoded, err := base64.Decode([]byte(*buf))
		if err != nil {
			return fmt.Errorf(`failed to base64 decode protected headers: %w`, err)
		}
		src := decoded

		prt := NewHeaders()
		//nolint:forcetypeassert
		prt.(*stdHeaders).SetDecodeCtx(s.DecodeCtx())
		if err := json.Unmarshal(src, prt); err != nil {
			return fmt.Errorf(`failed to unmarshal protected headers: %w`, err)
		}
		//nolint:forcetypeassert
		prt.(*stdHeaders).SetDecodeCtx(nil)
		s.protected = prt
	}

	if sup.Signature != nil {
		decoded, err := base64.DecodeString(*sup.Signature)
		if err != nil {
			return fmt.Errorf(`failed to base decode signature: %w`, err)
		}
		s.signature = decoded
	}
	return nil
}

// Sign populates the signature field, with a signature generated by
// the given signer object and payload.
//
// The first return value is the raw signature in binary format.
// The second return value is the full three-segment signature
// (e.g. "eyXXXX.XXXXX.XXXX")
//
// Deprecated: Signature objects will only be used as containers in the future.
// Use [Sign] or the jwsbb package to craft signatures manually.
func (s *Signature) Sign(payload []byte, signer Signer, key any) ([]byte, []byte, error) {
	return s.sign2(payload, signer, key)
}

func (s *Signature) sign2(payload []byte, signer interface{ Algorithm() jwa.SignatureAlgorithm }, key any) ([]byte, []byte, error) {
	// A jwk.UnsupportedKey placeholder carries no usable key material and
	// must never reach the Signer. This path builds a signatureBuilder
	// directly and bypasses validateAlgorithmForKey (only the
	// jws.Sign/WithKey path runs that guard), so reject the placeholder
	// here — the single entry to signatureBuilder for this path.
	if uk, ok := key.(jwk.UnsupportedKey); ok {
		return nil, nil, unsupportedKeyError(uk, `signing`)
	}

	// Create a signatureBuilder to use the shared signing logic
	sb := signatureBuilderPool.Get()
	defer signatureBuilderPool.Put(sb)

	sb.alg = signer.Algorithm()
	sb.key = key
	sb.protected = s.protected
	sb.public = s.headers

	// Set up the appropriate signer interface
	switch typedSigner := signer.(type) {
	case Signer2:
		sb.signer2 = typedSigner
	case Signer:
		sb.signer = typedSigner
	default:
		return nil, nil, fmt.Errorf(`invalid signer type: %T`, signer)
	}

	// Create a minimal sign context
	sc := signContextPool.Get()
	defer signContextPool.Put(sc)

	sc.detached = s.detached

	encoder := s.encoder
	if encoder == nil {
		encoder = base64.DefaultEncoder()
	}
	sc.encoder = encoder

	// Build the signature using signatureBuilder
	sig, err := sb.Build(sc, payload)
	if err != nil {
		return nil, nil, fmt.Errorf(`failed to build signature: %w`, err)
	}

	// Copy the signature result back to this signature instance
	s.signature = sig.signature
	s.protected = sig.protected
	s.headers = sig.headers

	// Build the complete JWS token for the return value
	buf := pool.BytesBuffer().Get()
	defer pool.BytesBuffer().Put(buf)

	// Marshal the merged headers for the final output
	hdrs, err := mergeHeaders(s.headers, s.protected)
	if err != nil {
		return nil, nil, fmt.Errorf(`failed to merge headers: %w`, err)
	}

	hdrbuf, err := json.Marshal(hdrs)
	if err != nil {
		return nil, nil, fmt.Errorf(`failed to marshal headers: %w`, err)
	}

	buf.WriteString(encoder.EncodeToString(hdrbuf))
	buf.WriteByte(tokens.Period)

	var plen int
	b64 := getB64Value(hdrs)
	if b64 {
		encoded := encoder.EncodeToString(payload)
		plen = len(encoded)
		buf.WriteString(encoded)
	} else {
		if !s.detached {
			if bytes.Contains(payload, []byte{tokens.Period}) {
				return nil, nil, fmt.Errorf(`compact serialization with b64=false requires payload to contain no "." characters per RFC 7797 §5.2; use jws.WithDetachedPayload to keep the payload out of the wire format`)
			}
		}
		plen = len(payload)
		buf.Write(payload)
	}

	// Handle detached payload
	if s.detached {
		buf.Truncate(buf.Len() - plen)
	}

	buf.WriteByte(tokens.Period)
	buf.WriteString(encoder.EncodeToString(s.signature))
	ret := bytes.Clone(buf.Bytes())

	return s.signature, ret, nil
}

func NewMessage() *Message {
	return &Message{}
}

// Clears the internal raw buffer that was accumulated during
// the verify phase
func (m *Message) clearRaw() {
	for _, sig := range m.signatures {
		if protected := sig.protected; protected != nil {
			if cr, ok := protected.(*stdHeaders); ok {
				cr.raw = nil
			}
		}
	}
}

func (m *Message) SetDecodeCtx(dc DecodeCtx) {
	m.dc = dc
}

func (m *Message) DecodeCtx() DecodeCtx {
	return m.dc
}

// Payload returns the decoded payload
func (m Message) Payload() []byte {
	return m.payload
}

func (m *Message) SetPayload(v []byte) *Message {
	m.payload = v
	return m
}

func (m Message) Signatures() []*Signature {
	return m.signatures
}

func (m *Message) AppendSignature(v *Signature) *Message {
	m.signatures = append(m.signatures, v)
	return m
}

func (m *Message) ClearSignatures() *Message {
	m.signatures = nil
	return m
}

// LookupSignature looks up a particular signature entry using
// the `kid` value
func (m Message) LookupSignature(kid string) []*Signature {
	var sigs []*Signature
	for _, sig := range m.signatures {
		if hdr := sig.PublicHeaders(); hdr != nil {
			hdrKeyID, ok := hdr.KeyID()
			if ok && hdrKeyID == kid {
				sigs = append(sigs, sig)
				continue
			}
		}

		if hdr := sig.ProtectedHeaders(); hdr != nil {
			hdrKeyID, ok := hdr.KeyID()
			if ok && hdrKeyID == kid {
				sigs = append(sigs, sig)
				continue
			}
		}
	}
	return sigs
}

// This struct is used to first probe for the structure of the
// incoming JSON object. We then decide how to parse it
// from the fields that are populated.
type messageUnmarshalProbe struct {
	Payload    *string           `json:"payload"`
	Signatures []json.RawMessage `json:"signatures,omitempty"`
	Header     json.RawMessage   `json:"header,omitempty"`
	Protected  *string           `json:"protected,omitempty"`
	Signature  *string           `json:"signature,omitempty"`
}

func (m *Message) UnmarshalJSON(buf []byte) error {
	m.payload = nil
	m.signatures = nil
	m.b64 = true
	m.detached = false

	var mup messageUnmarshalProbe
	if err := json.Unmarshal(buf, &mup); err != nil {
		return fmt.Errorf(`failed to unmarshal into temporary structure: %w`, err)
	}

	// Enforce the signature cap before we decode any signature entry.
	// The probe above leaves mup.Signatures as []json.RawMessage, so no
	// headers/base64 work has happened yet. Doing the check here prevents
	// a large signatures array from allocating O(input) work before Parse
	// gets a chance to reject it.
	maxSigs := m.maxSignatures
	if maxSigs == 0 {
		maxSigs = int(maxSignatures.Load())
	}
	if maxSigs > 0 && len(mup.Signatures) > maxSigs {
		return fmt.Errorf(`too many signatures in JWS message (%d > %d)`, len(mup.Signatures), maxSigs)
	}

	b64 := true
	if mup.Signature == nil { // flattened signature is NOT present
		if len(mup.Signatures) == 0 {
			return fmt.Errorf(`required field "signatures" not present`)
		}

		// RFC 7515 §7.2.1 places the unprotected JOSE header inside each
		// signature entry for the general (multi-signature) form; a
		// top-level "header" sibling of "signatures" is not defined.
		// Reject rather than silently drop — silent drop hid both typos
		// and an attacker-controlled trigger surface for any
		// RegisterCustomDecoder side effects on the dropped contents.
		if len(mup.Header) > 0 && !bytes.Equal(mup.Header, []byte("null")) {
			return fmt.Errorf(`general-form JWS must not contain top-level "header" sibling of "signatures" (RFC 7515 §7.2.1 places the unprotected header inside each signature entry)`)
		}

		m.signatures = make([]*Signature, 0, len(mup.Signatures))
		for i, rawsig := range mup.Signatures {
			var sig Signature
			sig.SetDecodeCtx(m.DecodeCtx())
			if err := json.Unmarshal(rawsig, &sig); err != nil {
				return fmt.Errorf(`failed to unmarshal signature #%d: %w`, i+1, err)
			}
			sig.SetDecodeCtx(nil)

			if sig.protected == nil {
				// Instead of barfing on a nil protected header, use an empty header
				sig.protected = NewHeaders()
			}

			if i == 0 {
				if !getB64Value(sig.protected) {
					b64 = false
				}
			} else {
				if got := getB64Value(sig.protected); b64 != got {
					return fmt.Errorf(`b64 value must be the same for all signatures; signature #%d declared b64=%t but earlier signature(s) declared b64=%t`, i+1, got, b64)
				}
			}

			m.signatures = append(m.signatures, &sig)
		}
	} else { // .signature is present, it's a flattened structure
		if len(mup.Signatures) != 0 {
			return fmt.Errorf(`invalid format ("signatures" and "signature" keys cannot both be present)`)
		}

		var sig Signature
		if len(mup.Header) > 0 && !bytes.Equal(mup.Header, []byte("null")) {
			hdrs := NewHeaders()
			if err := json.Unmarshal(mup.Header, hdrs); err != nil {
				return fmt.Errorf(`failed to unmarshal flattened unprotected header: %w`, err)
			}
			sig.headers = hdrs
		}
		if src := mup.Protected; src != nil {
			decoded, err := base64.DecodeString(*src)
			if err != nil {
				return fmt.Errorf(`failed to base64 decode flattened protected headers: %w`, err)
			}
			prt := NewHeaders()
			//nolint:forcetypeassert
			prt.(*stdHeaders).SetDecodeCtx(m.DecodeCtx())
			if err := json.Unmarshal(decoded, prt); err != nil {
				return fmt.Errorf(`failed to unmarshal flattened protected headers: %w`, err)
			}
			//nolint:forcetypeassert
			prt.(*stdHeaders).SetDecodeCtx(nil)
			sig.protected = prt
		}

		if sig.protected == nil {
			// Instead of barfing on a nil protected header, use an empty header
			sig.protected = NewHeaders()
		}

		decoded, err := base64.DecodeString(*mup.Signature)
		if err != nil {
			return fmt.Errorf(`failed to base64 decode flattened signature: %w`, err)
		}
		sig.signature = decoded

		m.signatures = []*Signature{&sig}
		b64 = getB64Value(sig.protected)
	}

	if mup.Payload == nil {
		// RFC 7515 Appendix F: detached content is signaled in the
		// JSON Serialization by omitting the "payload" member.
		m.detached = true
	} else {
		if !b64 { // NOT base64 encoded
			m.payload = []byte(*mup.Payload)
		} else {
			decoded, err := base64.DecodeString(*mup.Payload)
			if err != nil {
				return fmt.Errorf(`failed to base64 decode payload: %w`, err)
			}
			m.payload = decoded
		}
	}
	m.b64 = b64
	return nil
}

func (m Message) MarshalJSON() ([]byte, error) {
	if len(m.signatures) == 1 {
		return m.marshalFlattened()
	}
	return m.marshalFull()
}

func (m Message) marshalFlattened() ([]byte, error) {
	buf := pool.BytesBuffer().Get()
	defer pool.BytesBuffer().Put(buf)

	sig := m.signatures[0]

	buf.WriteRune(tokens.OpenCurlyBracket)
	var wrote bool

	if hdr := sig.headers; hdr != nil {
		hdrjs, err := json.Marshal(hdr)
		if err != nil {
			return nil, fmt.Errorf(`failed to marshal "header" (flattened format): %w`, err)
		}
		buf.WriteString(`"header":`)
		buf.Write(hdrjs)
		wrote = true
	}

	// RFC 7515 Appendix F: detached content is signaled in the JSON
	// Serialization by omitting the "payload" member.
	if !m.detached {
		if wrote {
			buf.WriteRune(tokens.Comma)
		}
		if !getB64Value(sig.protected) {
			// RFC 7797 b64=false: emit the raw payload as a JSON string
			// rather than re-base64-encoding it. json.Marshal handles
			// any necessary escaping for byte sequences that aren't
			// JSON-safe as-is.
			payloadbuf, err := json.Marshal(string(m.payload))
			if err != nil {
				return nil, fmt.Errorf(`failed to marshal "payload" (flattened, b64=false): %w`, err)
			}
			buf.WriteString(`"payload":`)
			buf.Write(payloadbuf)
		} else {
			buf.WriteString(`"payload":"`)
			buf.Write(base64.Encode(m.payload))
			buf.WriteRune('"')
		}
		wrote = true
	}

	if protected := sig.protected; protected != nil {
		protectedbuf, err := json.Marshal(protected)
		if err != nil {
			return nil, fmt.Errorf(`failed to marshal "protected" (flattened format): %w`, err)
		}
		if wrote {
			buf.WriteRune(tokens.Comma)
		}
		buf.WriteString(`"protected":"`)
		buf.Write(base64.Encode(protectedbuf))
		buf.WriteRune('"')
		wrote = true
	}

	if wrote {
		buf.WriteRune(tokens.Comma)
	}
	buf.WriteString(`"signature":"`)
	buf.Write(base64.Encode(sig.signature))
	buf.WriteRune('"')
	buf.WriteRune(tokens.CloseCurlyBracket)

	ret := bytes.Clone(buf.Bytes())
	return ret, nil
}

func (m Message) marshalFull() ([]byte, error) {
	buf := pool.BytesBuffer().Get()
	defer pool.BytesBuffer().Put(buf)

	// RFC 7515 Appendix F: detached content is signaled in the JSON
	// Serialization by omitting the "payload" member.
	if m.detached {
		buf.WriteString(`{"signatures":[`)
	} else {
		// RFC 7797 b64=false: emit the raw payload as a JSON string
		// rather than re-base64-encoding it. The general JWS form has
		// one shared payload across signatures; per RFC 7797, all
		// signers must agree on the b64 flag, so we consult the first
		// signature's protected header.
		var b64 = true
		if len(m.signatures) > 0 {
			b64 = getB64Value(m.signatures[0].protected)
		}
		if !b64 {
			payloadbuf, err := json.Marshal(string(m.payload))
			if err != nil {
				return nil, fmt.Errorf(`failed to marshal "payload" (full, b64=false): %w`, err)
			}
			buf.WriteString(`{"payload":`)
			buf.Write(payloadbuf)
			buf.WriteString(`,"signatures":[`)
		} else {
			buf.WriteString(`{"payload":"`)
			buf.Write(base64.Encode(m.payload))
			buf.WriteString(`","signatures":[`)
		}
	}
	for i, sig := range m.signatures {
		if i > 0 {
			buf.WriteRune(tokens.Comma)
		}

		buf.WriteRune(tokens.OpenCurlyBracket)
		var wrote bool
		if hdr := sig.headers; hdr != nil {
			hdrbuf, err := json.Marshal(hdr)
			if err != nil {
				return nil, fmt.Errorf(`failed to marshal "header" for signature #%d: %w`, i+1, err)
			}
			buf.WriteString(`"header":`)
			buf.Write(hdrbuf)
			wrote = true
		}

		if protected := sig.protected; protected != nil {
			protectedbuf, err := json.Marshal(protected)
			if err != nil {
				return nil, fmt.Errorf(`failed to marshal "protected" for signature #%d: %w`, i+1, err)
			}
			if wrote {
				buf.WriteRune(tokens.Comma)
			}
			buf.WriteString(`"protected":"`)
			buf.Write(base64.Encode(protectedbuf))
			buf.WriteRune('"')
			wrote = true
		}

		if len(sig.signature) > 0 {
			// If InsecureNoSignature is enabled, signature may not exist
			if wrote {
				buf.WriteRune(tokens.Comma)
			}
			buf.WriteString(`"signature":"`)
			buf.Write(base64.Encode(sig.signature))
			buf.WriteString(`"`)
		}
		buf.WriteString(`}`)
	}
	buf.WriteString(`]}`)

	ret := bytes.Clone(buf.Bytes())
	return ret, nil
}

// Compact generates a JWS message in compact serialization format from
// `*jws.Message` object. The object contain exactly one signature, or
// an error is returned.
//
// If using a detached payload, the payload must already be stored in
// the `*jws.Message` object, and the `jws.WithDetached()` option
// must be passed to the function.
func Compact(msg *Message, options ...CompactOption) ([]byte, error) {
	if l := len(msg.signatures); l != 1 {
		return nil, makeSignError(prefixJwsCompact, `cannot serialize message with %d signatures (must be one)`, l)
	}

	var detached bool
	var encoder Base64Encoder = base64.DefaultEncoder()
	for _, option := range options {
		switch option.Ident() {
		case identDetached{}:
			if err := option.Value(&detached); err != nil {
				return nil, makeSignError(prefixJwsCompact, `failed to retrieve detached option value: %w`, err)
			}
		case identBase64Encoder{}:
			if err := option.Value(&encoder); err != nil {
				return nil, makeSignError(prefixJwsCompact, `failed to retrieve base64 encoder option value: %w`, err)
			}
		}
	}

	s := msg.signatures[0]
	// XXX check if this is correct
	hdrs := s.ProtectedHeaders()

	hdrbuf, err := json.Marshal(hdrs)
	if err != nil {
		return nil, makeSignError(prefixJwsCompact, `failed to marshal headers: %w`, err)
	}

	buf := pool.BytesBuffer().Get()
	defer pool.BytesBuffer().Put(buf)

	buf.WriteString(encoder.EncodeToString(hdrbuf))
	buf.WriteByte(tokens.Period)

	if !detached {
		if getB64Value(hdrs) {
			encoded := encoder.EncodeToString(msg.payload)
			buf.WriteString(encoded)
		} else {
			if bytes.Contains(msg.payload, []byte{tokens.Period}) {
				return nil, makeSignError(prefixJwsCompact, `compact serialization with b64=false requires payload to contain no "." characters per RFC 7797 §5.2; use jws.WithDetachedPayload to keep the payload out of the wire format`)
			}
			buf.Write(msg.payload)
		}
	}

	buf.WriteByte(tokens.Period)
	buf.WriteString(encoder.EncodeToString(s.signature))
	ret := bytes.Clone(buf.Bytes())
	return ret, nil
}