package integration_test

import (
	"crypto/ecdsa"
	"crypto/sha256"
	"fmt"
	"testing"

	"github.com/rwadurian/mpc-system/pkg/tss"
	"github.com/stretchr/testify/require"
)

// TestVariousThresholds tests different threshold configurations
func TestVariousThresholds(t *testing.T) {
	testCases := []struct {
		name          string
		threshold     int // t in tss-lib (t+1 signers required)
		totalParties  int
		signersNeeded int // actual signers needed = threshold + 1
	}{
		{
			name:          "2-of-3 (t=1, n=3)",
			threshold:     1,
			totalParties:  3,
			signersNeeded: 2,
		},
		{
			name:          "3-of-5 (t=2, n=5)",
			threshold:     2,
			totalParties:  5,
			signersNeeded: 3,
		},
		{
			name:          "4-of-7 (t=3, n=7)",
			threshold:     3,
			totalParties:  7,
			signersNeeded: 4,
		},
	}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			fmt.Printf("\n========================================\n")
			fmt.Printf("  Testing %s\n", tc.name)
			fmt.Printf("========================================\n")

			// Step 1: Key Generation
			fmt.Printf("\n[Step 1] Running Distributed Key Generation...\n")
			fmt.Printf("  - Threshold (t): %d (meaning t+1=%d signers required)\n", tc.threshold, tc.signersNeeded)
			fmt.Printf("  - Total Parties (n): %d\n", tc.totalParties)

			keygenResults, err := tss.RunLocalKeygen(tc.threshold, tc.totalParties)
			require.NoError(t, err, "Keygen should succeed")
			require.Len(t, keygenResults, tc.totalParties, "Should have correct number of key shares")

			publicKey := keygenResults[0].PublicKey
			require.NotNil(t, publicKey, "Public key should not be nil")

			fmt.Printf("  [OK] Key generation completed with %d parties!\n", tc.totalParties)
			fmt.Printf("      Public Key X: %s...\n", publicKey.X.Text(16)[:32])

			// Verify all parties have the same public key
			for i, result := range keygenResults {
				require.Equal(t, publicKey.X, result.PublicKey.X, "Party %d should have same X", i)
				require.Equal(t, publicKey.Y, result.PublicKey.Y, "Party %d should have same Y", i)
			}
			fmt.Println("      All parties have consistent public key")

			// Step 2: Test signing with exactly threshold+1 parties
			fmt.Printf("\n[Step 2] Testing threshold signing with %d-of-%d...\n", tc.signersNeeded, tc.totalParties)

			message := []byte(fmt.Sprintf("Test message for %s", tc.name))
			messageHash := sha256.Sum256(message)

			// Use first signersNeeded parties
			signers := keygenResults[:tc.signersNeeded]
			signResult, err := tss.RunLocalSigning(tc.threshold, signers, messageHash[:])
			require.NoError(t, err, "Signing should succeed")
			require.NotNil(t, signResult.R, "R should not be nil")
			require.NotNil(t, signResult.S, "S should not be nil")

			// Verify signature
			valid := ecdsa.Verify(publicKey, messageHash[:], signResult.R, signResult.S)
			require.True(t, valid, "Signature should verify")

			fmt.Printf("  [OK] Signature with %d parties verified!\n", tc.signersNeeded)

			// Step 3: Verify fewer than threshold parties cannot sign
			if tc.signersNeeded > 2 {
				fmt.Printf("\n[Step 3] Verifying %d parties cannot sign (need %d)...\n", tc.signersNeeded-1, tc.signersNeeded)
				insufficientSigners := keygenResults[:tc.signersNeeded-1]
				_, err = tss.RunLocalSigning(tc.threshold, insufficientSigners, messageHash[:])
				require.Error(t, err, "Signing with insufficient parties should fail")
				fmt.Printf("  [OK] Correctly rejected signing with insufficient parties\n")
			}

			fmt.Printf("\n========================================\n")
			fmt.Printf("  %s: PASSED\n", tc.name)
			fmt.Printf("========================================\n")
		})
	}
}

