package verify

import (
	"crypto/ecdsa"
	"crypto/ed25519"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"github.com/Fuwn/plutia/internal/config"
	"github.com/Fuwn/plutia/internal/types"
	"github.com/decred/dcrd/dcrec/secp256k1/v4"
	secpECDSA "github.com/decred/dcrd/dcrec/secp256k1/v4/ecdsa"
	"github.com/fxamacker/cbor/v2"
	"github.com/mr-tron/base58"
	"testing"
)

func TestVerifyOperationValidSignature(t *testing.T) {
	pub, priv, err := ed25519.GenerateKey(rand.Reader)

	if err != nil {
		t.Fatalf("generate key: %v", err)
	}

	payloadDoc := []byte(`{"did":"did:plc:alice","didDoc":{"id":"did:plc:alice"}}`)
	sig := ed25519.Sign(priv, payloadDoc)
	opJSON := map[string]any{
		"did":        "did:plc:alice",
		"didDoc":     map[string]any{"id": "did:plc:alice"},
		"publicKey":  base64.RawURLEncoding.EncodeToString(pub),
		"sigPayload": base64.RawURLEncoding.EncodeToString(payloadDoc),
		"sig":        base64.RawURLEncoding.EncodeToString(sig),
	}
	raw, _ := json.Marshal(opJSON)
	op, err := types.ParseOperation(types.ExportRecord{
		Seq:       1,
		DID:       "did:plc:alice",
		Operation: raw,
	})

	if err != nil {
		t.Fatalf("parse operation: %v", err)
	}

	v := New(config.VerifyFull)

	if err := v.VerifyOperation(op, nil); err != nil {
		t.Fatalf("verify operation: %v", err)
	}
}

func TestVerifyOperationPrevMismatch(t *testing.T) {
	pub, priv, err := ed25519.GenerateKey(rand.Reader)

	if err != nil {
		t.Fatalf("generate key: %v", err)
	}

	payloadDoc := []byte(`{"did":"did:plc:alice","didDoc":{"id":"did:plc:alice"},"prev":"sha256:wrong"}`)
	sig := ed25519.Sign(priv, payloadDoc)
	opJSON := map[string]any{
		"did":        "did:plc:alice",
		"didDoc":     map[string]any{"id": "did:plc:alice"},
		"prev":       "sha256:wrong",
		"publicKey":  base64.RawURLEncoding.EncodeToString(pub),
		"sigPayload": base64.RawURLEncoding.EncodeToString(payloadDoc),
		"sig":        base64.RawURLEncoding.EncodeToString(sig),
	}
	raw, _ := json.Marshal(opJSON)
	op, err := types.ParseOperation(types.ExportRecord{Seq: 2, DID: "did:plc:alice", Operation: raw})

	if err != nil {
		t.Fatalf("parse operation: %v", err)
	}

	v := New(config.VerifyFull)
	existing := &types.StateV1{DID: "did:plc:alice", ChainTipHash: "sha256:right", LatestOpSeq: 1}

	if err := v.VerifyOperation(op, existing); err == nil {
		t.Fatalf("expected prev mismatch error")
	}
}

func TestVerifyOperationSecp256k1(t *testing.T) {
	priv, err := secp256k1.GeneratePrivateKey()

	if err != nil {
		t.Fatalf("generate secp256k1 key: %v", err)
	}

	pubKey := priv.PubKey()
	didKeyBytes := append([]byte{0xE7, 0x01}, pubKey.SerializeCompressed()...)
	didKey := "did:key:z" + base58.Encode(didKeyBytes)
	unsigned := map[string]any{
		"type":       "create",
		"prev":       nil,
		"handle":     "alice.example.com",
		"service":    "https://example.com",
		"signingKey": didKey,
	}
	enc, err := cbor.CanonicalEncOptions().EncMode()

	if err != nil {
		t.Fatalf("init cbor encoder: %v", err)
	}

	payload, err := enc.Marshal(unsigned)

	if err != nil {
		t.Fatalf("marshal cbor: %v", err)
	}

	sum := sha256.Sum256(payload)
	sig := secpECDSA.Sign(priv, sum[:]).Serialize()
	opJSON := map[string]any{
		"type":       "create",
		"prev":       nil,
		"handle":     "alice.example.com",
		"service":    "https://example.com",
		"signingKey": didKey,
		"sig":        base64.RawURLEncoding.EncodeToString(sig),
	}
	raw, _ := json.Marshal(opJSON)
	op, err := types.ParseOperation(types.ExportRecord{
		Seq:       1,
		DID:       "did:plc:alice",
		Operation: raw,
	})

	if err != nil {
		t.Fatalf("parse operation: %v", err)
	}

	v := New(config.VerifyFull)

	if err := v.VerifyOperation(op, nil); err != nil {
		t.Fatalf("verify operation: %v", err)
	}
}

func TestVerifyOperationP256(t *testing.T) {
	priv, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)

	if err != nil {
		t.Fatalf("generate p256 key: %v", err)
	}

	compressed := elliptic.MarshalCompressed(priv.Curve, priv.PublicKey.X, priv.PublicKey.Y)
	didKeyBytes := append([]byte{0x80, 0x24}, compressed...)
	didKey := "did:key:z" + base58.Encode(didKeyBytes)
	unsigned := map[string]any{
		"type":       "create",
		"prev":       nil,
		"handle":     "alice.example.com",
		"service":    "https://example.com",
		"signingKey": didKey,
	}
	enc, err := cbor.CanonicalEncOptions().EncMode()

	if err != nil {
		t.Fatalf("init cbor encoder: %v", err)
	}

	payload, err := enc.Marshal(unsigned)

	if err != nil {
		t.Fatalf("marshal cbor: %v", err)
	}

	sum := sha256.Sum256(payload)
	r, s, err := ecdsa.Sign(rand.Reader, priv, sum[:])

	if err != nil {
		t.Fatalf("sign p256: %v", err)
	}

	sig := make([]byte, 64)
	rb := r.Bytes()
	sb := s.Bytes()

	copy(sig[32-len(rb):32], rb)
	copy(sig[64-len(sb):], sb)

	opJSON := map[string]any{
		"type":       "create",
		"prev":       nil,
		"handle":     "alice.example.com",
		"service":    "https://example.com",
		"signingKey": didKey,
		"sig":        base64.RawURLEncoding.EncodeToString(sig),
	}
	raw, _ := json.Marshal(opJSON)
	op, err := types.ParseOperation(types.ExportRecord{
		Seq:       1,
		DID:       "did:plc:alice",
		Operation: raw,
	})

	if err != nil {
		t.Fatalf("parse operation: %v", err)
	}

	v := New(config.VerifyFull)

	if err := v.VerifyOperation(op, nil); err != nil {
		t.Fatalf("verify operation: %v", err)
	}
}