package pnyx

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/ed25519"
	"crypto/rand"
	"crypto/sha512"
	"encoding/binary"
	"hash/crc32"
	"io"
	mrand "math/rand"
	"net"
	"time"

	"fmt"
	"slices"

	"filippo.io/edwards25519"
	"github.com/google/uuid"
)

type SessionID uuid.UUID
func(id SessionID) String() string {
  return uuid.UUID(id).String()
}

type PeerID [32]byte

type Session struct {
  ID SessionID
  remote *net.UDPAddr
  Peer PeerID
  secret []byte
  cipher cipher.Block
  iv_generator *mrand.Rand
}

type SessionPacketType uint8
const (
  REQ_ID_LENGTH = 16
  PEER_ID_LENGTH = 32
  SESSION_ID_LENGTH = 16
  IV_LENGTH = aes.BlockSize
  PUBKEY_LENGTH = 32
  SIGNATURE_LENGTH = 64
  HMAC_LENGTH = 64
  COMMAND_LENGTH = 1
  TIME_LENGTH = 8
  INVITE_LENGTH = 32

  SESSION_OPEN_LENGTH = PUBKEY_LENGTH + PUBKEY_LENGTH + SIGNATURE_LENGTH
  SESSION_OPENED_LENGTH = PUBKEY_LENGTH + PUBKEY_LENGTH + SESSION_ID_LENGTH + SIGNATURE_LENGTH
  SESSION_INVITE_LENGTH = INVITE_LENGTH + PUBKEY_LENGTH
  SESSION_TIMED_LENGTH = SESSION_ID_LENGTH + TIME_LENGTH + SIGNATURE_LENGTH
  SESSION_TIMED_RESP_LENGTH = TIME_LENGTH + SIGNATURE_LENGTH

  /*
    pnyx session packets
  */
  SESSION_OPEN SessionPacketType = iota
  SESSION_OPENED
  SESSION_INVITE
  SESSION_INVITED
  SESSION_CONNECT
  SESSION_CONNECTED
  SESSION_CLOSE
  SESSION_CLOSED
  SESSION_DATA
)

func(packet_type SessionPacketType) String() string {
  switch packet_type {
  case SESSION_OPEN:
    return "SESSION_OPEN"
  case SESSION_OPENED:
    return "SESSION_OPENED"
  case SESSION_INVITE:
    return "SESSION_INVITE"
  case SESSION_INVITED:
    return "SESSION_INVITED"
  case SESSION_CONNECT:
    return "SESSION_CONNECT"
  case SESSION_CONNECTED:
    return "SESSION_CONNECTED"
  case SESSION_CLOSE:
    return "SESSION_CLOSE"
  case SESSION_CLOSED:
    return "SESSION_CLOSED"
  case SESSION_DATA:
    return "SESSION_DATA"
  default:
    return fmt.Sprintf("UNKNOWN 0x%02x", uint8(packet_type))
  }
}

func ECDH(public ed25519.PublicKey, private ed25519.PrivateKey) ([]byte, error) {
  public_point, err := (&edwards25519.Point{}).SetBytes(public)
  if err != nil {
    return nil, err
  } 


  h := sha512.Sum512(private.Seed())
  private_scalar, err := (&edwards25519.Scalar{}).SetBytesWithClamping(h[:32])
  if err != nil {
    return nil, err
  }

  shared_point := public_point.ScalarMult(private_scalar, public_point)

  return shared_point.BytesMontgomery(), nil
}

type Invite [INVITE_LENGTH]byte

func NewSessionInvite(key ed25519.PublicKey, invite Invite) ([]byte, error) {
  if key == nil {
    return nil, fmt.Errorf("Cannot create a SESSION_INVITE packet without a key")
  }

  packet := make([]byte, COMMAND_LENGTH + SESSION_INVITE_LENGTH)
  
  // TODO

  return packet, nil
}

func NewSessionOpen(key ed25519.PrivateKey) ([]byte, ed25519.PrivateKey, error) {
  if key == nil {
    return nil, nil, fmt.Errorf("Cannot create a SESSION_OPEN packet without a key")
  }

  public, private, err := ed25519.GenerateKey(rand.Reader)
  if err != nil {
    return nil, nil, fmt.Errorf("Failed to generate ecdh key: %w", err)
  }

  packet := make([]byte, COMMAND_LENGTH + SESSION_OPEN_LENGTH)
  cur := 0

  packet[0] = byte(SESSION_OPEN)
  cur += COMMAND_LENGTH

  copy(packet[cur:], []byte(key.Public().(ed25519.PublicKey)))
  cur += PUBKEY_LENGTH

  copy(packet[cur:], []byte(public)) 
  cur += PUBKEY_LENGTH

  signature := ed25519.Sign(key, packet[COMMAND_LENGTH:cur])
  copy(packet[cur:], signature)
  cur += SIGNATURE_LENGTH

  return packet, private, nil
}

