package main

import (
	"crypto/rand"
	"math/big"
	"net/http"
	"testing"
	"time"

	"party.at/party/internal/crypto"
)

func getVotes(t *testing.T, ts *testServer, issueID int64, optionID int64, adminToken string) int64 {
	 
	// -----------------------------------------------------------------------
	// 8. Tally
	// -----------------------------------------------------------------------

	code, _, res := ts.getWithToken(t, "/v1/issues/" + i64str(issueID), adminToken)
	if code != http.StatusOK {
		t.Fatalf("read issue: want 200 got %d: %s", code, res)
	}

	var resp struct {
		Options []struct {
			ID        int64  `json:"id"`
			VoteCount int64  `json:"vote_count"`
		} `json:"options"`
	}
	mustUnmarshal(t, res, &resp)

	for _, o := range resp.Options {
		if o.ID == optionID {
			return o.VoteCount
		}
	}

	return 0
}

func vote(t *testing.T, ts *testServer, issueID int64, optionID int64) {
	voterEmail := uniqueEmail()
	voterToken := ts.registerAndLogin(t, voterEmail, "pa$$word123")
	ts.assignRole(t, voterEmail, "contributor")

	code, _, res := ts.getWithToken(t, "/v1/issues/" + i64str(issueID)+"/pubkey", voterToken)
	if code != http.StatusOK {
		t.Fatalf("get pubkey: want 200 got %d: %s", code, res)
	}

	var pubKeyResp struct {
		PublicKey struct {
			N string `json:"n"`
			E int    `json:"e"`
		} `json:"public_key"`
	}
	mustUnmarshal(t, res, &pubKeyResp)

	nInt, ok := new(big.Int).SetString(pubKeyResp.PublicKey.N, 16)
	if !ok {
		t.Fatal("failed to parse public key modulus")
	}
	eInt := big.NewInt(int64(pubKeyResp.PublicKey.E))

	nonce := make([]byte, 16)
	if _, err := rand.Read(nonce); err != nil {
		t.Fatalf("generate nonce: %v", err)
	}

	m := crypto.VoteMessage(issueID, optionID, nonce)
	mInt := new(big.Int).SetBytes(m[:])
 
	r, err := generateBlindingFactor(nInt)
	if err != nil {
		t.Fatalf("generate blinding factor: %v", err)
	}

	// After fetching pubkey:
	t.Logf("n from pubkey (first 8 bytes): %x", nInt.Bytes()[:8])

 
	// m' = m * r^e mod n
	re := new(big.Int).Exp(r, eInt, nInt)
	mBlind := new(big.Int).Mod(new(big.Int).Mul(mInt, re), nInt)

	// -----------------------------------------------------------------------
	// 4. Client sends blinded vote to server (authenticated)
	// -----------------------------------------------------------------------
 
	code, _, res = ts.postJSONWithToken(t, "/v1/issues/" + i64str(issueID)+"/blind-sign", voterToken, map[string]any{
		"blinded_vote": mBlind.Bytes(),
	})
	if code != http.StatusOK {
		t.Fatalf("blind-sign: want 200 got %d: %s", code, res)
	}
 
	var blindSigResp struct {
		Signed []byte `json:"signed"`
	}
	mustUnmarshal(t, res, &blindSigResp)

	// -----------------------------------------------------------------------
	// 5. Client unblinds: s = s' * r^-1 mod n
	// -----------------------------------------------------------------------

	t.Logf("raw blind-sign response: %s", res)
 
	sPrime := new(big.Int).SetBytes(blindSigResp.Signed)
	// After getting blind sig back, before unblinding:
	t.Logf("s' (first 8 bytes): %x", sPrime.Bytes()[:8])
	rInv := new(big.Int).ModInverse(r, nInt)
	if rInv == nil {
		t.Fatal("could not compute modular inverse of r")
	}
	s := new(big.Int).Mod(new(big.Int).Mul(sPrime, rInv), nInt)
 
	// Sanity check: verify s^e mod n == m before submitting
	recovered := new(big.Int).Exp(s, eInt, nInt)
	if recovered.Cmp(mInt) != 0 {
		t.Fatal("unblinded signature verification failed: s^e mod n != m")
	}

	// After unblinding:
	t.Logf("s (first 8 bytes): %x", s.Bytes()[:8])
	t.Logf("m (expected):       %x", mInt.Bytes())
	t.Logf("s^e mod n:          %x", recovered.Bytes())

	// -----------------------------------------------------------------------
	// 6. Client logs out -- anonymous submission has no auth header
	// -----------------------------------------------------------------------
 
	code, _, res = ts.postJSON(t, "/v1/votes", map[string]any{
		"issue_id":  issueID,
		"option_id": optionID,
		"nonce":     nonce,
		"signature": s.Bytes(),
	})
	if code != http.StatusCreated {
		t.Fatalf("submit vote: want 201 got %d: %s", code, res)
	}

	// -----------------------------------------------------------------------
	// 7. Replay attack -- same signature should be rejected
	// -----------------------------------------------------------------------

	code, _, _ = ts.postJSON(t, "/v1/votes", map[string]any{
		"issue_id":  issueID,
		"option_id": optionID,
		"nonce":     nonce,
		"signature": s.Bytes(),
	})
	if code != http.StatusConflict {
		t.Errorf("replay attack: want 409 got %d", code)
	}
}

