package transport

import (
	"context"
	"net"
	"strconv"
	"sync"
	"sync/atomic"
	"time"

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

	"iop/apps/edge/internal/configrefresh"
	edgenode "iop/apps/edge/internal/node"
	"iop/packages/go/events"
	iop "iop/proto/gen/iop"
)

const (
	heartbeatIntervalSec = 30
	// heartbeatWaitSec mirrors the node side. See the comment in
	// apps/node/internal/transport/client.go for rationale.
	heartbeatWaitSec = 45
)

const configRefreshTimeoutSec = 30

func edgeParserMap() toki.ParserMap {
	return toki.ParserMap{
		toki.TypeNameOf(&iop.RunEvent{}): func(b []byte) (proto.Message, error) {
			m := &iop.RunEvent{}
			return m, proto.Unmarshal(b, m)
		},
		toki.TypeNameOf(&iop.EdgeNodeEvent{}): func(b []byte) (proto.Message, error) {
			m := &iop.EdgeNodeEvent{}
			return m, proto.Unmarshal(b, m)
		},
		toki.TypeNameOf(&iop.RegisterRequest{}): func(b []byte) (proto.Message, error) {
			m := &iop.RegisterRequest{}
			return m, proto.Unmarshal(b, m)
		},
		toki.TypeNameOf(&iop.NodeCommandResponse{}): func(b []byte) (proto.Message, error) {
			m := &iop.NodeCommandResponse{}
			return m, proto.Unmarshal(b, m)
		},
		toki.TypeNameOf(&iop.NodeConfigRefreshResponse{}): func(b []byte) (proto.Message, error) {
			m := &iop.NodeConfigRefreshResponse{}
			return m, proto.Unmarshal(b, m)
		},
	}
}

// Server wraps proto-socket TcpServer and manages node connections.
type Server struct {
	tcp               *toki.TcpServer
	listen            string
	registry          *edgenode.Registry
	nodeStoreMu       sync.RWMutex
	nodeStore         *edgenode.NodeStore
	logger            *zap.Logger
	handlerMu         sync.RWMutex
	onRunEvent        func(*iop.RunEvent)
	onNodeEvent       func(*iop.EdgeNodeEvent)
	stopping          atomic.Bool
	HeartbeatInterval int
	HeartbeatWait     int
}

func NewServer(listen string, registry *edgenode.Registry, nodeStore *edgenode.NodeStore, logger *zap.Logger) (*Server, error) {
	host, portStr, err := net.SplitHostPort(listen)
	if err != nil {
		return nil, err
	}
	port, err := strconv.Atoi(portStr)
	if err != nil {
		return nil, err
	}

	s := &Server{
		listen:            listen,
		registry:          registry,
		nodeStore:         nodeStore,
		logger:            logger,
		HeartbeatInterval: heartbeatIntervalSec,
		HeartbeatWait:     heartbeatWaitSec,
	}
	s.tcp = toki.NewTcpServer(host, port, func(conn net.Conn) *toki.TcpClient {
		return toki.NewTcpClient(conn, s.HeartbeatInterval, s.HeartbeatWait, edgeParserMap())
	})
	s.tcp.OnClientConnected = s.onNodeConnected
	return s, nil
}

func (s *Server) Start(ctx context.Context) error {
	s.stopping.Store(false)
	if err := s.tcp.Start(ctx); err != nil {
		return err
	}
	s.logger.Info("edge listening for nodes", zap.String("addr", s.listen))
	return nil
}

func (s *Server) SetNodeStore(store *edgenode.NodeStore) {
	s.nodeStoreMu.Lock()
	s.nodeStore = store
	s.nodeStoreMu.Unlock()
}

func (s *Server) nodeStoreSnapshot() *edgenode.NodeStore {
	s.nodeStoreMu.RLock()
	defer s.nodeStoreMu.RUnlock()
	return s.nodeStore
}

func (s *Server) Stop() error {
	s.stopping.Store(true)
	return s.tcp.Stop()
}

func (s *Server) SetRunEventHandler(handler func(*iop.RunEvent)) {
	s.handlerMu.Lock()
	s.onRunEvent = handler
	s.handlerMu.Unlock()
}

func (s *Server) SetNodeEventHandler(handler func(*iop.EdgeNodeEvent)) {
	s.handlerMu.Lock()
	s.onNodeEvent = handler
	s.handlerMu.Unlock()
}

func (s *Server) HasRunEventHandler() bool {
	s.handlerMu.RLock()
	defer s.handlerMu.RUnlock()
	return s.onRunEvent != nil
}

func (s *Server) HasNodeEventHandler() bool {
	s.handlerMu.RLock()
	defer s.handlerMu.RUnlock()
	return s.onNodeEvent != nil
}