func ParseSessionOpen(key ed25519.PrivateKey, session_open []byte) (Session, []byte, error) {
  if len(session_open) != SESSION_OPEN_LENGTH {
    return Session{}, nil, fmt.Errorf("Bad SESSION_OPEN length: %d/%d", len(session_open), SESSION_OPEN_LENGTH)
  }

  cur := 0

  client_pubkey := (ed25519.PublicKey)(session_open[cur:cur+PUBKEY_LENGTH])
  cur += PUBKEY_LENGTH

  client_ecdh := (ed25519.PublicKey)(session_open[cur:cur+PUBKEY_LENGTH])
  cur += PUBKEY_LENGTH
  
  digest := session_open[:cur]
  signature := session_open[cur:cur+SIGNATURE_LENGTH]
  cur += SIGNATURE_LENGTH

  if ed25519.Verify(client_pubkey, digest, signature) == false {
    return Session{}, nil, fmt.Errorf("SESSION_OPEN signature verification failed")
  }

  seed_bytes := make([]byte, 8)
  read, err := rand.Read(seed_bytes)
  if err != nil {
    return Session{}, nil, err
  } else if read != 8 {
    return Session{}, nil, fmt.Errorf("Not enough entropy to create session")
  }

  rand_gen := mrand.New(mrand.NewSource(int64(binary.BigEndian.Uint64(seed_bytes))).(mrand.Source64))
  session_uuid, err := uuid.NewRandomFromReader(rand_gen)
  if err != nil {
    return Session{}, nil, err
  }

  session_id := SessionID(session_uuid)

  ecdh_public, ecdh_private, err := ed25519.GenerateKey(rand.Reader)
  if err != nil {
    return Session{}, nil, err
  }

  session_secret, err := ECDH(client_ecdh, ecdh_private)
  if err != nil {
    return Session{}, nil, err
  }

  session_cipher, err := aes.NewCipher(session_secret)
  if err != nil {
    return Session{}, nil, err
  }

  session_opened := make([]byte, COMMAND_LENGTH + SESSION_OPENED_LENGTH)
  cur = 0

  session_opened[cur] = byte(SESSION_OPENED)
  cur += COMMAND_LENGTH

  copy(session_opened[cur:], key.Public().(ed25519.PublicKey))
  cur += PUBKEY_LENGTH

  copy(session_opened[cur:], ecdh_public)
  cur += PUBKEY_LENGTH

  copy(session_opened[cur:], session_id[:])
  cur += SESSION_ID_LENGTH


  signature = ed25519.Sign(key, session_opened[COMMAND_LENGTH:cur])
  copy(session_opened[cur:], signature)
  cur += SIGNATURE_LENGTH

  return Session{
    ID: session_id,
    remote: nil,
    Peer: PeerID(client_pubkey),
    iv_generator: rand_gen,
    cipher: session_cipher,
    secret: session_secret,
  }, session_opened, nil
}

func ParseSessionOpened(expected_pubkey ed25519.PublicKey, ecdh_private ed25519.PrivateKey, session_opened []byte) (Session, error) {
  if len(session_opened) != SESSION_OPENED_LENGTH {
    return Session{}, fmt.Errorf("Wrong SESSION_OPENED length: %d/%d", len(session_opened), SESSION_OPEN_LENGTH)
  }

  cur := 0

  server_pubkey := (ed25519.PublicKey)(session_opened[cur:cur+PUBKEY_LENGTH])
  cur += PUBKEY_LENGTH

  server_ecdh := (ed25519.PublicKey)(session_opened[cur:cur+PUBKEY_LENGTH])
  cur += PUBKEY_LENGTH

  session_id := SessionID(session_opened[cur:cur+SESSION_ID_LENGTH])
  cur += SESSION_ID_LENGTH

  if expected_pubkey != nil && slices.Compare(server_pubkey, expected_pubkey) != 0 {
    return Session{}, fmt.Errorf("server public key %x does not match expected %x", server_pubkey, expected_pubkey)
  }

  if ed25519.Verify(server_pubkey, session_opened[:cur], session_opened[cur:cur+SIGNATURE_LENGTH]) == false {
    return Session{}, fmt.Errorf("session opened signature verification failed")
  }

  seed_bytes := make([]byte, 8)
  read, err := rand.Read(seed_bytes)
  if err != nil {
    return Session{}, err
  } else if read != 8 {
    return Session{}, fmt.Errorf("Not enough entropy to create session")
  }

  session_secret, err := ECDH(server_ecdh, ecdh_private)
  if err != nil {
    return Session{}, err
  }

  cipher, err := aes.NewCipher(session_secret)
  if err != nil {
    return Session{}, nil
  }

  return Session{
    ID: session_id,
    remote: nil,
    Peer: PeerID(server_pubkey),
    secret: session_secret,
    cipher: cipher,
    iv_generator: mrand.New(mrand.NewSource(int64(binary.BigEndian.Uint64(seed_bytes))).(mrand.Source64)),
  }, nil
}

