package webauthn

import (
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"encoding/binary"
	"encoding/json"
	"testing"
)

// --- minimal CBOR encoders (only what a COSE key + attestation object need) ---

func cUint(u uint64) []byte {
	switch {
	case u < 24:
		return []byte{byte(u)}
	case u < 256:
		return []byte{0x18, byte(u)}
	default:
		return []byte{0x19, byte(u >> 8), byte(u)}
	}
}

// cNeg encodes a negative int n (n<0). arg = -1-n.
func cNeg(n int64) []byte {
	arg := uint64(-1 - n)
	b := cUint(arg)
	b[0] |= 0x20 // major type 1
	return b
}

func cBytes(b []byte) []byte {
	h := cUint(uint64(len(b)))
	h[0] |= 0x40 // major type 2
	return append(h, b...)
}

func cText(s string) []byte {
	h := cUint(uint64(len(s)))
	h[0] |= 0x60 // major type 3
	return append(h, []byte(s)...)
}

func cMapHeader(n int) []byte {
	h := cUint(uint64(n))
	h[0] |= 0xa0 // major type 5
	return h
}

// coseKey CBOR-encodes an ES256/P-256 public key as a COSE_Key map.
func coseKey(pub *ecdsa.PublicKey) []byte {
	x := pub.X.Bytes()
	y := pub.Y.Bytes()
	// left-pad to 32 bytes
	px := make([]byte, 32)
	py := make([]byte, 32)
	copy(px[32-len(x):], x)
	copy(py[32-len(y):], y)
	var out []byte
	kv := func(k, v []byte) { out = append(append(out, k...), v...) }
	out = append(out, cMapHeader(5)...)
	kv(cUint(1), cUint(2))   // kty: EC2
	kv(cUint(3), cNeg(-7))   // alg: ES256
	kv(cNeg(-1), cUint(1))   // crv: P-256
	kv(cNeg(-2), cBytes(px)) // x
	kv(cNeg(-3), cBytes(py)) // y
	return out
}

func b64(b []byte) string { return base64.RawURLEncoding.EncodeToString(b) }

// authData builds an authenticatorData blob. For registration it embeds the
// attested credential data (AT flag + COSE key); for assertion it's the 37-byte
// header only.
func authData(rpID string, flags byte, counter uint32, credID []byte, cose []byte) []byte {
	h := sha256.Sum256([]byte(rpID))
	d := append([]byte{}, h[:]...)
	d = append(d, flags)
	cb := make([]byte, 4)
	binary.BigEndian.PutUint32(cb, counter)
	d = append(d, cb...)
	if flags&(1<<6) != 0 { // AT set → attested credential data
		d = append(d, make([]byte, 16)...) // aaguid
		l := make([]byte, 2)
		binary.BigEndian.PutUint16(l, uint16(len(credID)))
		d = append(d, l...)
		d = append(d, credID...)
		d = append(d, cose...)
	}
	return d
}

func clientData(typ, challenge, origin string) []byte {
	b, _ := json.Marshal(map[string]string{"type": typ, "challenge": challenge, "origin": origin})
	return b
}

const testOrigin = "https://maven.test"
const testRPID = "maven.test"

// TestRegisterAssertRoundTrip drives the full passkey flow with a real P-256
// key: register a credential, then assert it and verify the ecdsa signature
// check passes end to end.
func TestRegisterAssertRoundTrip(t *testing.T) {
	rp := NewRP(Config{Origin: testOrigin, RPID: testRPID, RPName: "maven"})
	key, _ := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	cose := coseKey(&key.PublicKey)
	credID := []byte("cred-1")
	credIDb64 := b64(credID)

	// --- register ---
	_, regChal, err := rp.CreationOptions([]byte("u"), "user")
	if err != nil {
		t.Fatal(err)
	}
	att := append(cMapHeader(3), cText("fmt")...)
	att = append(att, cText("none")...)
	att = append(att, cText("attStmt")...)
	att = append(att, cMapHeader(0)...)
	att = append(att, cText("authData")...)
	att = append(att, cBytes(authData(testRPID, 1<<6|0x05, 0, credID, cose))...)

	var stored []byte
	save := func(id string, pk, _ []byte, _ string) error { stored = pk; return nil }
	gotID, err := rp.FinishRegistration(save, regChal, map[string]any{
		"id": credIDb64,
		"response": map[string]any{
			"clientDataJSON":    b64(clientData("webauthn.create", regChal, testOrigin)),
			"attestationObject": b64(att),
		},
	})
	if err != nil {
		t.Fatalf("register: %v", err)
	}
	if gotID != credIDb64 || len(stored) == 0 {
		t.Fatalf("register produced no credential")
	}

	// --- assert (valid signature) ---
	credID64 := gotID
	sign := func(chal string, flags byte, tamper bool) map[string]any {
		ad := authData(testRPID, flags, 5, nil, nil)
		cdj := clientData("webauthn.get", chal, testOrigin)
		hash := sha256.Sum256(cdj)
		sig, _ := ecdsa.SignASN1(rand.Reader, key, append(append([]byte{}, ad...), hash[:]...))
		if tamper {
			sig[len(sig)-1] ^= 0xff
		}
		return map[string]any{
			"id": credID64,
			"response": map[string]any{
				"clientDataJSON":    b64(cdj),
				"authenticatorData": b64(ad),
				"signature":         b64(sig),
			},
		}
	}
	lookup := func(id string) ([]byte, int64, error) { return stored, 0, nil }
	upd := func(id string, c int64) error { return nil }

	_, assertChal, _ := rp.AssertionOptions()
	if _, err := rp.FinishAssertion(lookup, upd, assertChal, sign(assertChal, 0x05, false)); err != nil {
		t.Fatalf("valid assertion should pass: %v", err)
	}

	// --- negative: tampered signature ---
	_, chal2, _ := rp.AssertionOptions()
	if _, err := rp.FinishAssertion(lookup, upd, chal2, sign(chal2, 0x05, true)); err == nil {
		t.Fatal("tampered signature must fail verification")
	}

	// --- negative: user-verification flag not set (no gesture) ---
	_, chal3, _ := rp.AssertionOptions()
	if _, err := rp.FinishAssertion(lookup, upd, chal3, sign(chal3, 0x01, false)); err == nil {
		t.Fatal("assertion without UV flag must fail (step-up requires a gesture)")
	}
}

// TestAssertRejectsWrongOrigin — a phished assertion from another origin fails.
func TestAssertRejectsWrongOrigin(t *testing.T) {
	err := verifyClientDataBytes(clientData("webauthn.get", "abc", "https://evil.test"), "webauthn.get", "abc", testOrigin)
	if err == nil {
		t.Fatal("wrong origin must be rejected")
	}
}