package pnyx

import (
	"encoding/binary"
	"errors"
	"fmt"
	"net"
  "os"
	"sync/atomic"

	"github.com/google/uuid"
)

const (
 SERVER_UDP_BUFFER_SIZE = 2048

 PUBKEY_LENGTH = 32
 ECDH_LENGTH = 32
 SIGNATURE_LENGTH = 64
 SESSION_OPEN_LENGTH = PUBKEY_LENGTH + ECDH_LENGTH + SIGNATURE_LENGTH
)

type PacketType uint16
const (
  SESSION_OPEN PacketType = iota
  SESSION_AUTHENTICATE
  SESSION_CONNECT
  SESSION_CLOSE
  SESSION_CLOSED
  SESSION_DATA
)

/*
Session Flow:
1. Client send a SESSION_OPEN packet to the server with first half for ECDH and it's public key
2. Server responds with an SESSION_AUTHENTICATE packet with the server's ECDH half and public key
aside - at this point if the client and server both hold the private parts of their public keys,
        they both hold the same ECDH secret which they put in a KDF to generate a symmetric session key
aside - at this point the server creates the session in memory, but there is no return address associated
        with it yet
3. Client sends a SESSION_CONNECT packet to the server with the session ID in cleartext and the
   return address hashed with the key(to prove the return address has not been modified without the key) 
4. Server adds the return address to the session info, and maps the address to the session for future packets
5. Server sends the HELLO packet to the client encrypted by the session key

If a client disconnects at any point and gets a new return address:
1. Client sends a SESSION_CONNECT packet to the server from the new socket
2. Server removes the old return address, and fills/maps the new return address

If a client wants to gracefully disconnect and notify the server to close the session:
1. Client sends a SESSION_CLOSE
2. Server responds with SESSION_CLOSED

Session Packets:
1. SESSION_OPEN
 Payload is CLIENT_PUBKEY + ECDH_HALF + SIGNATURE
3. SESSION_CONNECT
4. SESSION_CLOSE
5. SESSION_CLOSED
*/

type SessionID uuid.UUID
type ClientID uuid.UUID

type Connection struct {
  state string
  session SessionID
}

type Session struct {
  state string
  client ClientID
}

type Server struct {
  active atomic.Bool
  connection *net.UDPConn
  stopped chan error

  connections map[string]SessionID
  sessions map[SessionID]Session
}

func NewServer() *Server {
  server := &Server{
    connection: nil,
    active: atomic.Bool{},
    stopped: make(chan error, 0),
  }
  server.active.Store(false)
  return server
}

func (server *Server) Log(format string, fields ...interface{}) {
  fmt.Fprint(os.Stderr, fmt.Sprintf(format, fields...) + "\n")
}

func(server *Server) Stop() error {
  was_active := server.active.CompareAndSwap(true, false)
  if was_active {
    err := server.connection.Close()
    if err != nil {
      return err
    }
    return <-server.stopped
  } else {
    return fmt.Errorf("Called stop func on stopped server")
  }
}

func(server *Server) run() {
  server.Log("Started server on %s", server.connection.LocalAddr())

  var buf [SERVER_UDP_BUFFER_SIZE]byte
  for true {
    read, addr, err := server.connection.ReadFromUDP(buf[:])
    if err == nil {
      var packet_type PacketType = PacketType(binary.BigEndian.Uint16(buf[0:2]))
      switch packet_type {
      default:
        server.Log("Unhandled packet type 0x%04x from %s: %+v", packet_type, addr, buf[:read])
      }
    } else if errors.Is(err, net.ErrClosed) {
      server.Log("UDP_CLOSE: %s", server.connection.LocalAddr())
      break
    } else {
      server.Log("UDP_READ_ERROR: %s", err)
    }
  }

  server.Log("Shut down server on %s", server.connection.LocalAddr())
  server.stopped <- nil
}

func(server *Server) Start(listen string) error {
  was_inactive := server.active.CompareAndSwap(false, true)
  if was_inactive == false {
    return fmt.Errorf("Server already active")
  }

  address, err := net.ResolveUDPAddr("udp", listen)
  if err != nil {
    server.active.Store(false)
    return err
  }

  server.connection, err = net.ListenUDP("udp", address)
  if err != nil {
    server.active.Store(false)
    return fmt.Errorf("Failed to create udp server: %w", err)
  }

  go server.run()

  return nil
}