package transport

import (
	"context"
	"sync"

	toki "git.toki-labs.com/toki/common-proto-socket/go"
	"go.uber.org/zap"
	"google.golang.org/protobuf/proto"

	iop "iop/proto/gen/iop"
)

// Handler processes IOP messages received from edge.
type Handler interface {
	OnRunRequest(ctx context.Context, sess *Session, req *iop.RunRequest) error
	OnCapabilityRequest(ctx context.Context, sess *Session) (*iop.CapabilityResponse, error)
	OnCancel(ctx context.Context, sess *Session, req *iop.CancelRequest) error
}

// Session represents the node's persistent connection to edge.
type Session struct {
	client    *toki.TcpClient
	nodeID    string
	logger    *zap.Logger
	cancelFns sync.Map
}

func newSession(client *toki.TcpClient, handler Handler, nodeID string, logger *zap.Logger) *Session {
	s := &Session{client: client, nodeID: nodeID, logger: logger}

	toki.AddListenerTyped[*iop.RunRequest](&client.Communicator, func(req *iop.RunRequest) {
		go func() {
			if err := handler.OnRunRequest(context.Background(), s, req); err != nil {
				logger.Warn("run request error",
					zap.String("run_id", req.GetRunId()),
					zap.Error(err),
				)
			}
		}()
	})

	toki.AddRequestListenerTyped[*iop.CapabilityRequest, *iop.CapabilityResponse](
		&client.Communicator,
		func(_ *iop.CapabilityRequest) (*iop.CapabilityResponse, error) {
			return handler.OnCapabilityRequest(context.Background(), s)
		},
	)

	toki.AddListenerTyped[*iop.CancelRequest](&client.Communicator, func(req *iop.CancelRequest) {
		if err := handler.OnCancel(context.Background(), s, req); err != nil {
			logger.Warn("cancel error",
				zap.String("run_id", req.GetRunId()),
				zap.Error(err),
			)
		}
	})

	client.AddDisconnectListener(func(_ *toki.TcpClient) {
		logger.Info("disconnected from edge", zap.String("node_id", nodeID))
	})

	return s
}

// Send transmits a proto message to edge.
func (s *Session) Send(m proto.Message) error {
	return s.client.Send(m)
}

// IsAlive reports whether the connection is active.
func (s *Session) IsAlive() bool { return s.client.IsAlive() }

// Close terminates the connection to edge.
func (s *Session) Close() error { return s.client.Close() }

// RegisterCancel stores a cancel function for a running execution.
func (s *Session) RegisterCancel(runID string, cancel context.CancelFunc) {
	s.cancelFns.Store(runID, cancel)
}

// DeregisterCancel removes a run's cancel function after completion.
func (s *Session) DeregisterCancel(runID string) {
	s.cancelFns.Delete(runID)
}

// CancelRun invokes the cancel function for a run if one is registered.
func (s *Session) CancelRun(runID string) {
	if v, ok := s.cancelFns.Load(runID); ok {
		v.(context.CancelFunc)()
	}
}