func TestVote(t *testing.T) {
	app := newTestApplication(t)
	ts := newTestServer(t, app, routes(app))
	defer ts.Close()

	// -----------------------------------------------------------------------
	// 1. Admin registers, logs in, creates an issue with options
	// -----------------------------------------------------------------------

	adminEmail := uniqueEmail()
	adminToken := ts.registerAndLogin(t, adminEmail, "pa$$word123")

	issueBody := map[string]any{
		"title":       "Should we adopt a four-day work week?",
		"description": "Vote yes or no.",
		"start_time":  time.Now().Format(time.RFC3339),
		"end_time":    time.Now().Add(24 * time.Hour).Format(time.RFC3339),
		"options":     []string{"Yes", "No"},
	}
	code, _, res := ts.postJSONWithToken(t, "/v1/issues", adminToken, issueBody)
		if code != http.StatusCreated {
		t.Fatalf("create issue: want 201 got %d: %s", code, res)
	}

	var issueResp struct {
		Issue struct {
			ID int64 `json:"id"`
		} `json:"issue"`
		Options []struct {
			ID    int64  `json:"id"`
			Label string `json:"label"`
		} `json:"options"`
	}

	mustUnmarshal(t, res, &issueResp)
	issueID := issueResp.Issue.ID
 
	var yesOptionID, noOptionID int64
	for _, opt := range issueResp.Options {
		switch opt.Label {
		case "Yes":
			yesOptionID = opt.ID
		case "No":
			noOptionID = opt.ID
		}
	}
	if yesOptionID == 0 || noOptionID == 0 {
		t.Fatalf("expected Yes and No options in response, got: %v", issueResp.Options)
	}

	// -----------------------------------------------------------------------
	// 2. Normal user registers, logs in, fetches the public key
	// -----------------------------------------------------------------------

	count := getVotes(t, ts, issueID, yesOptionID, adminToken)
	if count != 0 {
		t.Fatalf("expected 0 yes votes")
	}

	count = getVotes(t, ts, issueID, noOptionID, adminToken)
	if count != 0 {
		t.Fatalf("expected 0 no votes")
	}

	vote(t, ts, issueID, yesOptionID)

	count = getVotes(t, ts, issueID, yesOptionID, adminToken)
	if count != 1 {
		t.Fatalf("expected 1 yes votes")
	}

	vote(t, ts, issueID, yesOptionID)

	count = getVotes(t, ts, issueID, yesOptionID, adminToken)
	if count != 2 {
		t.Fatalf("expected 2 yes votes")
	}

	vote(t, ts, issueID, noOptionID)

	count = getVotes(t, ts, issueID, noOptionID, adminToken)
	if count != 1 {
		t.Fatalf("expected 1 no votes")
	}
}
