package peer

import (
	"log"
	"net/netip"
	"time"
	"vppn/m"
)

func logState(s connState, msg string, args ...any) {
	log.Printf("["+s.Name()+"] "+msg, args...)
}

// ----------------------------------------------------------------------------

// The connection state corresponds to what we're connected TO.
type connState interface {
	Name() string
	HandleMediatorUpdate(ip byte) connState
	HandleSendPing()
	HandlePing(wrapper[Ping]) connState
	HandlePong(wrapper[Pong]) connState
	HandleTimeout() connState
}

// Helper function.

func newConnStateFromPeer(update peerUpdate, data *connData) connState {
	peer := update.Peer

	if peer == nil {
		return newConnNull(data)
	}

	if _, isPublic := netip.AddrFromSlice(peer.PublicIP); isPublic {
		return newConnUnconnectedServer(data, peer)
	} else if data.server {
		return newConnUnconnectedClient(data, peer)
	} else {
		return newConnUnconnectedMediator(data, peer)
	}
}

/////////////////////
// Null Connection //
/////////////////////

type connNull struct {
	*connData
}

func newConnNull(data *connData) connState {
	c := connNull{data}
	c.peer = nil
	c.encSharedKey = nil
	c.publicAddr = netip.AddrPort{}
	c.pingTimer.Stop()
	c.timeoutTimer.Stop()
	c.addr = c.publicAddr
	c.viaIP = 0
	c.up = false
	c.route.Store(nil)
	return c
}

func (c connNull) Name() string {
	return "NoPeer"
}

func (c connNull) HandleMediatorUpdate(ip byte) connState {
	c.mediatorIP = ip
	return c
}

func (c connNull) HandlePing(w wrapper[Ping]) connState {
	logState(c, "Ignoring ping.")
	return c
}

func (c connNull) HandlePong(w wrapper[Pong]) connState {
	logState(c, "Ignoring pong.")
	return c
}

func (c connNull) HandleTimeout() connState {
	logState(c, "Unexpected timeout.")
	return c
}

////////////////////////
// Unconnected Server //
////////////////////////

type connUnconnectedServer struct {
	*connData
}

func newConnUnconnectedServer(data *connData, peer *m.Peer) connState {
	addr, _ := netip.AddrFromSlice(peer.PublicIP)
	pubAddr := netip.AddrPortFrom(addr, peer.Port)

	c := connUnconnectedServer{data}
	c.peer = peer
	c.encSharedKey = computeSharedKey(peer.EncPubKey, c.encPrivKey)
	c.publicAddr = pubAddr
	c.pingTimer.Reset(time.Millisecond) // Ping right away to bring up.
	c.timeoutTimer.Stop()               // No timeouts yet.
	c.addr = c.publicAddr
	c.viaIP = 0
	c.up = false
	c.route.Store(c.Route())

	return c
}

func (c connUnconnectedServer) Name() string {
	return "ServerUnconnected"
}

func (c connUnconnectedServer) HandleMediatorUpdate(ip byte) connState {
	// Server connection doesn't use a mediator.
	c.mediatorIP = ip
	return c
}

func (c connUnconnectedServer) HandlePing(w wrapper[Ping]) connState {
	logState(c, "Ignoring ping.")
	return c
}

func (c connUnconnectedServer) HandlePong(w wrapper[Pong]) connState {
	return newConnConnectedServer(c.connData, w)
}

func (c connUnconnectedServer) HandleTimeout() connState {
	logState(c, "Unexpected timeout.")
	return c
}

//////////////////////
// Connected Server //
//////////////////////

type connConnectedServer struct {
	*connData
}

func newConnConnectedServer(data *connData, w wrapper[Pong]) connState {
	c := connConnectedServer{data}
	c.pingTimer.Reset(pingInterval)
	c.timeoutTimer.Reset(timeoutInterval)
	c.addr = w.SrcAddr
	c.viaIP = 0
	c.up = true
	c.route.Store(c.Route())
	return c
}

func (c connConnectedServer) Name() string {
	return "ServerConnected"
}

func (c connConnectedServer) HandleMediatorUpdate(ip byte) connState {
	// Server connection doesn't use a mediator.
	c.mediatorIP = ip
	return c
}

func (c connConnectedServer) HandlePing(w wrapper[Ping]) connState {
	logState(c, "Ignoring ping.")
	return c
}

func (c connConnectedServer) HandlePong(w wrapper[Pong]) connState {
	c.timeoutTimer.Reset(timeoutInterval)
	return c
}

func (c connConnectedServer) HandleTimeout() connState {
	return newConnUnconnectedServer(c.connData, c.peer)
}

////////////////////////
// Unconnected Client //
////////////////////////

type connUnconnectedClient struct {
	*connData
}

func newConnUnconnectedClient(data *connData, peer *m.Peer) connState {
	addr, _ := netip.AddrFromSlice(peer.PublicIP)
	pubAddr := netip.AddrPortFrom(addr, peer.Port)

	c := connUnconnectedClient{data}
	c.peer = peer
	c.publicAddr = pubAddr
	c.encPrivKey = data.encPrivKey
	c.encSharedKey = computeSharedKey(peer.EncPubKey, c.encPrivKey)
	c.addr = c.publicAddr
	c.viaIP = 0
	c.up = false
	c.route.Store(c.Route())

	c.pingTimer.Stop()    // Conncection is from client => pings incoming.
	c.timeoutTimer.Stop() // No timeouts yet.

	return c
}

