package pnyx

import (
	"crypto/ed25519"
	"crypto/rand"
	"slices"
	"testing"
)

func fatalErr(t *testing.T, err error) {
  if err != nil {
    t.Fatal(err)
  }
}

func TestSessionOpen(t *testing.T) {
  client_pubkey, client_key, err := ed25519.GenerateKey(rand.Reader)
  fatalErr(t, err)

  server_pubkey, server_key, err := ed25519.GenerateKey(rand.Reader)
  fatalErr(t, err)

  session_open, client_ecdh, err := NewSessionOpen(client_key)
  fatalErr(t, err)
  if session_open[0] != byte(SESSION_OPEN) {
    t.Fatalf("Session open command byte mismatch(%x != %x)", session_open[0], SESSION_OPEN)
  }

  server_session, session_opened, err := ParseSessionOpen(server_key, session_open[COMMAND_LENGTH:])
  fatalErr(t, err)

  client_session, err := ParseSessionOpened(server_pubkey, client_ecdh, session_opened[COMMAND_LENGTH:])
  fatalErr(t, err)

  if PeerID(client_pubkey) != server_session.Peer {
    t.Fatalf("Server session(%x) has wrong peer ID(%x)", server_session.Peer, client_pubkey)
  }
  
  if PeerID(server_pubkey) != client_session.Peer {
    t.Fatalf("Client session(%x) has wrong peer ID(%x)", client_session.Peer, server_pubkey)
  }

  if slices.Compare(server_session.secret, client_session.secret) != 0 {
    t.Fatalf("Client secret(%x) and server secret(%x) do not match", client_session.secret, server_session.secret)
  }
}

func TestECDH(t *testing.T) {
  client_public, client_private, err := ed25519.GenerateKey(rand.Reader)
  fatalErr(t, err)

  server_public, server_private, err := ed25519.GenerateKey(rand.Reader)
  fatalErr(t, err)

  server_secret, err := ECDH(client_public, server_private)
  fatalErr(t, err)

  client_secret, err := ECDH(server_public, client_private)
  fatalErr(t, err)

  if slices.Compare(server_secret, client_secret) != 0 {
    t.Fatalf("Server and Client secrets do not match")
  }
}

func TestSessionData(t *testing.T) {
  _, client_key, err := ed25519.GenerateKey(rand.Reader)
  fatalErr(t, err)

  server_pubkey, server_key, err := ed25519.GenerateKey(rand.Reader)
  fatalErr(t, err)

  session_open, client_ecdh, err := NewSessionOpen(client_key)
  fatalErr(t, err)

  server_session, session_opened, err := ParseSessionOpen(server_key, session_open[COMMAND_LENGTH:])
  fatalErr(t, err)

  client_session, err := ParseSessionOpened(server_pubkey, client_ecdh, session_opened[COMMAND_LENGTH:])
  fatalErr(t, err)

  message := []byte("hello")
  session_data, err := NewSessionData(&server_session, message)
  fatalErr(t, err)

  if session_data[0] != byte(SESSION_DATA) {
    t.Fatalf("Session data command byte mismatch(%x != %x)", session_data[0], SESSION_DATA)
  } else if SessionID(session_data[COMMAND_LENGTH:COMMAND_LENGTH+SESSION_ID_LENGTH]) != server_session.ID {
    t.Fatal("Session data ID mismatch")
  }

  parsed_message, err := ParseSessionData(&client_session, session_data[COMMAND_LENGTH+SESSION_ID_LENGTH:])
  fatalErr(t, err)

  if slices.Compare(message, parsed_message) != 0 {
    t.Fatalf("Parsed message(%s) != test message(%s)", parsed_message, message)
  }
}