package integration

import (
	"net"
	"testing"
	"time"

	"hash-of-wisdom/internal/protocol"
	"hash-of-wisdom/internal/server"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestSlowlorisProtection_SlowReader(t *testing.T) {
	// Setup server with very short read timeout for testing
	serverConfig := &server.Config{
		Address:  ":0",
		Timeouts: server.TimeoutConfig{
			Read:       100 * time.Millisecond,
			Write:      5 * time.Second,
			Connection: 15 * time.Second,
		},
	}

	srv := setupTestServerWithConfig(t, serverConfig)
	defer srv.Stop()

	// Connect to server
	conn, err := net.Dial("tcp", srv.Address())
	require.NoError(t, err)
	defer conn.Close()

	// Send partial message header very slowly (slowloris attack)
	_, err = conn.Write([]byte{0x01}) // Challenge request type
	require.NoError(t, err)

	// Wait longer than read timeout before sending length
	time.Sleep(200 * time.Millisecond)

	// Try to send more data - connection should be timed out
	_, err = conn.Write([]byte{0x00, 0x00, 0x00, 0x00}) // Payload length

	// Verify connection is closed by trying to read
	buffer := make([]byte, 1024)
	conn.SetReadDeadline(time.Now().Add(500 * time.Millisecond))
	_, err = conn.Read(buffer)
	assert.Error(t, err, "Connection should be closed due to slow reading")
}

func TestSlowlorisProtection_SlowWriter(t *testing.T) {
	// Setup server with very short write timeout for testing
	serverConfig := &server.Config{
		Address:  ":0",
		Timeouts: server.TimeoutConfig{
			Read:       5 * time.Second,
			Write:      100 * time.Millisecond,
			Connection: 15 * time.Second,
		},
	}

	srv := setupTestServerWithConfig(t, serverConfig)
	defer srv.Stop()

	// Connect to server but don't read responses (simulate slow writer client)
	conn, err := net.Dial("tcp", srv.Address())
	require.NoError(t, err)
	defer conn.Close()

	// Send a complete challenge request
	challengeReq := &protocol.ChallengeRequest{}
	err = challengeReq.Encode(conn)
	require.NoError(t, err)

	// Don't read the response - this should trigger write timeout
	time.Sleep(200 * time.Millisecond)

	// Try to write again - connection should be timed out
	_, err = conn.Write([]byte{0x01})

	// Verify connection is closed
	buffer := make([]byte, 1024)
	conn.SetReadDeadline(time.Now().Add(500 * time.Millisecond))
	_, err = conn.Read(buffer)
	assert.Error(t, err, "Connection should be closed due to slow writing")
}

func TestSlowlorisProtection_SlowConnectionTimeout(t *testing.T) {
	// Setup server with very short connection timeout for testing
	serverConfig := &server.Config{
		Address:  ":0",
		Timeouts: server.TimeoutConfig{
			Read:       5 * time.Second,
			Write:      5 * time.Second,
			Connection: 200 * time.Millisecond,
		},
	}

	srv := setupTestServerWithConfig(t, serverConfig)
	defer srv.Stop()

	// Connect to server but do nothing
	conn, err := net.Dial("tcp", srv.Address())
	require.NoError(t, err)
	defer conn.Close()

	// Wait longer than connection timeout
	time.Sleep(300 * time.Millisecond)

	// Try to read - connection should be closed
	buffer := make([]byte, 1024)
	conn.SetReadDeadline(time.Now().Add(500 * time.Millisecond))
	_, err = conn.Read(buffer)
	assert.Error(t, err, "Connection should be closed due to connection timeout")
}

func TestSlowlorisProtection_MultipleSlowClients(t *testing.T) {
	// Setup server with short timeouts
	serverConfig := &server.Config{
		Address:  ":0",
		Timeouts: server.TimeoutConfig{
			Read:       1 * time.Second,
			Write:      1 * time.Second,
			Connection: 2 * time.Second,
		},
	}

	srv := setupTestServerWithConfig(t, serverConfig)
	defer srv.Stop()

	// Create multiple slow connections
	var connections []net.Conn
	for i := 0; i < 5; i++ {
		conn, err := net.Dial("tcp", srv.Address())
		require.NoError(t, err)
		connections = append(connections, conn)
		// Send partial data on each connection
		conn.Write([]byte{0x01}) // Challenge request type only
	}

	// Wait for timeouts to kick in
	time.Sleep(1500 * time.Millisecond)

	// All connections should be closed
	buffer := make([]byte, 1024)
	for i, conn := range connections {
		conn.SetReadDeadline(time.Now().Add(500 * time.Millisecond))
		_, err := conn.Read(buffer)
		assert.Error(t, err, "Connection %d should be closed", i)
		conn.Close()
	}
}

func TestSlowlorisProtection_AttackMitigation(t *testing.T) {
	// Test that server can still serve legitimate clients during an attack
	serverConfig := &server.Config{
		Address:  ":0",
		Timeouts: server.TimeoutConfig{
			Read:       500 * time.Millisecond,
			Write:      500 * time.Millisecond,
			Connection: 1 * time.Second,
		},
	}

	srv := setupTestServerWithConfig(t, serverConfig)
	defer srv.Stop()

	// Start slowloris attack - multiple slow connections
	var attackConnections []net.Conn
	for i := 0; i < 3; i++ {
		conn, err := net.Dial("tcp", srv.Address())
		require.NoError(t, err)
		attackConnections = append(attackConnections, conn)
		// Send partial data
		conn.Write([]byte{0x01})
	}

	// Meanwhile, legitimate client should still work
	legitimateConn, err := net.Dial("tcp", srv.Address())
	require.NoError(t, err)
	defer legitimateConn.Close()

	// Send complete request quickly
	challengeReq := &protocol.ChallengeRequest{}
	err = challengeReq.Encode(legitimateConn)
	require.NoError(t, err)

	// Should receive response despite ongoing attack
	decoder := protocol.NewMessageDecoder()
	msg, err := decoder.Decode(legitimateConn)
	assert.NoError(t, err)
	assert.Equal(t, protocol.ChallengeResponseType, msg.Type)

	// Clean up attack connections
	for _, conn := range attackConnections {
		conn.Close()
	}
}

func TestSlowlorisProtection_SlowChunkReading(t *testing.T) {
	// Test protection against clients that read responses very slowly
	serverConfig := &server.Config{
		Address:  ":0",
		Timeouts: server.TimeoutConfig{
			Read:       5 * time.Second,
			Write:      200 * time.Millisecond,
			Connection: 10 * time.Second,
		},
	}

	srv := setupTestServerWithConfig(t, serverConfig)
	defer srv.Stop()

	// Connect and send valid request
	conn, err := net.Dial("tcp", srv.Address())
	require.NoError(t, err)
	defer conn.Close()

	challengeReq := &protocol.ChallengeRequest{}
	err = challengeReq.Encode(conn)
	require.NoError(t, err)

	// Read only first byte of response very slowly
	buffer := make([]byte, 1)
	_, err = conn.Read(buffer)
	require.NoError(t, err)

	// Wait longer than write timeout before reading more
	time.Sleep(300 * time.Millisecond)

	// Try to read more - connection should be closed due to slow reading
	conn.SetReadDeadline(time.Now().Add(500 * time.Millisecond))
	_, err = conn.Read(buffer)
	assert.Error(t, err, "Connection should be closed due to slow chunk reading")
}