// JCR-1 independent reference verifier — Go, standard library only.
//
// No external modules: a from-scratch canonical CBOR encoder (RFC 8949 4.2.1),
// a minimal COSE_Sign1 reader, crypto/ed25519, crypto/hmac, crypto/sha256.
// This is the *second implementation* whose agreement with the Python suite
// makes JCR-1 a standard rather than a house format.
//
// It does more than verify: it RE-ENCODES each signed core from the receipt
// fields and checks the bytes equal the Python-produced signed core exactly.
// Two independent encoders producing identical canonical bytes is the proof.
//
// Usage: go run jcr1_verify.go jcr1-vectors.json jcr1-test-keys.json
package main

import (
	"bytes"
	"crypto/ed25519"
	"crypto/hmac"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"os"
	"sort"
	"strconv"
)

// ---------- canonical CBOR encoder (subset: uint, text, bstr, array, map) ----

func cborUint(buf *bytes.Buffer, major byte, n uint64) {
	mt := major << 5
	switch {
	case n < 24:
		buf.WriteByte(mt | byte(n))
	case n < 1<<8:
		buf.WriteByte(mt | 24)
		buf.WriteByte(byte(n))
	case n < 1<<16:
		buf.WriteByte(mt | 25)
		buf.WriteByte(byte(n >> 8))
		buf.WriteByte(byte(n))
	case n < 1<<32:
		buf.WriteByte(mt | 26)
		for s := 24; s >= 0; s -= 8 {
			buf.WriteByte(byte(n >> uint(s)))
		}
	default:
		buf.WriteByte(mt | 27)
		for s := 56; s >= 0; s -= 8 {
			buf.WriteByte(byte(n >> uint(s)))
		}
	}
}

func cborText(buf *bytes.Buffer, s string) {
	cborUint(buf, 3, uint64(len(s)))
	buf.WriteString(s)
}

func cborBytes(buf *bytes.Buffer, b []byte) {
	cborUint(buf, 2, uint64(len(b)))
	buf.Write(b)
}

// encodeValue handles the value tree produced by json with UseNumber.
func encodeValue(buf *bytes.Buffer, v interface{}) error {
	switch x := v.(type) {
	case string:
		cborText(buf, x)
	case json.Number:
		// signed core integers are all non-negative
		u, err := strconv.ParseUint(x.String(), 10, 64)
		if err != nil {
			return fmt.Errorf("non-integer in signed core: %s", x.String())
		}
		cborUint(buf, 0, u)
	case bool, float64, nil:
		return fmt.Errorf("disallowed type in signed core: %T", x)
	case []interface{}:
		cborUint(buf, 4, uint64(len(x)))
		for _, e := range x {
			if err := encodeValue(buf, e); err != nil {
				return err
			}
		}
	case map[string]interface{}:
		return encodeMap(buf, x)
	default:
		return fmt.Errorf("unsupported type %T", x)
	}
	return nil
}

// encodeMap: RFC 8949 4.2.1 ordering — length-first, then bytewise on key.
func encodeMap(buf *bytes.Buffer, m map[string]interface{}) error {
	keys := make([]string, 0, len(m))
	for k := range m {
		keys = append(keys, k)
	}
	sort.Slice(keys, func(i, j int) bool {
		if len(keys[i]) != len(keys[j]) {
			return len(keys[i]) < len(keys[j])
		}
		return keys[i] < keys[j]
	})
	cborUint(buf, 5, uint64(len(keys)))
	for _, k := range keys {
		cborText(buf, k)
		if err := encodeValue(buf, m[k]); err != nil {
			return err
		}
	}
	return nil
}

// signedCore: drop token_attributions and null fields, canonical-encode.
func signedCore(receipt map[string]interface{}) ([]byte, error) {
	core := map[string]interface{}{}
	for k, v := range receipt {
		if k == "token_attributions" || v == nil {
			continue
		}
		core[k] = v
	}
	var buf bytes.Buffer
	if err := encodeMap(&buf, core); err != nil {
		return nil, err
	}
	return buf.Bytes(), nil
}

// ---------- minimal CBOR reader (only what COSE_Sign1 needs) ----------------

type reader struct {
	b   []byte
	pos int
}

func (r *reader) byte() (byte, error) {
	if r.pos >= len(r.b) {
		return 0, fmt.Errorf("eof")
	}
	c := r.b[r.pos]
	r.pos++
	return c, nil
}