func (s *Server) onNodeConnected(client *toki.TcpClient) {
	s.logger.Info("node connection established")

	toki.AddListenerTyped[*iop.RunEvent](&client.Communicator, func(e *iop.RunEvent) {
		s.logger.Debug("run event received",
			zap.String("run_id", e.GetRunId()),
			zap.String("type", e.GetType()),
		)
		s.enrichRunEvent(e)
		s.handlerMu.RLock()
		handler := s.onRunEvent
		s.handlerMu.RUnlock()
		if handler != nil {
			handler(e)
		}
	})

	toki.AddRequestListenerTyped[*iop.RegisterRequest, *iop.RegisterResponse](
		&client.Communicator,
		func(req *iop.RegisterRequest) (*iop.RegisterResponse, error) {
			nodeStore := s.nodeStoreSnapshot()
			if nodeStore == nil {
				return &iop.RegisterResponse{Accepted: false, Reason: "edge node store unavailable"}, nil
			}
			rec, ok := nodeStore.FindByToken(req.GetToken())
			if !ok {
				s.logger.Warn("unknown token", zap.String("token_prefix", safePrefix(req.GetToken())))
				s.emitNodeEvent(events.NewEdgeNodeEvent(
					events.SourceEdge,
					events.TypeNodeRegistrationFailed,
					"",
					"",
					events.ReasonUnknownToken,
					map[string]string{
						events.MetadataFailureReason: events.ReasonUnknownToken,
						events.MetadataTokenPrefix:   safePrefix(req.GetToken()),
					},
				))
				return &iop.RegisterResponse{Accepted: false, Reason: "unknown token"}, nil
			}

			cfg, err := edgenode.BuildConfigPayload(rec)
			if err != nil {
				s.logger.Error("build config payload failed",
					zap.String("node_id", rec.ID),
					zap.Error(err),
				)
				return &iop.RegisterResponse{Accepted: false, Reason: "internal config error"}, nil
			}

			entry := &edgenode.NodeEntry{
				NodeID:         rec.ID,
				Alias:          rec.Alias,
				AgentKind:      rec.AgentKind,
				LifecycleState: edgenode.LifecycleConnected,
				Client:         client,
				Index:          rec.Index,
				HasIndex:       true,
			}
			toki.AddListenerTyped[*iop.EdgeNodeEvent](&client.Communicator, func(e *iop.EdgeNodeEvent) {
				if e.NodeId == "" {
					e.NodeId = rec.ID
				}
				if e.Alias == "" {
					e.Alias = rec.Alias
				}
				s.emitNodeEvent(e)
			})
			client.AddDisconnectListener(func(_ *toki.TcpClient) {
				transportInfo := client.DisconnectInfo()
				fields := []zap.Field{zap.String("node_id", rec.ID)}
				fields = append(fields, transportDisconnectFields(transportInfo)...)
				s.logger.Info("node unregistered", fields...)
				reason := events.ReasonTransportClosed
				if transportInfo.Reason == "heartbeat_timeout" {
					reason = events.ReasonHeartbeatTimeout
					s.registry.UpdateLifecycle(rec.ID, edgenode.LifecycleFailed)
				} else if s.stopping.Load() {
					reason = events.ReasonEdgeShutdown
				}

				meta := transportDisconnectMetadata(transportInfo)
				if meta == nil {
					meta = make(map[string]string)
				}
				if reason == events.ReasonHeartbeatTimeout {
					meta[events.MetadataFailureReason] = events.ReasonHeartbeatTimeout
					meta[events.MetadataLifecycleState] = edgenode.LifecycleFailed
				} else {
					meta[events.MetadataLifecycleState] = edgenode.LifecycleFailed
				}
				meta[events.MetadataAgentKind] = rec.AgentKind

				s.emitNodeEvent(events.NewEdgeNodeEvent(
					events.SourceEdge,
					events.TypeNodeDisconnected,
					rec.ID,
					rec.Alias,
					reason,
					meta,
				))
				s.registry.UnregisterIfClient(rec.ID, client)
			})
			if !s.registry.RegisterIfAbsent(entry) {
				reason := "node already connected"
				s.logger.Warn("duplicate registration rejected",
					zap.String("node_id", rec.ID),
					zap.String("agent_kind", rec.AgentKind),
				)
				s.emitNodeEvent(events.NewEdgeNodeEvent(
					events.SourceEdge,
					events.TypeNodeRegistrationFailed,
					rec.ID,
					rec.Alias,
					events.ReasonDuplicateConnection,
					map[string]string{
						events.MetadataFailureReason: events.ReasonDuplicateConnection,
						events.MetadataAgentKind:     rec.AgentKind,
					},
				))
				return &iop.RegisterResponse{Accepted: false, Reason: reason}, nil
			}
			s.logger.Info("node registered",
				zap.String("node_id", rec.ID),
				zap.String("alias", rec.Alias),
			)
			s.emitNodeEvent(events.NewEdgeNodeEvent(
				events.SourceEdge,
				events.TypeNodeConnected,
				rec.ID,
				rec.Alias,
				events.ReasonRegistered,
				map[string]string{
					events.MetadataAgentKind:      rec.AgentKind,
					events.MetadataLifecycleState: edgenode.LifecycleConnected,
				},
			))

			return &iop.RegisterResponse{
				Accepted: true,
				NodeId:   rec.ID,
				Alias:    rec.Alias,
				Config:   cfg,
			}, nil
		},
	)
}