func (c connUnconnectedClient) Name() string {
	return "ClientUnconnected"
}

func (c connUnconnectedClient) HandleMediatorUpdate(ip byte) connState {
	// Client connection doesn't use a mediator.
	c.mediatorIP = ip
	return c
}

func (c connUnconnectedClient) HandlePing(w wrapper[Ping]) connState {
	next := newConnConnectedClient(c.connData, w)
	c.sendPong(w) // Have to send after transitionsing so route is ok.
	return next
}

func (c connUnconnectedClient) HandlePong(w wrapper[Pong]) connState {
	logState(c, "Ignorning pong.")
	return c
}

func (c connUnconnectedClient) HandleTimeout() connState {
	logState(c, "Unexpected timeout.")
	return c
}

//////////////////////
// Connected Client //
//////////////////////

type connConnectedClient struct {
	*connData
}

func newConnConnectedClient(data *connData, w wrapper[Ping]) connState {
	c := connConnectedClient{data}
	c.addr = w.SrcAddr
	c.viaIP = 0
	c.up = true
	c.route.Store(c.Route())

	c.pingTimer.Stop() // Conncection is from client => pings incoming.
	c.timeoutTimer.Reset(timeoutInterval)
	return c
}

func (c connConnectedClient) Name() string {
	return "ClientConnected"
}

func (c connConnectedClient) HandleMediatorUpdate(ip byte) connState {
	// Client connection doesn't use a mediator.
	c.mediatorIP = ip
	return c
}

func (c connConnectedClient) HandlePing(w wrapper[Ping]) connState {
	// The connection is from a client. If the client's address changes, we
	// should follow that change.
	if c.addr != w.SrcAddr {
		c.addr = w.SrcAddr
		c.route.Store(c.Route())
	}
	c.sendPong(w)
	c.timeoutTimer.Reset(timeoutInterval)
	return c
}

func (c connConnectedClient) HandlePong(w wrapper[Pong]) connState {
	logState(c, "Ignoring pong.")
	return c
}

func (c connConnectedClient) HandleTimeout() connState {
	return newConnUnconnectedClient(c.connData, c.peer)
}

//////////////////////////
// Unconnected Mediator //
//////////////////////////

type connUnconnectedMediator struct {
	*connData
}

func newConnUnconnectedMediator(data *connData, peer *m.Peer) connState {
	addr, _ := netip.AddrFromSlice(peer.PublicIP)
	pubAddr := netip.AddrPortFrom(addr, peer.Port)

	c := connUnconnectedMediator{data}
	c.peer = peer
	c.publicAddr = pubAddr
	c.encPrivKey = data.encPrivKey
	c.encSharedKey = computeSharedKey(peer.EncPubKey, c.encPrivKey)
	c.addr = c.publicAddr
	c.viaIP = 0
	c.up = false
	c.route.Store(c.Route())

	c.pingTimer.Stop()    // No pings for mediators.
	c.timeoutTimer.Stop() // No timeouts yet.

	// If we have a mediator route, we can connect.
	if mRoute := c.routes[c.mediatorIP].Load(); mRoute != nil {
		return newConnConnectedMediator(data, mRoute)
	}

	return c
}

func (c connUnconnectedMediator) Name() string {
	return "MediatorUnconnected"
}

func (c connUnconnectedMediator) HandleMediatorUpdate(ip byte) connState {
	c.mediatorIP = ip
	if mRoute := c.routes[c.mediatorIP].Load(); mRoute != nil {
		return newConnConnectedMediator(c.connData, mRoute)
	}
	return c
}

func (c connUnconnectedMediator) HandlePing(w wrapper[Ping]) connState {
	logState(c, "Ignorning ping.")
	return c
}

func (c connUnconnectedMediator) HandlePong(w wrapper[Pong]) connState {
	logState(c, "Ignorning pong.")
	return c
}

func (c connUnconnectedMediator) HandleTimeout() connState {
	logState(c, "Unexpected timeout.")
	return c
}

////////////////////////
// Connected Mediator //
////////////////////////

type connConnectedMediator struct {
	*connData
}

func newConnConnectedMediator(data *connData, route *route) connState {
	c := connConnectedMediator{data}
	c.addr = route.Addr
	c.viaIP = route.PeerIP
	c.up = true
	c.route.Store(c.Route())

	// No pings for mediated routes.
	c.pingTimer.Stop()
	c.timeoutTimer.Stop()
	return c
}

func (c connConnectedMediator) Name() string {
	return "MediatorConnected"
}

func (c connConnectedMediator) HandleMediatorUpdate(ip byte) connState {
	c.mediatorIP = ip
	if mRoute := c.routes[c.mediatorIP].Load(); mRoute != nil {
		return newConnConnectedMediator(c.connData, mRoute)
	}
	return newConnUnconnectedMediator(c.connData, c.peer)
}

func (c connConnectedMediator) HandlePing(w wrapper[Ping]) connState {
	logState(c, "Ignoring ping.")
	return c
}

func (c connConnectedMediator) HandlePong(w wrapper[Pong]) connState {
	logState(c, "Ignoring pong.")
	return c
}

func (c connConnectedMediator) HandleTimeout() connState {
	return newConnUnconnectedMediator(c.connData, c.peer)
}