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
	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 newStateServerDown(data, peer)
	} else if data.server {
		return newStateClientDown(data, peer)
	} else {
		return newStateMediatedDown(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 stateServerDown struct {
	*connData
}

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

	c := stateServerDown{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 stateServerDown) Name() string {
	return "Server:DOWN"
}

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

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

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

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

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

type stateServerUp struct {
	*connData
}

func newStateServerUp(data *connData, w wrapper[Pong]) connState {
	c := stateServerUp{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 stateServerUp) Name() string {
	return "Server:UP"
}

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

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

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

func (c stateServerUp) HandleTimeout() connState {
	return newStateServerDown(c.connData, c.peer)
}

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

type stateClientDown struct {
	*connData
}

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

	c := stateClientDown{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 stateClientDown) Name() string {
	return "Client:DOWN"
}

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

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

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

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

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

type stateClientUp struct {
	*connData
}

func newStateClientUp(data *connData, w wrapper[Ping]) connState {
	c := stateClientUp{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 stateClientUp) Name() string {
	return "Client:UP"
}

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

func (c stateClientUp) 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 stateClientUp) HandlePong(w wrapper[Pong]) connState {
	logState(c, "Ignoring pong.")
	return c
}

func (c stateClientUp) HandleTimeout() connState {
	return newStateClientDown(c.connData, c.peer)
}

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

type stateMediatedDown struct {
	*connData
}

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

	c := stateMediatedDown{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 newStateMediatedUp(data, mRoute)
	}

	return c
}

func (c stateMediatedDown) Name() string {
	return "Mediated:DOWN"
}

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

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

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

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

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

type stateMediatedUp struct {
	*connData
}

func newStateMediatedUp(data *connData, route *route) connState {
	c := stateMediatedUp{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 stateMediatedUp) Name() string {
	return "Mediated:UP"
}

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

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

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

func (c stateMediatedUp) HandleTimeout() connState {
	return newStateMediatedDown(c.connData, c.peer)
}