// PushConfigRefresh sends a NodeConfigRefreshRequest to every configured node
// and collects per-node results. Nodes that are not in the live registry or
// have no alive client are recorded as skipped. Connected nodes receive a
// node-specific payload built from the refreshed NodeStore record. Transport
// failures are recorded as failed and do not roll back the local config commit.
func (s *Server) PushConfigRefresh(ctx context.Context, req *iop.NodeConfigRefreshRequest) []configrefresh.NodeResult {
	ns := s.nodeStoreSnapshot()
	var records []*edgenode.NodeRecord
	if ns != nil {
		records = ns.All()
	}
	results := make([]configrefresh.NodeResult, 0, len(records))
	timeout := time.Duration(configRefreshTimeoutSec) * time.Second
	for _, rec := range records {
		nr := configrefresh.NodeResult{
			NodeID: rec.ID,
			Alias:  rec.Alias,
		}
		entry, ok := s.registry.Get(rec.ID)
		if !ok || entry.Client == nil || !entry.Client.IsAlive() {
			nr.Status = "skipped"
			results = append(results, nr)
			continue
		}
		nodeReq := &iop.NodeConfigRefreshRequest{
			RequestId:    req.GetRequestId(),
			ChangedPaths: req.GetChangedPaths(),
		}
		if payload, err := edgenode.BuildConfigPayload(rec); err == nil {
			nodeReq.Config = payload
		}
		resp, err := toki.SendRequestTyped[*iop.NodeConfigRefreshRequest, *iop.NodeConfigRefreshResponse](
			&entry.Client.Communicator,
			nodeReq,
			timeout,
		)
		if err != nil {
			nr.Status = "failed"
			nr.Error = err.Error()
			results = append(results, nr)
			continue
		}
		switch resp.GetStatus() {
		case iop.NodeConfigRefreshStatus_NODE_CONFIG_REFRESH_STATUS_APPLIED:
			nr.Status = "applied"
		case iop.NodeConfigRefreshStatus_NODE_CONFIG_REFRESH_STATUS_RESTART_REQUIRED:
			nr.Status = "restart_required"
			nr.RestartRequiredPaths = resp.GetRestartRequiredPaths()
		case iop.NodeConfigRefreshStatus_NODE_CONFIG_REFRESH_STATUS_FAILED:
			nr.Status = "failed"
			nr.Error = resp.GetError()
		case iop.NodeConfigRefreshStatus_NODE_CONFIG_REFRESH_STATUS_SKIPPED:
			nr.Status = "skipped"
		default:
			nr.Status = "skipped"
		}
		results = append(results, nr)
	}
	return results
}

func (s *Server) enrichRunEvent(event *iop.RunEvent) {
	if event == nil || event.GetNodeAlias() != "" || s.registry == nil {
		return
	}
	nodeID := event.GetNodeId()
	if nodeID == "" {
		return
	}
	if entry, ok := s.registry.Get(nodeID); ok && entry.Alias != "" {
		event.NodeAlias = entry.Alias
	}
}

func (s *Server) emitNodeEvent(event *iop.EdgeNodeEvent) {
	s.logger.Debug("node event received",
		zap.String("node_id", event.GetNodeId()),
		zap.String("type", event.GetType()),
		zap.String("reason", event.GetReason()),
	)
	s.handlerMu.RLock()
	handler := s.onNodeEvent
	s.handlerMu.RUnlock()
	if handler != nil {
		handler(event)
	}
}

func safePrefix(s string) string {
	if len(s) > 8 {
		return s[:8] + "..."
	}
	return s
}

func transportDisconnectMetadata(info toki.DisconnectInfo) map[string]string {
	metadata := make(map[string]string, 2)
	if info.Reason != "" {
		metadata[events.MetadataTransportCloseReason] = info.Reason
	}
	if info.Error != "" {
		metadata[events.MetadataTransportCloseError] = info.Error
	}
	if len(metadata) == 0 {
		return nil
	}
	return metadata
}

func transportDisconnectFields(info toki.DisconnectInfo) []zap.Field {
	fields := make([]zap.Field, 0, 2)
	if info.Reason != "" {
		fields = append(fields, zap.String("transport_close_reason", info.Reason))
	}
	if info.Error != "" {
		fields = append(fields, zap.String("transport_close_error", info.Error))
	}
	return fields
}
