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
	//HandleConnReq(wrapper[ConnReq]) connState
	HandlePing(wrapper[Ping]) connState
	HandlePong(wrapper[Pong]) connState
	HandleTimeout() connState
}

// Helper functions.

func newStateFromPeerUpdate(update peerUpdate, data *connData) connState {
	if update.Peer != nil {
		return newStateFromPeer(update.Peer, data)
	}
	return newConnNull(data)
}

func newStateFromPeer(peer *m.Peer, data *connData) connState {
	if _, isPublic := netip.AddrFromSlice(peer.PublicIP); isPublic {
		return newStateServerDown(data, peer)
	} else if data.server {
		return newStateClientDown(data, peer)
	} else {
		return newStateMediated(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.useMediator = false
	c.up = false
	c.route.Store(nil)
	return c
}

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

func (c connNull) HandleConnReq(w wrapper[ConnReq]) connState {
	logState(c, "Ignoring conn request.")
	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.Second) // Ping right away to bring up.
	c.timeoutTimer.Stop()          // No timeouts yet.
	c.addr = c.publicAddr
	c.useMediator = false
	c.up = false
	c.route.Store(c.Route())

	return c
}

func (c stateServerDown) Name() string {
	return "Server:DOWN"
}

func (c stateServerDown) HandleConnReq(w wrapper[ConnReq]) connState {
	// Send ConnResp.
	// TODO
	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.useMediator = false
	c.up = true
	c.route.Store(c.Route())
	return c
}

func (c stateServerUp) Name() string {
	return "Server:UP"
}

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 newStateFromPeer(c.peer, c.connData)
}

////////////////////////
// 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.useMediator = false
	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) HandlePing(w wrapper[Ping]) connState {
	log.Printf("Got ping...")
	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.useMediator = false
	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) 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 newStateFromPeer(c.peer, c.connData)
}

//////////////
// Mediated //
//////////////

type stateMediated struct {
	*connData
}

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

	c := stateMediated{data}
	c.peer = peer
	c.publicAddr = pubAddr
	c.encPrivKey = data.encPrivKey
	c.encSharedKey = computeSharedKey(peer.EncPubKey, c.encPrivKey)
	c.addr = c.publicAddr
	c.useMediator = true
	c.up = true
	c.route.Store(c.Route())

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

func (c stateMediated) Name() string {
	return "Mediated:UP"
}

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

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

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