func (r *reader) head() (major byte, arg uint64, err error) {
	c, err := r.byte()
	if err != nil {
		return
	}
	major = c >> 5
	ai := c & 0x1f
	switch {
	case ai < 24:
		arg = uint64(ai)
	case ai == 24:
		var v byte
		v, err = r.byte()
		arg = uint64(v)
	case ai == 25:
		for i := 0; i < 2; i++ {
			var v byte
			if v, err = r.byte(); err != nil {
				return
			}
			arg = arg<<8 | uint64(v)
		}
	case ai == 26:
		for i := 0; i < 4; i++ {
			var v byte
			if v, err = r.byte(); err != nil {
				return
			}
			arg = arg<<8 | uint64(v)
		}
	case ai == 27:
		for i := 0; i < 8; i++ {
			var v byte
			if v, err = r.byte(); err != nil {
				return
			}
			arg = arg<<8 | uint64(v)
		}
	default:
		err = fmt.Errorf("unsupported additional info %d", ai)
	}
	return
}

func (r *reader) bstr() ([]byte, error) {
	major, n, err := r.head()
	if err != nil {
		return nil, err
	}
	if major != 2 {
		return nil, fmt.Errorf("expected bstr, got major %d", major)
	}
	if r.pos+int(n) > len(r.b) {
		return nil, fmt.Errorf("bstr overrun")
	}
	out := r.b[r.pos : r.pos+int(n)]
	r.pos += int(n)
	return out, nil
}

// parseCOSESign1 extracts protected, payload, signature from a tag-18 array.
func parseCOSESign1(b []byte) (protected, payload, signature []byte, err error) {
	r := &reader{b: b}
	major, tag, err := r.head()
	if err != nil {
		return
	}
	if major == 6 { // optional tag 18
		if tag != 18 {
			err = fmt.Errorf("unexpected tag %d", tag)
			return
		}
		major, tag, err = r.head()
		if err != nil {
			return
		}
	}
	if major != 4 || tag != 4 {
		err = fmt.Errorf("expected 4-element array, got major %d len %d", major, tag)
		return
	}
	if protected, err = r.bstr(); err != nil { // [0] protected
		return
	}
	if m, a, e := r.head(); e != nil || m != 5 { // [1] unprotected map
		err = fmt.Errorf("expected map for unprotected header (major %d)", m)
		_ = a
		return
	}
	if payload, err = r.bstr(); err != nil { // [2] payload
		return
	}
	signature, err = r.bstr() // [3] signature
	return
}

// ---------- verification --------------------------------------------------

var provOrder = map[string]int{"HwShunt": 3, "ModelBased": 2, "Estimator": 1}

func worst(provs []string) string {
	w := provs[0]
	for _, p := range provs[1:] {
		if provOrder[p] < provOrder[w] {
			w = p
		}
	}
	return w
}

type Vectors struct {
	VectorCount int `json:"vector_count"`
	Envelope    string `json:"envelope"`
	Vectors     []struct {
		N           int                    `json:"n"`
		Description string                 `json:"description"`
		Receipt     map[string]interface{} `json:"receipt"`
		CoreHex     string                 `json:"signed_core_cbor_hex"`
		CoseHex     string                 `json:"cose_sign1_hex"`
	} `json:"vectors"`
}

type Keys struct {
	PubHex string            `json:"ed25519_public_raw_hex"`
	GwKeys map[string]string `json:"gateway_hmac_keys"`
}

func num(v interface{}) uint64 {
	u, _ := strconv.ParseUint(v.(json.Number).String(), 10, 64)
	return u
}

