package protocol_test

import (
	"encoding/hex"
	"encoding/json"
	"testing"

	"git.n1ko.dev/Niko/niko_trust/pkg/address"
	"git.n1ko.dev/Niko/niko_trust/pkg/protocol"
	"git.n1ko.dev/Niko/niko_trust/pkg/tce"
)

// Golden tests against the frozen reference vectors.
//
// These are the heart of cross-implementation agreement: the Go code must
// reproduce every byte of the Python reference implementation's output, must
// verify every reference signature, and must assign every object the same
// SHA-256 object ID. Any disagreement here means the Go implementation is
// wrong, not the vectors (docs/PROTOCOL.md section 14).

type vectorJSON struct {
	Type string `json:"type"`
}

// jsonString and jsonUint pull fields out of the JSON view for comparison.
func jsonString(t *testing.T, raw json.RawMessage, key string) string {
	t.Helper()
	var m map[string]any
	if err := json.Unmarshal(raw, &m); err != nil {
		t.Fatalf("json view: %v", err)
	}
	s, _ := m[key].(string)
	return s
}

func jsonNonce(t *testing.T, raw json.RawMessage) string {
	return jsonString(t, raw, "nonce")
}

// reencodeAndCompare decodes one object, encodes its public fields back, and
// requires the result to equal the received bytes byte for byte. This is the
// cross-implementation identity: the encoder and the Python reference must
// agree, and the decoder must have read exactly the reference's field order.
func reencodeAndCompare(t *testing.T, label string, got []byte, want []byte) {
	t.Helper()
	if !eqBytes(got, want) {
		t.Fatalf("%s: encoder did not reproduce the reference bytes\n got %x\nwant %x", label, got, want)
	}
}

func eqBytes(a, b []byte) bool {
	if len(a) != len(b) {
		return false
	}
	for i := range a {
		if a[i] != b[i] {
			return false
		}
	}
	return true
}

func TestVectorPartiesDerivedFromSeeds(t *testing.T) {
	vf := loadVectors(t)
	if len(vf.Parties) != 2 {
		t.Fatalf("expect exactly two parties, got %d", len(vf.Parties))
	}
	for name, p := range vf.Parties {
		seed := mustHex(t, p.SeedHex)
		s := seedSigner(t, seed)
		if hex.EncodeToString(s.Public()) != p.PubkeyHex {
			t.Errorf("%s: public key does not match the reference", name)
		}
		if s.Address().String() != p.Address {
			t.Errorf("%s: address does not match the reference", name)
		}
	}
}