// Test3of5Flow tests 3-of-5 specifically with multiple combinations
func Test3of5Flow(t *testing.T) {
	fmt.Println("\n========================================")
	fmt.Println("  3-of-5 MPC Full Flow Test")
	fmt.Println("========================================")

	threshold := 2 // t=2 means t+1=3 signers required
	totalParties := 5

	// Key Generation
	fmt.Println("\n[Keygen] Generating keys for 5 parties...")
	keygenResults, err := tss.RunLocalKeygen(threshold, totalParties)
	require.NoError(t, err)
	require.Len(t, keygenResults, 5)

	publicKey := keygenResults[0].PublicKey
	fmt.Printf("  [OK] 5 key shares generated\n")
	fmt.Printf("      Public Key: %s...\n", publicKey.X.Text(16)[:32])

	message := []byte("3-of-5 threshold signing test")
	messageHash := sha256.Sum256(message)

	// Test multiple 3-party combinations
	combinations := [][]int{
		{0, 1, 2},
		{0, 1, 3},
		{0, 2, 4},
		{1, 3, 4},
		{2, 3, 4},
	}

	fmt.Println("\n[Signing] Testing various 3-party combinations...")
	for _, combo := range combinations {
		signers := []*tss.LocalKeygenResult{
			keygenResults[combo[0]],
			keygenResults[combo[1]],
			keygenResults[combo[2]],
		}

		signResult, err := tss.RunLocalSigning(threshold, signers, messageHash[:])
		require.NoError(t, err, "Signing with parties %v should succeed", combo)

		valid := ecdsa.Verify(publicKey, messageHash[:], signResult.R, signResult.S)
		require.True(t, valid, "Signature from parties %v should verify", combo)

		fmt.Printf("  [OK] Parties %v: signature verified\n", combo)
	}

	fmt.Println("\n========================================")
	fmt.Println("  3-of-5 Flow: ALL PASSED")
	fmt.Println("========================================")
}

// Test4of7Flow tests 4-of-7 specifically
func Test4of7Flow(t *testing.T) {
	fmt.Println("\n========================================")
	fmt.Println("  4-of-7 MPC Full Flow Test")
	fmt.Println("========================================")

	threshold := 3 // t=3 means t+1=4 signers required
	totalParties := 7

	// Key Generation
	fmt.Println("\n[Keygen] Generating keys for 7 parties...")
	keygenResults, err := tss.RunLocalKeygen(threshold, totalParties)
	require.NoError(t, err)
	require.Len(t, keygenResults, 7)

	publicKey := keygenResults[0].PublicKey
	fmt.Printf("  [OK] 7 key shares generated\n")
	fmt.Printf("      Public Key: %s...\n", publicKey.X.Text(16)[:32])

	message := []byte("4-of-7 threshold signing test")
	messageHash := sha256.Sum256(message)

	// Test a few 4-party combinations
	combinations := [][]int{
		{0, 1, 2, 3},
		{0, 2, 4, 6},
		{1, 3, 5, 6},
		{3, 4, 5, 6},
	}

	fmt.Println("\n[Signing] Testing various 4-party combinations...")
	for _, combo := range combinations {
		signers := []*tss.LocalKeygenResult{
			keygenResults[combo[0]],
			keygenResults[combo[1]],
			keygenResults[combo[2]],
			keygenResults[combo[3]],
		}

		signResult, err := tss.RunLocalSigning(threshold, signers, messageHash[:])
		require.NoError(t, err, "Signing with parties %v should succeed", combo)

		valid := ecdsa.Verify(publicKey, messageHash[:], signResult.R, signResult.S)
		require.True(t, valid, "Signature from parties %v should verify", combo)

		fmt.Printf("  [OK] Parties %v: signature verified\n", combo)
	}

	// Verify 3 parties cannot sign
	fmt.Println("\n[Security] Verifying 3 parties cannot sign...")
	insufficientSigners := keygenResults[:3]
	_, err = tss.RunLocalSigning(threshold, insufficientSigners, messageHash[:])
	require.Error(t, err, "3 parties should not be able to sign in 4-of-7")
	fmt.Println("  [OK] Correctly rejected 3-party signing")

	fmt.Println("\n========================================")
	fmt.Println("  4-of-7 Flow: ALL PASSED")
	fmt.Println("========================================")
}