func verify(rec map[string]interface{}, coreHex, coseHex string, pub ed25519.PublicKey, gw map[string][]byte) []string {
	var errs []string

	// 1 — re-encode determinism (independent encoder must match Python bytes)
	core, err := signedCore(rec)
	if err != nil {
		errs = append(errs, "encode error: "+err.Error())
	}
	want, _ := hex.DecodeString(coreHex)
	if !bytes.Equal(core, want) {
		errs = append(errs, "signed core re-encode mismatch")
	}

	// 2 — COSE_Sign1 envelope signature
	coseBytes, _ := hex.DecodeString(coseHex)
	protected, payload, signature, perr := parseCOSESign1(coseBytes)
	if perr != nil {
		errs = append(errs, "cose parse: "+perr.Error())
	} else {
		if !bytes.Equal(payload, want) {
			errs = append(errs, "envelope payload != signed core")
		}
		var sig bytes.Buffer
		sig.WriteByte(0x84) // array(4)
		cborText(&sig, "Signature1")
		cborBytes(&sig, protected)
		cborBytes(&sig, []byte{})
		cborBytes(&sig, payload)
		if !ed25519.Verify(pub, sig.Bytes(), signature) {
			errs = append(errs, "ed25519 verification failed")
		}
	}

	// 3 — gateway tool-receipt HMACs
	if trs, ok := rec["tool_receipts"].([]interface{}); ok {
		for _, t := range trs {
			tr := t.(map[string]interface{})
			gid := tr["gateway_id"].(string)
			key, ok := gw[gid]
			if !ok {
				errs = append(errs, "no key for gateway "+gid)
				continue
			}
			payload := bytes.Join([][]byte{
				[]byte(tr["tool_id"].(string)),
				[]byte(strconv.FormatUint(num(tr["joules_uj"]), 10)),
				[]byte(tr["energy_provenance"].(string)),
				[]byte(tr["minted_at"].(string)),
				[]byte(gid),
			}, []byte{0})
			mac := hmac.New(sha256.New, key)
			mac.Write(payload)
			got := hex.EncodeToString(mac.Sum(nil))
			if !hmac.Equal([]byte(got), []byte(tr["hmac_sha256_hex"].(string))) {
				errs = append(errs, "HMAC mismatch "+tr["tool_id"].(string))
			}
		}
	}

	// 4 — honesty consistency
	var provs []string
	collect := func(field string) {
		if arr, ok := rec[field].([]interface{}); ok {
			for _, e := range arr {
				provs = append(provs, e.(map[string]interface{})["energy_provenance"].(string))
			}
		}
	}
	collect("tools_touched")
	collect("tool_receipts")
	if len(provs) > 0 && rec["energy_provenance"].(string) != worst(provs) {
		errs = append(errs, "energy_provenance != worst counter")
	}
	var sum uint64
	if arr, ok := rec["tools_touched"].([]interface{}); ok {
		for _, e := range arr {
			sum += num(e.(map[string]interface{})["joules_uj"])
		}
	}
	total := num(rec["joules_uj"])
	if total < sum {
		errs = append(errs, "joules_uj < sum(tools_touched)")
	}
	prov := rec["energy_provenance"].(string)
	if (prov == "ModelBased" || prov == "Estimator") && total%1000 != 0 {
		errs = append(errs, "mJ resolution floor violated")
	}
	return errs
}

func mustDecodeJSON(path string, v interface{}) {
	f, err := os.Open(path)
	if err != nil {
		panic(err)
	}
	defer f.Close()
	dec := json.NewDecoder(f)
	dec.UseNumber()
	if err := dec.Decode(v); err != nil {
		panic(err)
	}
}

func main() {
	vecPath := "jcr1-vectors.json"
	keyPath := "jcr1-test-keys.json"
	if len(os.Args) > 2 {
		vecPath, keyPath = os.Args[1], os.Args[2]
	}
	var vs Vectors
	var ks Keys
	mustDecodeJSON(vecPath, &vs)
	mustDecodeJSON(keyPath, &ks)

	pubBytes, _ := hex.DecodeString(ks.PubHex)
	pub := ed25519.PublicKey(pubBytes)
	gw := map[string][]byte{}
	for k, v := range ks.GwKeys {
		gw[k] = []byte(v)
	}

	fmt.Printf("JCR-1 independent Go verifier — %d vectors — %s\n\n", vs.VectorCount, vs.Envelope)
	pass := 0
	for _, v := range vs.Vectors {
		errs := verify(v.Receipt, v.CoreHex, v.CoseHex, pub, gw)
		ok := len(errs) == 0
		if ok {
			pass++
		}
		status := "PASS"
		if !ok {
			status = "FAIL"
		}
		fmt.Printf("  [%s] vector %2d: %s\n", status, v.N, v.Description)
		for _, e := range errs {
			fmt.Printf("         - %s\n", e)
		}
	}
	fmt.Printf("\n%d/%d vectors verified by the independent Go implementation\n", pass, vs.VectorCount)
	if pass != vs.VectorCount {
		os.Exit(1)
	}
}