func TestVectorGolden(t *testing.T) {
	vf := loadVectors(t)
	if len(vf.Vectors) != 13 {
		t.Fatalf("expect 9 object vectors, got %d", len(vf.Vectors))
	}

	byName := map[string]vectorEntry{}
	for _, v := range vf.Vectors {
		byName[v.Name] = v
	}
	claimBool := byName["claim/boolean"]
	req := byName["approval_request/ban"]
	rotReq := byName["key_rotation/request"]

	var revClaim *protocol.Claim
	var rev *protocol.Revocation

	for _, v := range vf.Vectors {
		t.Run(v.Name, func(t *testing.T) {
			tceBytes := mustHex(t, v.TCEHex)
			if len(tceBytes) != v.TCELen {
				t.Fatalf("tce_len: field says %d, hex has %d", v.TCELen, len(tceBytes))
			}
			sig := mustHex(t, v.SignatureHex)
			if got, want := hex.EncodeToString(tce.ComputeID(tceBytes).Bytes()), v.ObjectIDHex; got != want {
				t.Fatalf("object ID mismatch\n got %s\nwant %s", got, want)
			}

			var typ vectorJSON
			if err := json.Unmarshal(v.JSON, &typ); err != nil {
				t.Fatal(err)
			}

			switch typ.Type {
			case "identity":
				obj, err := protocol.VerifyIdentity(tceBytes, sig)
				if err != nil {
					t.Fatalf("VerifyIdentity: %v", err)
				}
				reencodeAndCompare(t, "identity", mustTCE(t, obj), tceBytes)
				if obj.Alias != jsonString(t, v.JSON, "alias") {
					t.Errorf("alias decoded as %q", obj.Alias)
				}
				if obj.CreatedAt != 1_700_000_000 {
					t.Errorf("created_at = %d", obj.CreatedAt)
				}
				checkAddress(t, obj.PubKey, v.SignerAddress)

			case "claim":
				obj, err := protocol.VerifyClaim(tceBytes, sig)
				if err != nil {
					t.Fatalf("VerifyClaim: %v", err)
				}
				reencodeAndCompare(t, "claim", mustTCE(t, obj), tceBytes)
				if len(obj.Claims) < 1 || len(obj.Claims) > 32 {
					t.Errorf("claims size %d out of bounds", len(obj.Claims))
				}
				checkAddress(t, obj.Issuer, v.SignerAddress)
				if hex.EncodeToString(obj.Nonce) != jsonNonce(t, v.JSON) {
					t.Errorf("nonce mismatch")
				}
				// The two response vectors and the revocation reference this
				// claim by ID; capture the verified claim for later checks.
				if v.Name == "claim/boolean" {
					revClaim = obj
				}

			case "revocation":
				obj, err := protocol.VerifyRevocation(tceBytes, sig)
				if err != nil {
					t.Fatalf("VerifyRevocation: %v", err)
				}
				reencodeAndCompare(t, "revocation", mustTCE(t, obj), tceBytes)
				checkAddress(t, obj.Issuer, v.SignerAddress)
				rev = obj
				// The frozen revocation withdraws claim/boolean by ID; verify
				// the binding passes against the decoded claim.
				if claimBool.ObjectIDHex != "" && obj.ClaimID.String() != claimBool.ObjectIDHex {
					t.Errorf("revocation targets %s, want the boolean claim", obj.ClaimID)
				}

			case "approval_request":
				obj, err := protocol.VerifyApprovalRequest(tceBytes, sig)
				if err != nil {
					t.Fatalf("VerifyApprovalRequest: %v", err)
				}
				reencodeAndCompare(t, "approval_request", mustTCE(t, obj), tceBytes)
				checkAddress(t, obj.Sender, v.SignerAddress)
				if obj.ExpiresAt-obj.CreatedAt > 60 {
					t.Errorf("request lifetime %d exceeds the format bound", obj.ExpiresAt-obj.CreatedAt)
				}

			case "approval_response":
				reqBytes := mustHex(t, req.TCEHex)
				reqSig := mustHex(t, req.SignatureHex)
				obj, err := protocol.VerifyApprovalResponse(reqBytes, reqSig, tceBytes, sig)
				if err != nil {
					t.Fatalf("VerifyApprovalResponse: %v", err)
				}
				reencodeAndCompare(t, "approval_response", mustTCE(t, obj), tceBytes)
				checkAddress(t, obj.Responder, v.SignerAddress)
				if jsonString(t, v.JSON, "decision") != obj.Decision.String() {
					t.Errorf("decision field mismatch")
				}
				// request_hash is the object ID of the exact request.
				if obj.RequestHash.String() != req.ObjectIDHex {
					t.Errorf("request_hash does not match the request object ID")
				}

			case "auth_assertion":
				obj, err := protocol.VerifyAuthAssertion(tceBytes, sig, "trust.n1ko.dev")
				if err != nil {
					t.Fatalf("VerifyAuthAssertion: %v", err)
				}
				reencodeAndCompare(t, "auth_assertion", mustTCE(t, obj), tceBytes)
				checkAddress(t, obj.PubKey, v.SignerAddress)
				if jsonString(t, v.JSON, "challenge") != hex.EncodeToString(obj.Challenge) {
					t.Errorf("challenge mismatch")
				}

			case "delegation":
				obj, err := protocol.VerifyDelegationClaim(tceBytes, sig)
				if err != nil {
					t.Fatalf("VerifyDelegationClaim: %v", err)
				}
				reencodeAndCompare(t, "delegation", mustTCE(t, obj), tceBytes)
				checkAddress(t, obj.Granter, v.SignerAddress)
				if len(obj.Predicates) < 1 || len(obj.Predicates) > 32 {
					t.Errorf("predicates size %d out of bounds", len(obj.Predicates))
				}
				for k := range obj.Predicates {
					if !obj.Covers(k) {
						t.Errorf("Covers(%q) false for a delegated predicate", k)
					}
				}
			case "key_rotation_request":
				obj, err := protocol.VerifyKeyRotationRequest(tceBytes, sig)
				if err != nil {
					t.Fatalf("VerifyKeyRotationRequest: %v", err)
				}
				reencodeAndCompare(t, "key_rotation_request", mustTCE(t, obj), tceBytes)
				checkAddress(t, obj.Successor, v.SignerAddress)
				if obj.ExpiresAt-obj.CreatedAt > 60 {
					t.Errorf("lifetime %d exceeds the format bound", obj.ExpiresAt-obj.CreatedAt)
				}

			case "key_rotation_confirm":
				reqBytes := mustHex(t, rotReq.TCEHex)
				reqSig := mustHex(t, rotReq.SignatureHex)
				obj, err := protocol.VerifyKeyRotationConfirm(reqBytes, reqSig, tceBytes, sig)
				if err != nil {
					t.Fatalf("VerifyKeyRotationConfirm: %v", err)
				}
				reencodeAndCompare(t, "key_rotation_confirm", mustTCE(t, obj), tceBytes)
				wantHash := rotReq.ObjectIDHex
				if obj.RotationHash.String() != wantHash {
					t.Errorf("rotation_hash does not match the request object ID")
				}
				_ = wantHash

			default:
				t.Fatalf("unknown json type %q", typ.Type)
			}
		})
	}

	// Cross-object bindings that only become checkable once every vector is
	// verified.
	t.Run("revocation withdraws the boolean claim", func(t *testing.T) {
		if revClaim == nil || rev == nil {
			t.Fatal("precondition: claim/boolean and revocation not both decoded")
		}
		if err := protocol.VerifyRevocationOf(rev, revClaim); err != nil {
			t.Fatalf("VerifyRevocationOf: %v", err)
		}
	})
}

// mustTCE returns the verified object's canonical bytes.
func mustTCE(t *testing.T, obj interface{ TCE() []byte }) []byte {
	t.Helper()
	return obj.TCE()
}

// checkAddress asserts that the decoded key renders as the reference address.
func checkAddress(t *testing.T, pub []byte, wantAddr string) {
	t.Helper()
	a, err := address.FromPubKey(pub)
	if err != nil {
		t.Fatalf("decoded key is not a valid address: %v", err)
	}
	if a.String() != wantAddr {
		t.Fatalf("address %s, want %s", a.String(), wantAddr)
	}
}

// TestVectorSignatureMutation is the signature half of the cross-check: a
// signature that verifies over the unchanged bytes must fail over any
// mutation of the TCE bytes (see mutation_test.go for the byte-by-byte
// sweep).