package transport

import (
	"context"
	"fmt"
	"net"
	"strconv"
	"time"

	toki "git.toki-labs.com/toki/proto-socket/go"
	"go.uber.org/zap"

	iop "iop/proto/gen/iop"
)

const (
	heartbeatIntervalSec = 30
	// heartbeatWaitSec is kept above heartbeatIntervalSec as defence in depth:
	// the library wait-timer callback already self-heals stale state (see
	// proto-socket go/base_client.go sendHeartBeat), but a larger wait window
	// gives the peer's next heartbeat an extra chance to overwrite any stray
	// timer on slow or jittery links before it fires.
	heartbeatWaitSec = 45
	registerTimeout  = 10 * time.Second
)

// RegisterResult is returned by DialEdge after successful registration.
type RegisterResult struct {
	Session *Session
	NodeID  string
	Alias   string
	Config  *iop.NodeConfigPayload
}

// DialEdge connects to edge, performs the registration handshake, and returns
// a RegisterResult. Call result.Session.SetHandler after creating node.Node.
func DialEdge(ctx context.Context, addr, token string, logger *zap.Logger) (*RegisterResult, error) {
	host, portStr, err := net.SplitHostPort(addr)
	if err != nil {
		return nil, fmt.Errorf("transport: invalid addr %q: %w", addr, err)
	}
	port, err := strconv.Atoi(portStr)
	if err != nil {
		return nil, fmt.Errorf("transport: invalid port %q: %w", portStr, err)
	}

	client, err := toki.DialTcp(ctx, host, port, heartbeatIntervalSec, heartbeatWaitSec, nodeParserMap())
	if err != nil {
		return nil, fmt.Errorf("transport: dial edge %s: %w", addr, err)
	}

	resp, err := toki.SendRequestTyped[*iop.RegisterRequest, *iop.RegisterResponse](
		&client.Communicator,
		&iop.RegisterRequest{Token: token},
		registerTimeout,
	)
	if err != nil {
		_ = client.Close()
		return nil, fmt.Errorf("transport: register: %w", err)
	}
	if !resp.GetAccepted() {
		_ = client.Close()
		return nil, fmt.Errorf("transport: register rejected: %s", resp.GetReason())
	}

	sess := newSession(client, logger, resp.GetNodeId(), resp.GetAlias())
	logger.Info("registered with edge",
		zap.String("node_id", resp.GetNodeId()),
		zap.String("alias", resp.GetAlias()),
	)
	return &RegisterResult{
		Session: sess,
		NodeID:  resp.GetNodeId(),
		Alias:   resp.GetAlias(),
		Config:  resp.GetConfig(),
	}, nil
}