func NewSessionTimed(command SessionPacketType, key ed25519.PrivateKey, session *Session, t time.Time) []byte {
  packet := make([]byte, COMMAND_LENGTH + SESSION_TIMED_LENGTH)
  cur := 0

  packet[cur] = byte(command)
  cur += COMMAND_LENGTH

  copy(packet[cur:], session.ID[:])
  cur += SESSION_ID_LENGTH

  binary.BigEndian.PutUint64(packet[cur:], uint64(t.UnixMilli()))
  cur += TIME_LENGTH

  signature := ed25519.Sign(key, packet[COMMAND_LENGTH:cur])
  copy(packet[cur:], signature)
  cur += SIGNATURE_LENGTH

  return packet
}

func ParseSessionTimed(resp_type SessionPacketType, key ed25519.PrivateKey, packet []byte, session *Session, last_t int64) ([]byte, error) {
  if len(packet) != SESSION_TIMED_LENGTH {
    return nil, fmt.Errorf("Bad timed packet length: %d/%d", len(packet), SESSION_TIMED_LENGTH)
  }

  cur := SESSION_ID_LENGTH
  t := int64(binary.BigEndian.Uint64(packet[cur:]))
  cur += TIME_LENGTH

  if t < last_t {
    return nil, fmt.Errorf("Time in packet to old: %d < %d", t, last_t)
  }

  if ed25519.Verify(ed25519.PublicKey(session.Peer[:]), packet[:cur], packet[cur:cur+SIGNATURE_LENGTH]) == false {
    return nil, fmt.Errorf("Failed to verify packet signature")
  }

  resp := make([]byte, COMMAND_LENGTH + SESSION_TIMED_RESP_LENGTH)
  cur = 0

  resp[cur] = byte(resp_type)
  cur += COMMAND_LENGTH
  
  binary.BigEndian.PutUint64(resp[cur:], uint64(t))
  cur += TIME_LENGTH

  signature := ed25519.Sign(key, resp[COMMAND_LENGTH:cur])
  copy(resp[cur:], signature)
  cur += SIGNATURE_LENGTH

  return resp, nil
}

func NewSessionData(session *Session, packet []byte) ([]byte, error) {
  iv := make([]byte, IV_LENGTH)
  for i := 0; i < IV_LENGTH/8; i++ {
    binary.BigEndian.PutUint64(iv[i*8:], session.iv_generator.Uint64())
  }

  stream := cipher.NewOFB(session.cipher, iv[:])
  header := make([]byte, COMMAND_LENGTH + SESSION_ID_LENGTH + IV_LENGTH)
  header[0] = byte(SESSION_DATA)
  copy(header[COMMAND_LENGTH:], session.ID[:])
  copy(header[COMMAND_LENGTH+SESSION_ID_LENGTH:], iv)

  packet_encrypted := bytes.NewBuffer(header)
  writer := &cipher.StreamWriter{S: stream, W: packet_encrypted}

  // Encrypt the packet with a crc32 checksum appended to the end
  _, err := io.Copy(writer, bytes.NewBuffer(binary.BigEndian.AppendUint32(packet, crc32.ChecksumIEEE(packet))))
  if err != nil {
    return nil, err
  }

  return packet_encrypted.Bytes(), nil
}

func ParseSessionData(session *Session, encrypted []byte) ([]byte, error) {
  iv := encrypted[0:IV_LENGTH]

  stream := cipher.NewOFB(session.cipher, iv)
  var packet bytes.Buffer
  reader := &cipher.StreamReader{S: stream, R: bytes.NewBuffer(encrypted[IV_LENGTH:])}
  _, err := io.Copy(&packet, reader)
  if err != nil {
    return nil, err
  }

  packet_clear := packet.Bytes()
  checksum := binary.BigEndian.Uint32(packet_clear[len(packet_clear)-4:])
  data := packet_clear[:len(packet_clear)-4]
  calculated := crc32.ChecksumIEEE(data)
  if checksum != calculated {
    return nil, fmt.Errorf("SESSION_DATA checksum mismatch: %x != %x", checksum, calculated)
  }

  return data, nil
}