// Licensed to Apache Software Foundation (ASF) under one or more contributor
// license agreements. See the NOTICE file distributed with
// this work for additional information regarding copyright
// ownership. Apache Software Foundation (ASF) licenses this file to you under
// the Apache License, Version 2.0 (the "License"); you may
// not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

// Package proxy provides a client for communicating with the FODC Proxy.
package proxy

import (
	"context"
	"errors"
	"fmt"
	"io"
	"strings"
	"sync"
	"time"

	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/metadata"
	"google.golang.org/grpc/status"
	"google.golang.org/protobuf/types/known/timestamppb"

	databasev1 "github.com/apache/skywalking-banyandb/api/proto/banyandb/database/v1"
	fodcv1 "github.com/apache/skywalking-banyandb/api/proto/banyandb/fodc/v1"
	"github.com/apache/skywalking-banyandb/fodc/agent/internal/cluster"
	"github.com/apache/skywalking-banyandb/fodc/agent/internal/crashcollector"
	"github.com/apache/skywalking-banyandb/fodc/agent/internal/flightrecorder"
	"github.com/apache/skywalking-banyandb/fodc/agent/internal/lifecycle"
	"github.com/apache/skywalking-banyandb/fodc/agent/internal/metrics"
	"github.com/apache/skywalking-banyandb/fodc/agent/internal/pressureprofiler"
	"github.com/apache/skywalking-banyandb/pkg/logger"
	"github.com/apache/skywalking-banyandb/pkg/panicdiag"
	"github.com/apache/skywalking-banyandb/pkg/run"
)

// defaultReconnectRetryInterval is the fallback delay scheduleReconnect uses when the client was
// created without a positive reconnect interval. In practice c.reconnectInterval is always set from
// the --reconnect-interval flag; this only guards a zero value.
const defaultReconnectRetryInterval = 10 * time.Second

// MetricsRequestFilter defines filters for metrics requests.
type MetricsRequestFilter struct {
	StartTime *time.Time
	EndTime   *time.Time
}

// Client manages connection and communication with the FODC Proxy.
type Client struct {
	logger             *logger.Logger
	heartbeatTicker    *time.Ticker
	flightRecorder     *flightrecorder.FlightRecorder
	connManager        *connManager
	clusterCollector   *cluster.Collector
	lifecycleCollector *lifecycle.Collector
	collectionLister   crashcollector.CollectionLister
	stopCh             chan struct{}
	reconnectCh        chan struct{}
	labels             map[string]string
	metricsStream      fodcv1.FODCService_StreamMetricsClient
	registrationStream fodcv1.FODCService_RegisterAgentClient
	clusterStateStream fodcv1.FODCService_StreamClusterTopologyClient
	lifecycleStream    fodcv1.FODCService_StreamLifecycleClient
	crashStream        fodcv1.FODCService_StreamCrashDiagnosticsClient
	pressureStream     fodcv1.FODCService_StreamPressureProfilesClient
	profileSource      ProfileSource
	client             fodcv1.FODCServiceClient
	reconnectFn        func(context.Context) // overridable in tests; nil means use reconnect

	proxyAddr      string
	nodeRole       string
	podName        string
	agentID        string
	containerNames []string

	heartbeatInterval time.Duration
	reconnectInterval time.Duration
	streamsMu         sync.RWMutex   // Protects streams and stream operations
	heartbeatWg       sync.WaitGroup // Tracks heartbeat goroutine
	disconnected      bool
}

// NewClient creates a new Client instance.
func NewClient(
	proxyAddr string,
	nodeRole string,
	podName string,
	containerNames []string,
	labels map[string]string,
	heartbeatInterval time.Duration,
	reconnectInterval time.Duration,
	flightRecorder *flightrecorder.FlightRecorder,
	clusterCollector *cluster.Collector,
	lifecycleCollector *lifecycle.Collector,
	collectionLister crashcollector.CollectionLister,
	profileSource ProfileSource,
	logger *logger.Logger,
) *Client {
	connMgr := newConnManager(proxyAddr, reconnectInterval, logger)
	client := &Client{
		connManager:        connMgr,
		proxyAddr:          proxyAddr,
		nodeRole:           nodeRole,
		podName:            podName,
		containerNames:     containerNames,
		labels:             labels,
		heartbeatInterval:  heartbeatInterval,
		reconnectInterval:  reconnectInterval,
		flightRecorder:     flightRecorder,
		clusterCollector:   clusterCollector,
		lifecycleCollector: lifecycleCollector,
		collectionLister:   collectionLister,
		profileSource:      profileSource,
		logger:             logger,
		stopCh:             make(chan struct{}),
		reconnectCh:        make(chan struct{}, 1),
	}
	connMgr.setHeartbeatChecker(client.SendHeartbeat)
	return client
}

// StartConnManager is useful for tests or scenarios where you want to manually control connection lifecycle.
func (c *Client) StartConnManager(ctx context.Context) {
	c.connManager.start(ctx)
}

// Connect establishes a gRPC connection to Proxy.
func (c *Client) Connect(ctx context.Context) error {
	resultCh := c.connManager.RequestConnect(ctx)
	result := <-resultCh
	if result.err != nil {
		return result.err
	}

	c.streamsMu.Lock()
	wasDisconnected := c.disconnected
	c.streamsMu.Unlock()

	if wasDisconnected {
		if !waitWaitGroupWithTimeout(&c.heartbeatWg, heartbeatWaitTimeout) {
			c.logger.Warn().Msg("Heartbeat goroutine did not exit in time during Connect, proceeding anyway")
		}
	}

	c.streamsMu.Lock()
	c.client = fodcv1.NewFODCServiceClient(result.conn)
	if wasDisconnected {
		c.disconnected = false
		c.stopCh = make(chan struct{})
	}
	c.streamsMu.Unlock()

	return nil
}

// StartRegistrationStream establishes bi-directional registration stream with Proxy.
func (c *Client) StartRegistrationStream(ctx context.Context) error {
	c.streamsMu.Lock()
	if c.client == nil {
		c.streamsMu.Unlock()
		return fmt.Errorf("client not connected, call Connect() first")
	}
	client := c.client
	c.streamsMu.Unlock()

	stream, streamErr := client.RegisterAgent(ctx)
	if streamErr != nil {
		return fmt.Errorf("failed to create registration stream: %w", streamErr)
	}

	req := &fodcv1.RegisterAgentRequest{
		NodeRole:       c.nodeRole,
		Labels:         c.labels,
		PodName:        c.podName,
		ContainerNames: c.containerNames,
	}

	if sendErr := stream.Send(req); sendErr != nil {
		return fmt.Errorf("failed to send registration request: %w", sendErr)
	}

	resp, recvErr := stream.Recv()
	if recvErr != nil {
		return fmt.Errorf("failed to receive registration response: %w", recvErr)
	}

	if !resp.Success {
		return fmt.Errorf("registration failed: %s", resp.Message)
	}

	if resp.AgentId == "" {
		return fmt.Errorf("registration response missing agent ID")
	}

	c.streamsMu.Lock()
	c.registrationStream = stream
	c.agentID = resp.AgentId
	if resp.HeartbeatIntervalSeconds > 0 {
		c.heartbeatInterval = time.Duration(resp.HeartbeatIntervalSeconds) * time.Second
	}
	c.streamsMu.Unlock()

	c.logger.Info().
		Str("proxy_addr", c.proxyAddr).
		Str("agent_id", resp.AgentId).
		Str("node_role", c.nodeRole).
		Dur("heartbeat_interval", c.heartbeatInterval).
		Msg("Agent registered with Proxy")

	c.startHeartbeat(ctx)

	run.Go(ctx, "fodc.agent.proxy.registration-stream", c.logger, func(ctx context.Context) {
		c.handleRegistrationStream(ctx, stream)
	})

	return nil
}

// StartMetricsStream establishes bi-directional metrics stream with Proxy.
func (c *Client) StartMetricsStream(ctx context.Context) error {
	c.streamsMu.Lock()
	if c.client == nil {
		c.streamsMu.Unlock()
		return fmt.Errorf("client not connected, call Connect() first")
	}
	client := c.client
	agentID := c.agentID
	c.streamsMu.Unlock()

	if agentID == "" {
		return fmt.Errorf("agent ID not available, register agent first")
	}

	md := metadata.New(map[string]string{"agent_id": agentID})
	ctxWithMetadata := metadata.NewOutgoingContext(ctx, md)

	stream, streamErr := client.StreamMetrics(ctxWithMetadata)
	if streamErr != nil {
		return fmt.Errorf("failed to create metrics stream: %w", streamErr)
	}

	c.streamsMu.Lock()
	c.metricsStream = stream
	c.streamsMu.Unlock()

	run.Go(ctx, "fodc.agent.proxy.metrics-stream", c.logger, func(ctx context.Context) {
		c.handleMetricsStream(ctx, stream)
	})

	c.logger.Info().
		Str("agent_id", agentID).
		Msg("Metrics stream established with Proxy")

	return nil
}

// StartClusterStateStream establishes bi-directional cluster state stream with Proxy.
func (c *Client) StartClusterStateStream(ctx context.Context) error {
	c.streamsMu.Lock()
	if c.client == nil {
		c.streamsMu.Unlock()
		return fmt.Errorf("client not connected, call Connect() first")
	}
	client := c.client
	agentID := c.agentID
	c.streamsMu.Unlock()

	if agentID == "" {
		return fmt.Errorf("agent ID not available, register agent first")
	}

	md := metadata.New(map[string]string{"agent_id": agentID})
	ctxWithMetadata := metadata.NewOutgoingContext(ctx, md)

	stream, streamErr := client.StreamClusterTopology(ctxWithMetadata)
	if streamErr != nil {
		return fmt.Errorf("failed to create cluster state stream: %w", streamErr)
	}

	c.streamsMu.Lock()
	c.clusterStateStream = stream
	c.streamsMu.Unlock()

	run.Go(ctx, "fodc.agent.proxy.cluster-state-stream", c.logger, func(ctx context.Context) {
		c.handleClusterStateStream(ctx, stream)
	})

	c.logger.Info().
		Str("agent_id", agentID).
		Msg("Cluster state stream established with Proxy")

	return nil
}

// sendClusterTopology sends cluster topology to Proxy.
func (c *Client) sendClusterTopology() error {
	c.streamsMu.RLock()
	if c.disconnected || c.clusterStateStream == nil {
		c.streamsMu.RUnlock()
		return fmt.Errorf("cluster state stream not established")
	}
	clusterStateStream := c.clusterStateStream
	collector := c.clusterCollector
	c.streamsMu.RUnlock()
	var topology cluster.TopologyMap
	if collector == nil {
		c.logger.Debug().Msg("Cluster collector not available, sending empty topology")
		topology = cluster.TopologyMap{
			Nodes: make([]*databasev1.Node, 0),
			Calls: make([]*fodcv1.Call, 0),
		}
	} else {
		topology = collector.GetClusterTopology()
		if len(topology.Nodes) == 0 && len(topology.Calls) == 0 {
			c.logger.Debug().Msg("No cluster topology available to send")
		}
	}
	req := &fodcv1.StreamClusterTopologyRequest{
		Topology: &fodcv1.Topology{
			Nodes: topology.Nodes,
			Calls: topology.Calls,
		},
		Timestamp: timestamppb.Now(),
	}
	c.streamsMu.Lock()
	if c.clusterStateStream != clusterStateStream {
		c.streamsMu.Unlock()
		return fmt.Errorf("cluster state stream changed during send")
	}
	sendErr := clusterStateStream.Send(req)
	c.streamsMu.Unlock()
	if sendErr != nil {
		return fmt.Errorf("failed to send cluster topology: %w", sendErr)
	}
	c.logger.Info().
		Int("nodes_count", len(topology.Nodes)).
		Int("calls_count", len(topology.Calls)).
		Msg("Successfully sent cluster topology to proxy")
	return nil
}

// StartLifecycleStream establishes bi-directional lifecycle stream with Proxy.
func (c *Client) StartLifecycleStream(ctx context.Context) error {
	c.streamsMu.Lock()
	if c.client == nil {
		c.streamsMu.Unlock()
		return fmt.Errorf("client not connected, call Connect() first")
	}
	client := c.client
	agentID := c.agentID
	c.streamsMu.Unlock()

	if agentID == "" {
		return fmt.Errorf("agent ID not available, register agent first")
	}

	md := metadata.New(map[string]string{"agent_id": agentID})
	ctxWithMetadata := metadata.NewOutgoingContext(ctx, md)

	stream, streamErr := client.StreamLifecycle(ctxWithMetadata)
	if streamErr != nil {
		return fmt.Errorf("failed to create lifecycle stream: %w", streamErr)
	}

	c.streamsMu.Lock()
	c.lifecycleStream = stream
	c.streamsMu.Unlock()

	run.Go(ctx, "fodc.agent.proxy.lifecycle-stream", c.logger, func(ctx context.Context) {
		c.handleLifecycleStream(ctx, stream)
	})

	c.logger.Info().
		Str("agent_id", agentID).
		Msg("Lifecycle stream established with Proxy")

	return nil
}

// sendLifecycleData sends lifecycle data to Proxy.
func (c *Client) sendLifecycleData(ctx context.Context) error {
	c.streamsMu.RLock()
	if c.disconnected || c.lifecycleStream == nil {
		c.streamsMu.RUnlock()
		return fmt.Errorf("lifecycle stream not established")
	}
	lifecycleStream := c.lifecycleStream
	collector := c.lifecycleCollector
	c.streamsMu.RUnlock()
	if collector == nil {
		c.logger.Warn().Msg("Lifecycle collector is nil, skipping send")
		return nil
	}
	data, collectErr := collector.Collect(ctx)
	if collectErr != nil {
		return fmt.Errorf("failed to collect lifecycle data: %w", collectErr)
	}
	if len(data.Reports) == 0 && len(data.Groups) == 0 {
		c.logger.Debug().Msg("No lifecycle data available to send")
	}
	req := &fodcv1.StreamLifecycleRequest{
		Timestamp:     timestamppb.Now(),
		PodName:       c.podName,
		LifecycleData: data,
	}
	c.streamsMu.Lock()
	if c.lifecycleStream != lifecycleStream {
		c.streamsMu.Unlock()
		return fmt.Errorf("lifecycle stream changed during send")
	}
	sendErr := lifecycleStream.Send(req)
	c.streamsMu.Unlock()
	if sendErr != nil {
		return fmt.Errorf("failed to send lifecycle data: %w", sendErr)
	}
	c.logger.Info().
		Int("reports_count", len(data.Reports)).
		Int("groups_count", len(data.Groups)).
		Msg("Successfully sent lifecycle data to proxy")
	return nil
}

// StartCrashStream establishes bi-directional crash diagnostics stream with Proxy.
func (c *Client) StartCrashStream(ctx context.Context) error {
	c.streamsMu.Lock()
	if c.client == nil {
		c.streamsMu.Unlock()
		return fmt.Errorf("client not connected, call Connect() first")
	}
	client := c.client
	agentID := c.agentID
	c.streamsMu.Unlock()

	if agentID == "" {
		return fmt.Errorf("agent ID not available, register agent first")
	}

	md := metadata.New(map[string]string{"agent_id": agentID})
	ctxWithMetadata := metadata.NewOutgoingContext(ctx, md)

	stream, streamErr := client.StreamCrashDiagnostics(ctxWithMetadata)
	if streamErr != nil {
		return fmt.Errorf("failed to create crash diagnostics stream: %w", streamErr)
	}

	c.streamsMu.Lock()
	c.crashStream = stream
	c.streamsMu.Unlock()

	run.Go(ctx, "fodc.agent.proxy.crash-stream", c.logger, func(ctx context.Context) {
		c.handleCrashStream(ctx, stream)
	})

	c.logger.Info().
		Str("agent_id", agentID).
		Msg("Crash diagnostics stream established with Proxy")

	return nil
}

// sendCrashCollections sends all current crash collections to Proxy.
func (c *Client) sendCrashCollections() error {
	c.streamsMu.RLock()
	if c.disconnected || c.crashStream == nil {
		c.streamsMu.RUnlock()
		return fmt.Errorf("crash diagnostics stream not established")
	}
	crashStream := c.crashStream
	lister := c.collectionLister
	c.streamsMu.RUnlock()

	if lister == nil {
		c.logger.Warn().Msg("Crash collection lister is nil, skipping send")
		return nil
	}

	records := lister.ListCollections()
	if len(records) == 0 {
		c.logger.Debug().Msg("No crash collections available to send")
		return nil
	}

	for _, record := range records {
		req := &fodcv1.StreamCrashDiagnosticsRequest{
			FetchedAt:      timestamppb.New(record.FetchedAt),
			SourceEndpoint: record.SourceEndpoint,
			ArtifactDir:    record.Collection.ArtifactDir,
			Files:          record.Collection.Files,
		}
		if record.Collection.Record != nil {
			req.PanicRecord = &fodcv1.CrashPanicRecord{
				OccurredAt:      timestamppb.New(record.Collection.Record.OccurredAt),
				Component:       record.Collection.Record.Component,
				PanicValue:      record.Collection.Record.PanicValue,
				Recovered:       record.Collection.Record.Recovered,
				GoroutineStack:  record.Collection.Record.GoroutineStack,
				Breadcrumbs:     buildCrashBreadcrumbs(record.Collection.Record.Breadcrumbs),
				ProcessMetadata: cloneStringMap(record.Collection.Record.ProcessMetadata),
			}
		}
		c.streamsMu.Lock()
		if c.crashStream != crashStream {
			c.streamsMu.Unlock()
			return fmt.Errorf("crash diagnostics stream changed during send")
		}
		sendErr := crashStream.Send(req)
		c.streamsMu.Unlock()
		if sendErr != nil {
			return fmt.Errorf("failed to send crash collection: %w", sendErr)
		}
	}

	c.logger.Info().
		Int("records_count", len(records)).
		Msg("Successfully sent crash collections to proxy")
	return nil
}

func buildCrashBreadcrumbs(breadcrumbs []panicdiag.Breadcrumb) []*fodcv1.CrashBreadcrumb {
	if len(breadcrumbs) == 0 {
		return nil
	}
	result := make([]*fodcv1.CrashBreadcrumb, 0, len(breadcrumbs))
	for _, breadcrumb := range breadcrumbs {
		result = append(result, &fodcv1.CrashBreadcrumb{
			Time:      timestamppb.New(breadcrumb.Time),
			Stage:     breadcrumb.Stage,
			Component: breadcrumb.Component,
			Fields:    cloneBreadcrumbFields(breadcrumb.Fields),
		})
	}
	return result
}

func cloneBreadcrumbFields(fields map[string]string) map[string]string {
	return cloneStringMap(fields)
}

func cloneStringMap(src map[string]string) map[string]string {
	if len(src) == 0 {
		return nil
	}
	cloned := make(map[string]string, len(src))
	for key, value := range src {
		cloned[key] = value
	}
	return cloned
}

// handleCrashStream handles the crash diagnostics stream.
func (c *Client) handleCrashStream(ctx context.Context, stream fodcv1.FODCService_StreamCrashDiagnosticsClient) {
	c.streamsMu.RLock()
	stopCh := c.stopCh
	c.streamsMu.RUnlock()
	for {
		select {
		case <-ctx.Done():
			return
		case <-stopCh:
			return
		default:
		}
		resp, recvErr := stream.Recv()
		if errors.Is(recvErr, io.EOF) {
			c.logger.Warn().Msg("Crash diagnostics stream closed by Proxy, reconnecting...")
			run.Go(ctx, "fodc.agent.proxy.reconnect", c.logger, func(ctx context.Context) {
				c.reconnect(ctx)
			})
			return
		}
		if recvErr != nil {
			c.streamsMu.RLock()
			disconnected := c.disconnected
			c.streamsMu.RUnlock()
			if disconnected {
				c.logger.Debug().Err(recvErr).Msg("Crash diagnostics stream closed")
				return
			}
			c.logger.Error().Err(recvErr).Msg("Error receiving from crash diagnostics stream, reconnecting...")
			run.Go(ctx, "fodc.agent.proxy.reconnect", c.logger, func(ctx context.Context) {
				c.reconnect(ctx)
			})
			return
		}
		if resp != nil && resp.RequestDiagnostics {
			c.logger.Debug().Msg("Received crash diagnostics request from proxy")
			if sendErr := c.sendCrashCollections(); sendErr != nil {
				c.logger.Error().Err(sendErr).Msg("Failed to send crash collections to proxy")
			}
		}
	}
}

// pressureChunkSize bounds each pprof download chunk so a single gRPC message stays
// far below the server's max-message-size regardless of profile size.
const pressureChunkSize = 1 << 20 // 1MB

// ProfileSource is the read side of the pressure profiler the proxy client serves:
// it lists capture-event metadata and opens a profile file for chunked download.
type ProfileSource interface {
	ListProfileRecords() []pressureprofiler.ProfileRecord
	OpenProfile(path string) (io.ReadCloser, error)
}

// StartPressureProfilesStream establishes the bi-directional pressure-profiles stream with Proxy.
func (c *Client) StartPressureProfilesStream(ctx context.Context) error {
	c.streamsMu.Lock()
	if c.client == nil {
		c.streamsMu.Unlock()
		return fmt.Errorf("client not connected, call Connect() first")
	}
	client := c.client
	agentID := c.agentID
	c.streamsMu.Unlock()

	if agentID == "" {
		return fmt.Errorf("agent ID not available, register agent first")
	}

	md := metadata.New(map[string]string{"agent_id": agentID})
	ctxWithMetadata := metadata.NewOutgoingContext(ctx, md)

	stream, streamErr := client.StreamPressureProfiles(ctxWithMetadata)
	if streamErr != nil {
		return fmt.Errorf("failed to create pressure profiles stream: %w", streamErr)
	}

	c.streamsMu.Lock()
	c.pressureStream = stream
	c.streamsMu.Unlock()

	run.Go(ctx, "fodc.agent.proxy.pressure-stream", c.logger, func(ctx context.Context) {
		c.handlePressureProfilesStream(ctx, stream)
	})

	c.logger.Info().Str("agent_id", agentID).Msg("Pressure profiles stream established with Proxy")
	return nil
}

// handlePressureProfilesStream serves proxy commands on the pressure-profiles stream:
// list metadata, or stream one profile's bytes.
func (c *Client) handlePressureProfilesStream(ctx context.Context, stream fodcv1.FODCService_StreamPressureProfilesClient) {
	c.streamsMu.RLock()
	stopCh := c.stopCh
	c.streamsMu.RUnlock()
	for {
		select {
		case <-ctx.Done():
			return
		case <-stopCh:
			return
		default:
		}
		resp, recvErr := stream.Recv()
		if errors.Is(recvErr, io.EOF) {
			c.logger.Warn().Msg("Pressure profiles stream closed by Proxy, reconnecting...")
			run.Go(ctx, "fodc.agent.proxy.reconnect", c.logger, func(ctx context.Context) {
				c.reconnect(ctx)
			})
			return
		}
		if recvErr != nil {
			c.streamsMu.RLock()
			disconnected := c.disconnected
			c.streamsMu.RUnlock()
			if disconnected {
				c.logger.Debug().Err(recvErr).Msg("Pressure profiles stream closed")
				return
			}
			c.logger.Error().Err(recvErr).Msg("Error receiving from pressure profiles stream, reconnecting...")
			run.Go(ctx, "fodc.agent.proxy.reconnect", c.logger, func(ctx context.Context) {
				c.reconnect(ctx)
			})
			return
		}
		if resp == nil {
			continue
		}
		switch cmd := resp.Command.(type) {
		case *fodcv1.StreamPressureProfilesResponse_ListProfiles:
			if sendErr := c.sendProfileRecords(stream, cmd.ListProfiles.GetRequestId()); sendErr != nil {
				c.logger.Error().Err(sendErr).Msg("Failed to send pressure profile records to proxy")
			}
		case *fodcv1.StreamPressureProfilesResponse_FetchProfile:
			if sendErr := c.sendProfileChunks(stream, cmd.FetchProfile); sendErr != nil {
				c.logger.Error().Err(sendErr).Msg("Failed to send pressure profile chunks to proxy")
			}
		}
	}
}

// sendProfileRecords answers a list_profiles command with one metadata record per complete event.
func (c *Client) sendProfileRecords(stream fodcv1.FODCService_StreamPressureProfilesClient, requestID string) error {
	c.streamsMu.RLock()
	source := c.profileSource
	disconnected := c.disconnected
	current := c.pressureStream
	c.streamsMu.RUnlock()
	if disconnected || current != stream {
		return fmt.Errorf("pressure profiles stream not established")
	}

	if source != nil {
		for _, record := range source.ListProfileRecords() {
			req := &fodcv1.StreamPressureProfilesRequest{
				Payload: &fodcv1.StreamPressureProfilesRequest_Record{Record: toProtoRecord(record)},
			}
			if sendErr := c.sendPressureRequest(stream, req); sendErr != nil {
				return sendErr
			}
		}
	}

	// Always signal completion (including the zero-record / nil-source case) so the proxy
	// returns as soon as this agent is done instead of waiting out its timeout.
	done := &fodcv1.StreamPressureProfilesRequest{
		Payload: &fodcv1.StreamPressureProfilesRequest_ListComplete{
			ListComplete: &fodcv1.ListComplete{RequestId: requestID},
		},
	}
	return c.sendPressureRequest(stream, done)
}

// sendProfileChunks answers a fetch_profile command by streaming one profile's bytes
// in bounded chunks; a non-empty chunk error reports a file that cannot be served.
func (c *Client) sendProfileChunks(stream fodcv1.FODCService_StreamPressureProfilesClient, fetch *fodcv1.FetchPressureProfile) error {
	c.streamsMu.RLock()
	source := c.profileSource
	c.streamsMu.RUnlock()
	if source == nil {
		return c.sendPressureRequest(stream, chunkError(fetch, "pressure profiler not available"))
	}

	file, openErr := source.OpenProfile(fetch.Filepath)
	if openErr != nil {
		return c.sendPressureRequest(stream, chunkError(fetch, openErr.Error()))
	}
	defer func() { _ = file.Close() }()

	buf := make([]byte, pressureChunkSize)
	for {
		n, readErr := file.Read(buf)
		isEOF := errors.Is(readErr, io.EOF)
		if n > 0 || isEOF {
			chunk := &fodcv1.PressureProfileChunk{
				RequestId: fetch.RequestId,
				ProfileId: fetch.ProfileId,
				Type:      fetch.Type,
				Data:      buf[:n],
				Last:      isEOF,
			}
			req := &fodcv1.StreamPressureProfilesRequest{
				Payload: &fodcv1.StreamPressureProfilesRequest_Chunk{Chunk: chunk},
			}
			if sendErr := c.sendPressureRequest(stream, req); sendErr != nil {
				return sendErr
			}
		}
		if isEOF {
			return nil
		}
		if readErr != nil {
			return c.sendPressureRequest(stream, chunkError(fetch, readErr.Error()))
		}
	}
}

// sendPressureRequest sends one message on the pressure stream, guarding against a
// stream that has been swapped out by a concurrent reconnect.
func (c *Client) sendPressureRequest(stream fodcv1.FODCService_StreamPressureProfilesClient, req *fodcv1.StreamPressureProfilesRequest) error {
	c.streamsMu.Lock()
	defer c.streamsMu.Unlock()
	if c.pressureStream != stream {
		return fmt.Errorf("pressure profiles stream changed during send")
	}
	if sendErr := stream.Send(req); sendErr != nil {
		return fmt.Errorf("failed to send pressure profiles message: %w", sendErr)
	}
	return nil
}

func chunkError(fetch *fodcv1.FetchPressureProfile, msg string) *fodcv1.StreamPressureProfilesRequest {
	return &fodcv1.StreamPressureProfilesRequest{
		Payload: &fodcv1.StreamPressureProfilesRequest_Chunk{
			Chunk: &fodcv1.PressureProfileChunk{
				RequestId: fetch.RequestId,
				ProfileId: fetch.ProfileId,
				Type:      fetch.Type,
				Error:     msg,
				Last:      true,
			},
		},
	}
}

func toProtoRecord(record pressureprofiler.ProfileRecord) *fodcv1.PressureProfileRecord {
	profiles := make([]*fodcv1.PressureProfileInfo, 0, len(record.Profiles))
	for _, p := range record.Profiles {
		profiles = append(profiles, &fodcv1.PressureProfileInfo{
			Type:      p.Type,
			Filename:  p.Filename,
			Filepath:  p.Filepath,
			Format:    p.Format,
			SizeBytes: p.SizeBytes,
		})
	}
	return &fodcv1.PressureProfileRecord{
		ProfileId:        record.ProfileID,
		CapturedAt:       timestamppb.New(record.CapturedAt),
		SourceEndpoint:   record.SourceEndpoint,
		RssBytes:         record.RSSBytes,
		CgroupLimitBytes: record.CgroupLimitBytes,
		TriggerPercent:   record.TriggerPercent,
		ThresholdBytes:   record.ThresholdBytes,
		Profiles:         profiles,
	}
}

// handleLifecycleStream handles the lifecycle stream.
func (c *Client) handleLifecycleStream(ctx context.Context, stream fodcv1.FODCService_StreamLifecycleClient) {
	c.streamsMu.RLock()
	stopCh := c.stopCh
	c.streamsMu.RUnlock()
	for {
		select {
		case <-ctx.Done():
			return
		case <-stopCh:
			return
		default:
		}
		resp, recvErr := stream.Recv()
		if errors.Is(recvErr, io.EOF) {
			c.logger.Warn().Msg("Lifecycle stream closed by Proxy, reconnecting...")
			run.Go(ctx, "fodc.agent.proxy.reconnect", c.logger, func(ctx context.Context) {
				c.reconnect(ctx)
			})
			return
		}
		if recvErr != nil {
			c.streamsMu.RLock()
			disconnected := c.disconnected
			c.streamsMu.RUnlock()
			if disconnected {
				c.logger.Debug().Err(recvErr).Msg("Lifecycle stream closed")
				return
			}
			c.logger.Error().Err(recvErr).Msg("Error receiving from lifecycle stream, reconnecting...")
			run.Go(ctx, "fodc.agent.proxy.reconnect", c.logger, func(ctx context.Context) {
				c.reconnect(ctx)
			})
			return
		}
		if resp != nil && resp.RequestLifecycle {
			c.logger.Debug().Msg("Received lifecycle data request from proxy")
			if sendErr := c.sendLifecycleData(ctx); sendErr != nil {
				c.logger.Error().Err(sendErr).Msg("Failed to send lifecycle data to proxy")
			}
		}
	}
}

// handleClusterStateStream handles the cluster state stream.
func (c *Client) handleClusterStateStream(ctx context.Context, stream fodcv1.FODCService_StreamClusterTopologyClient) {
	c.streamsMu.RLock()
	stopCh := c.stopCh
	c.streamsMu.RUnlock()
	for {
		select {
		case <-ctx.Done():
			return
		case <-stopCh:
			return
		default:
		}
		resp, recvErr := stream.Recv()
		if errors.Is(recvErr, io.EOF) {
			c.logger.Warn().Msg("Cluster state stream closed by Proxy, reconnecting...")
			run.Go(ctx, "fodc.agent.proxy.reconnect", c.logger, func(ctx context.Context) {
				c.reconnect(ctx)
			})
			return
		}
		if recvErr != nil {
			c.streamsMu.RLock()
			disconnected := c.disconnected
			c.streamsMu.RUnlock()
			if disconnected {
				c.logger.Debug().Err(recvErr).Msg("Cluster state stream closed")
				return
			}
			c.logger.Error().Err(recvErr).Msg("Error receiving from cluster state stream, reconnecting...")
			run.Go(ctx, "fodc.agent.proxy.reconnect", c.logger, func(ctx context.Context) {
				c.reconnect(ctx)
			})
			return
		}
		if resp != nil && resp.RequestTopology {
			c.logger.Debug().Msg("Received cluster data request from proxy")
			if sendErr := c.sendClusterTopology(); sendErr != nil {
				c.logger.Error().Err(sendErr).Msg("Failed to send cluster topology to proxy")
			}
		}
	}
}

// RetrieveAndSendMetrics retrieves metrics from Flight Recorder when requested by Proxy.
func (c *Client) RetrieveAndSendMetrics(_ context.Context, filter *MetricsRequestFilter) error {
	c.streamsMu.RLock()
	if c.disconnected || c.metricsStream == nil {
		c.streamsMu.RUnlock()
		return fmt.Errorf("metrics stream not established")
	}
	metricsStream := c.metricsStream
	c.streamsMu.RUnlock()

	datasources := c.flightRecorder.GetDatasources()
	if len(datasources) == 0 {
		req := &fodcv1.StreamMetricsRequest{
			Metrics:   []*fodcv1.Metric{},
			Timestamp: timestamppb.Now(),
		}
		c.streamsMu.Lock()
		if c.metricsStream != metricsStream {
			c.streamsMu.Unlock()
			return fmt.Errorf("metrics stream changed during send")
		}
		sendErr := metricsStream.Send(req)
		c.streamsMu.Unlock()
		if sendErr != nil {
			return fmt.Errorf("failed to send empty metrics response: %w", sendErr)
		}
		return nil
	}

	ds := datasources[0]
	allMetrics := ds.GetMetrics()
	timestamps := ds.GetTimestamps()
	descriptions := ds.GetDescriptions()
	types := ds.GetTypes()

	if filter != nil && (filter.StartTime != nil || filter.EndTime != nil) {
		if timestamps == nil {
			req := &fodcv1.StreamMetricsRequest{
				Metrics:   []*fodcv1.Metric{},
				Timestamp: timestamppb.Now(),
			}
			c.streamsMu.Lock()
			if c.metricsStream != metricsStream {
				c.streamsMu.Unlock()
				return fmt.Errorf("metrics stream changed during send")
			}
			sendErr := metricsStream.Send(req)
			c.streamsMu.Unlock()
			if sendErr != nil {
				return fmt.Errorf("failed to send empty metrics response: %w", sendErr)
			}
			return nil
		}

		timestampValues := timestamps.GetAllValues()
		if len(timestampValues) == 0 {
			req := &fodcv1.StreamMetricsRequest{
				Metrics:   []*fodcv1.Metric{},
				Timestamp: timestamppb.Now(),
			}
			c.streamsMu.Lock()
			if c.metricsStream != metricsStream {
				c.streamsMu.Unlock()
				return fmt.Errorf("metrics stream changed during send")
			}
			sendErr := metricsStream.Send(req)
			c.streamsMu.Unlock()
			if sendErr != nil {
				return fmt.Errorf("failed to send empty metrics response: %w", sendErr)
			}
			return nil
		}

		c.streamsMu.Lock()
		if c.metricsStream != metricsStream {
			c.streamsMu.Unlock()
			return fmt.Errorf("metrics stream changed during send")
		}
		sendErr := c.sendFilteredMetrics(metricsStream, allMetrics, timestampValues, descriptions, types, filter)
		c.streamsMu.Unlock()
		return sendErr
	}

	c.streamsMu.Lock()
	if c.metricsStream != metricsStream {
		c.streamsMu.Unlock()
		return fmt.Errorf("metrics stream changed during send")
	}
	sendErr := c.sendLatestMetrics(metricsStream, allMetrics, descriptions, types)
	c.streamsMu.Unlock()
	return sendErr
}

// toProtoMetricType converts a lowercase prometheus type string to the proto enum.
func toProtoMetricType(t string) fodcv1.MetricType {
	switch t {
	case "counter":
		return fodcv1.MetricType_METRIC_TYPE_COUNTER
	case "gauge":
		return fodcv1.MetricType_METRIC_TYPE_GAUGE
	case "histogram":
		return fodcv1.MetricType_METRIC_TYPE_HISTOGRAM
	case "summary":
		return fodcv1.MetricType_METRIC_TYPE_SUMMARY
	case "untyped":
		return fodcv1.MetricType_METRIC_TYPE_UNTYPED
	default:
		return fodcv1.MetricType_METRIC_TYPE_UNSPECIFIED
	}
}

// sendLatestMetrics sends the latest metrics (most recent values).
func (c *Client) sendLatestMetrics(
	stream fodcv1.FODCService_StreamMetricsClient,
	allMetrics map[string]*flightrecorder.MetricRingBuffer,
	descriptions map[string]string,
	types map[string]string,
) error {
	protoMetrics := make([]*fodcv1.Metric, 0)

	for metricKey, metricBuffer := range allMetrics {
		metricValue := metricBuffer.GetCurrentValue()
		allValues := metricBuffer.GetAllValues()

		if len(allValues) == 0 && metricValue == 0 {
			continue
		}

		parsedKey, parseErr := c.parseMetricKey(metricKey)
		if parseErr != nil {
			c.logger.Warn().Err(parseErr).Str("metric_key", metricKey).Msg("Failed to parse metric key")
			continue
		}

		labelsMap := make(map[string]string)
		for _, label := range parsedKey.Labels {
			labelsMap[label.Name] = label.Value
		}

		protoMetric := &fodcv1.Metric{
			Name:        parsedKey.Name,
			Labels:      labelsMap,
			Value:       metricValue,
			Description: descriptions[parsedKey.Name],
			Type:        toProtoMetricType(metrics.ResolveMetricType(types, parsedKey.Name)),
		}

		protoMetrics = append(protoMetrics, protoMetric)
	}

	req := &fodcv1.StreamMetricsRequest{
		Metrics:   protoMetrics,
		Timestamp: timestamppb.Now(),
	}

	if sendErr := stream.Send(req); sendErr != nil {
		return fmt.Errorf("failed to send metrics: %w", sendErr)
	}

	return nil
}

// sendFilteredMetrics sends metrics filtered by time window.
func (c *Client) sendFilteredMetrics(
	stream fodcv1.FODCService_StreamMetricsClient,
	allMetrics map[string]*flightrecorder.MetricRingBuffer,
	timestampValues []int64,
	descriptions map[string]string,
	types map[string]string,
	filter *MetricsRequestFilter,
) error {
	protoMetrics := make([]*fodcv1.Metric, 0)

	for metricKey, metricBuffer := range allMetrics {
		metricValues := metricBuffer.GetAllValues()
		if len(metricValues) == 0 {
			continue
		}

		parsedKey, parseErr := c.parseMetricKey(metricKey)
		if parseErr != nil {
			c.logger.Warn().Err(parseErr).Str("metric_key", metricKey).Msg("Failed to parse metric key")
			continue
		}

		description := descriptions[parsedKey.Name]

		minLen := len(metricValues)
		if len(timestampValues) < minLen {
			minLen = len(timestampValues)
		}

		labelsMap := make(map[string]string)
		for _, label := range parsedKey.Labels {
			labelsMap[label.Name] = label.Value
		}

		for idx := 0; idx < minLen; idx++ {
			timestampUnix := timestampValues[idx]
			timestamp := time.Unix(timestampUnix, 0)

			if filter.StartTime != nil && timestamp.Before(*filter.StartTime) {
				continue
			}
			if filter.EndTime != nil && timestamp.After(*filter.EndTime) {
				continue
			}

			protoMetric := &fodcv1.Metric{
				Name:        parsedKey.Name,
				Labels:      labelsMap,
				Value:       metricValues[idx],
				Description: description,
				Timestamp:   timestamppb.New(timestamp),
				Type:        toProtoMetricType(metrics.ResolveMetricType(types, parsedKey.Name)),
			}

			protoMetrics = append(protoMetrics, protoMetric)
		}
	}

	req := &fodcv1.StreamMetricsRequest{
		Metrics:   protoMetrics,
		Timestamp: timestamppb.Now(),
	}

	if sendErr := stream.Send(req); sendErr != nil {
		return fmt.Errorf("failed to send metrics: %w", sendErr)
	}

	return nil
}

// SendHeartbeat sends heartbeat to Proxy. The provided context bounds the wait so that a
// stream whose Send is wedged (for example, after the proxy entered graceful_stop and the
// underlying TCP write queue is no longer being drained) cannot pin the heartbeat goroutine
// indefinitely.
func (c *Client) SendHeartbeat(ctx context.Context) error {
	c.streamsMu.RLock()
	if c.disconnected || c.registrationStream == nil {
		c.streamsMu.RUnlock()
		return fmt.Errorf("registration stream not established")
	}
	registrationStream := c.registrationStream
	c.streamsMu.RUnlock()

	req := &fodcv1.RegisterAgentRequest{
		NodeRole:       c.nodeRole,
		Labels:         c.labels,
		PodName:        c.podName,
		ContainerNames: c.containerNames,
	}

	// gRPC client-stream Send does not accept a context, so wrap it in a goroutine and
	// race the result against ctx.Done(). When ctx fires first the in-flight Send is
	// orphaned. The leak is bounded: every reconnect/cleanupStreams/Disconnect path
	// closes the registration stream after waiting on heartbeatWg with bounded timeout,
	// which forces Send to return so the orphan goroutine can write to the (buffered)
	// channel and exit. Without an upstream tear-down the orphan would persist, so the
	// reliability of the heartbeatWaitTimeout-bounded cleanup paths is load-bearing here.
	sendErrCh := make(chan error, 1)
	run.Go(ctx, "fodc.agent.proxy.heartbeat-send", c.logger, func(_ context.Context) {
		sendErrCh <- registrationStream.Send(req)
	})

	var sendErr error
	select {
	case sendErr = <-sendErrCh:
	case <-ctx.Done():
		return fmt.Errorf("heartbeat send aborted: %w", ctx.Err())
	}
	if sendErr != nil {
		// Check if error is due to stream being closed/disconnected
		if errors.Is(sendErr, io.EOF) || errors.Is(sendErr, context.Canceled) {
			return fmt.Errorf("registration stream closed")
		}
		if st, ok := status.FromError(sendErr); ok {
			if st.Code() == codes.Canceled {
				return fmt.Errorf("registration stream closed")
			}
		}
		return fmt.Errorf("failed to send heartbeat: %w", sendErr)
	}

	return nil
}

// cleanupStreams cleans up streams without stopping the connection manager.
func (c *Client) cleanupStreams() {
	c.streamsMu.Lock()
	if c.heartbeatTicker != nil {
		c.heartbeatTicker.Stop()
		c.heartbeatTicker = nil
	}

	oldStopCh := c.stopCh
	c.stopCh = make(chan struct{})
	c.streamsMu.Unlock()

	safeClose(oldStopCh)

	if !waitWaitGroupWithTimeout(&c.heartbeatWg, heartbeatWaitTimeout) {
		c.logger.Warn().Msg("Heartbeat goroutine did not exit in time during cleanupStreams, proceeding anyway")
	}

	c.streamsMu.Lock()
	if c.registrationStream != nil {
		if closeErr := c.registrationStream.CloseSend(); closeErr != nil {
			c.logger.Warn().Err(closeErr).Msg("Error closing registration stream")
		}
		c.registrationStream = nil
	}

	if c.metricsStream != nil {
		if closeErr := c.metricsStream.CloseSend(); closeErr != nil {
			c.logger.Warn().Err(closeErr).Msg("Error closing metrics stream")
		}
		c.metricsStream = nil
	}

	if c.clusterStateStream != nil {
		if closeErr := c.clusterStateStream.CloseSend(); closeErr != nil {
			c.logger.Warn().Err(closeErr).Msg("Error closing cluster state stream")
		}
		c.clusterStateStream = nil
	}

	if c.lifecycleStream != nil {
		if closeErr := c.lifecycleStream.CloseSend(); closeErr != nil {
			c.logger.Warn().Err(closeErr).Msg("Error closing lifecycle stream")
		}
		c.lifecycleStream = nil
	}

	if c.crashStream != nil {
		if closeErr := c.crashStream.CloseSend(); closeErr != nil {
			c.logger.Warn().Err(closeErr).Msg("Error closing crash diagnostics stream")
		}
		c.crashStream = nil
	}
	if c.pressureStream != nil {
		if closeErr := c.pressureStream.CloseSend(); closeErr != nil {
			c.logger.Warn().Err(closeErr).Msg("Error closing pressure profiles stream")
		}
		c.pressureStream = nil
	}
	c.client = nil
	c.streamsMu.Unlock()
}

// Disconnect closes connection to Proxy.
func (c *Client) Disconnect() error {
	c.streamsMu.Lock()
	if c.disconnected {
		c.streamsMu.Unlock()
		return nil
	}

	c.disconnected = true
	close(c.stopCh)

	if c.heartbeatTicker != nil {
		c.heartbeatTicker.Stop()
		c.heartbeatTicker = nil
	}

	c.streamsMu.Unlock()

	// Wait for heartbeat goroutine to exit before closing streams. Bounded so a wedged
	// in-flight Send cannot keep Disconnect from progressing.
	if !waitWaitGroupWithTimeout(&c.heartbeatWg, heartbeatWaitTimeout) {
		c.logger.Warn().Msg("Heartbeat goroutine did not exit in time during Disconnect, proceeding anyway")
	}

	c.streamsMu.Lock()
	if c.registrationStream != nil {
		if closeErr := c.registrationStream.CloseSend(); closeErr != nil {
			c.logger.Warn().Err(closeErr).Msg("Error closing registration stream")
		}
		c.registrationStream = nil
	}

	if c.metricsStream != nil {
		if closeErr := c.metricsStream.CloseSend(); closeErr != nil {
			c.logger.Warn().Err(closeErr).Msg("Error closing metrics stream")
		}
		c.metricsStream = nil
	}

	if c.clusterStateStream != nil {
		if closeErr := c.clusterStateStream.CloseSend(); closeErr != nil {
			c.logger.Warn().Err(closeErr).Msg("Error closing cluster state stream")
		}
		c.clusterStateStream = nil
	}

	if c.lifecycleStream != nil {
		if closeErr := c.lifecycleStream.CloseSend(); closeErr != nil {
			c.logger.Warn().Err(closeErr).Msg("Error closing lifecycle stream")
		}
		c.lifecycleStream = nil
	}

	if c.crashStream != nil {
		if closeErr := c.crashStream.CloseSend(); closeErr != nil {
			c.logger.Warn().Err(closeErr).Msg("Error closing crash diagnostics stream")
		}
		c.crashStream = nil
	}
	if c.pressureStream != nil {
		if closeErr := c.pressureStream.CloseSend(); closeErr != nil {
			c.logger.Warn().Err(closeErr).Msg("Error closing pressure profiles stream")
		}
		c.pressureStream = nil
	}
	c.streamsMu.Unlock()

	c.connManager.stop()

	c.streamsMu.Lock()
	c.client = nil
	c.streamsMu.Unlock()

	c.logger.Info().Msg("Disconnected from FODC Proxy")

	return nil
}

// Start starts the proxy client with automatic reconnection.
func (c *Client) Start(ctx context.Context) error {
	c.connManager.start(ctx)

	for {
		c.streamsMu.RLock()
		stopCh := c.stopCh
		c.streamsMu.RUnlock()

		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-stopCh:
			return nil
		default:
		}

		if connectErr := c.Connect(ctx); connectErr != nil {
			c.logger.Error().Err(connectErr).Msg("Failed to connect to Proxy, retrying...")
			time.Sleep(c.reconnectInterval)
			continue
		}

		if regErr := c.StartRegistrationStream(ctx); regErr != nil {
			c.logger.Error().Err(regErr).Msg("Failed to start registration stream, reconnecting...")
			c.cleanupStreams()
			time.Sleep(c.reconnectInterval)
			continue
		}

		if metricsErr := c.StartMetricsStream(ctx); metricsErr != nil {
			c.logger.Error().Err(metricsErr).Msg("Failed to start metrics stream, reconnecting...")
			c.cleanupStreams()
			time.Sleep(c.reconnectInterval)
			continue
		}

		if clusterStateErr := c.StartClusterStateStream(ctx); clusterStateErr != nil {
			c.logger.Error().Err(clusterStateErr).Msg("Failed to start cluster state stream, reconnecting...")
			c.cleanupStreams()
			time.Sleep(c.reconnectInterval)
			continue
		}

		if c.lifecycleCollector != nil {
			if lifecycleErr := c.StartLifecycleStream(ctx); lifecycleErr != nil {
				c.logger.Error().Err(lifecycleErr).Msg("Failed to start lifecycle stream, reconnecting...")
				c.cleanupStreams()
				time.Sleep(c.reconnectInterval)
				continue
			}
		}

		if c.collectionLister != nil {
			if crashErr := c.StartCrashStream(ctx); crashErr != nil {
				c.logger.Error().Err(crashErr).Msg("Failed to start crash diagnostics stream, reconnecting...")
				c.cleanupStreams()
				time.Sleep(c.reconnectInterval)
				continue
			}
		}

		if c.profileSource != nil {
			if pressureErr := c.StartPressureProfilesStream(ctx); pressureErr != nil {
				c.logger.Error().Err(pressureErr).Msg("Failed to start pressure profiles stream, reconnecting...")
				c.cleanupStreams()
				time.Sleep(c.reconnectInterval)
				continue
			}
		}

		c.logger.Info().Msg("Proxy client started successfully")
		return nil
	}
}

// handleRegistrationStream handles the registration stream.
func (c *Client) handleRegistrationStream(ctx context.Context, stream fodcv1.FODCService_RegisterAgentClient) {
	c.streamsMu.RLock()
	stopCh := c.stopCh
	c.streamsMu.RUnlock()

	for {
		select {
		case <-ctx.Done():
			return
		case <-stopCh:
			return
		default:
		}

		_, recvErr := stream.Recv()
		if errors.Is(recvErr, io.EOF) {
			c.logger.Warn().Msg("Registration stream closed by Proxy, reconnecting...")
			run.Go(ctx, "fodc.agent.proxy.reconnect", c.logger, func(ctx context.Context) {
				c.reconnect(ctx)
			})
			return
		}
		if recvErr != nil {
			c.streamsMu.RLock()
			disconnected := c.disconnected
			c.streamsMu.RUnlock()
			if disconnected {
				c.logger.Debug().Err(recvErr).Msg("Registration stream closed")
				return
			}
			c.logger.Error().Err(recvErr).Msg("Error receiving from registration stream, reconnecting...")
			run.Go(ctx, "fodc.agent.proxy.reconnect", c.logger, func(ctx context.Context) {
				c.reconnect(ctx)
			})
			return
		}
	}
}

// handleMetricsStream handles the metrics stream.
func (c *Client) handleMetricsStream(ctx context.Context, stream fodcv1.FODCService_StreamMetricsClient) {
	c.streamsMu.RLock()
	stopCh := c.stopCh
	c.streamsMu.RUnlock()

	for {
		select {
		case <-ctx.Done():
			return
		case <-stopCh:
			return
		default:
		}

		resp, recvErr := stream.Recv()
		if errors.Is(recvErr, io.EOF) {
			c.logger.Warn().Msg("Metrics stream closed by Proxy, reconnecting...")
			run.Go(ctx, "fodc.agent.proxy.reconnect", c.logger, func(ctx context.Context) {
				c.reconnect(ctx)
			})
			return
		}
		if recvErr != nil {
			c.streamsMu.RLock()
			disconnected := c.disconnected
			c.streamsMu.RUnlock()
			if disconnected {
				c.logger.Debug().Err(recvErr).Msg("Metrics stream closed")
				return
			}
			c.logger.Error().Err(recvErr).Msg("Error receiving from metrics stream, reconnecting...")
			run.Go(ctx, "fodc.agent.proxy.reconnect", c.logger, func(ctx context.Context) {
				c.reconnect(ctx)
			})
			return
		}

		filter := &MetricsRequestFilter{}
		if resp.StartTime != nil {
			startTime := resp.StartTime.AsTime()
			filter.StartTime = &startTime
		}
		if resp.EndTime != nil {
			endTime := resp.EndTime.AsTime()
			filter.EndTime = &endTime
		}

		if retrieveErr := c.RetrieveAndSendMetrics(ctx, filter); retrieveErr != nil {
			c.logger.Error().Err(retrieveErr).Msg("Failed to retrieve and send metrics")
		}
	}
}

// reconnect uses a buffered channel to ensure only one reconnect goroutine runs at a time.
func (c *Client) reconnect(ctx context.Context) {
	select {
	case <-ctx.Done():
		c.logger.Debug().Msg("Context canceled, skipping reconnection...")
		return
	case c.reconnectCh <- struct{}{}:
		// Acquired the slot, proceed with reconnection
	default:
		c.logger.Debug().Msg("Reconnection already in progress, skipping...")
		return
	}

	defer func() {
		<-c.reconnectCh
	}()

	c.streamsMu.Lock()
	if c.disconnected {
		c.streamsMu.Unlock()
		c.logger.Warn().Msg("Already disconnected intentionally, skipping reconnection...")
		return
	}

	c.logger.Info().Msg("Starting reconnection process...")

	if c.heartbeatTicker != nil {
		c.heartbeatTicker.Stop()
		c.heartbeatTicker = nil
	}

	// Swap the stop channel under the lock so subsequent callers see the new one, then
	// release the lock and close the old channel from outside the critical section.
	// This is what allows the heartbeat goroutine to exit via <-stopCh BEFORE we wait on
	// heartbeatWg; closing stopCh AFTER the wait (the previous behavior) deadlocks any
	// heartbeat goroutine that is parked on the select.
	oldStopCh := c.stopCh
	c.stopCh = make(chan struct{})
	c.streamsMu.Unlock()

	safeClose(oldStopCh)

	if !waitWaitGroupWithTimeout(&c.heartbeatWg, heartbeatWaitTimeout) {
		c.logger.Warn().Msg("Heartbeat goroutine did not exit in time during reconnect, proceeding anyway")
	}

	c.streamsMu.Lock()
	if c.registrationStream != nil {
		_ = c.registrationStream.CloseSend()
		c.registrationStream = nil
	}
	if c.metricsStream != nil {
		_ = c.metricsStream.CloseSend()
		c.metricsStream = nil
	}
	if c.clusterStateStream != nil {
		_ = c.clusterStateStream.CloseSend()
		c.clusterStateStream = nil
	}
	if c.lifecycleStream != nil {
		_ = c.lifecycleStream.CloseSend()
		c.lifecycleStream = nil
	}
	if c.crashStream != nil {
		_ = c.crashStream.CloseSend()
		c.crashStream = nil
	}
	if c.pressureStream != nil {
		_ = c.pressureStream.CloseSend()
		c.pressureStream = nil
	}
	c.streamsMu.Unlock()

	connResultCh := c.connManager.RequestConnect(ctx)
	var connResult connResult
	select {
	case <-ctx.Done():
		c.logger.Warn().Msg("Context canceled during reconnection")
		return
	case connResult = <-connResultCh:
	}

	if connResult.err != nil {
		// Transient failure (e.g. the proxy pod is mid-rollout). Do NOT call Disconnect here: that
		// latches c.disconnected and the guard at the top of reconnect would suppress every future
		// attempt. Schedule another try so the agent recovers once the proxy is ready again.
		c.logger.Error().Err(connResult.err).Msg("Failed to reconnect to Proxy; will retry")
		c.scheduleReconnect(ctx)
		return
	}

	if connResult.conn != nil {
		c.streamsMu.Lock()
		c.client = fodcv1.NewFODCServiceClient(connResult.conn)
		c.streamsMu.Unlock()
	}

	// Each stream (re)start below can fail transiently when the new proxy pod is not yet accepting
	// RPCs. reconnect() has already torn down the heartbeat ticker and all stream receive-loops -
	// the only mechanisms that would otherwise re-trigger a reconnect - so returning here without
	// rescheduling leaves the agent permanently disconnected (agents_online drops to 0 until a
	// manual pod restart). Always schedule another attempt on failure.
	if regErr := c.StartRegistrationStream(ctx); regErr != nil {
		c.logger.Error().Err(regErr).Msg("Failed to restart registration stream; will retry")
		c.scheduleReconnect(ctx)
		return
	}

	if metricsErr := c.StartMetricsStream(ctx); metricsErr != nil {
		c.logger.Error().Err(metricsErr).Msg("Failed to restart metrics stream; will retry")
		c.scheduleReconnect(ctx)
		return
	}

	if clusterStateErr := c.StartClusterStateStream(ctx); clusterStateErr != nil {
		c.logger.Error().Err(clusterStateErr).Msg("Failed to restart cluster state stream; will retry")
		c.scheduleReconnect(ctx)
		return
	}

	if c.lifecycleCollector != nil {
		if lifecycleErr := c.StartLifecycleStream(ctx); lifecycleErr != nil {
			c.logger.Error().Err(lifecycleErr).Msg("Failed to restart lifecycle stream; will retry")
			c.scheduleReconnect(ctx)
			return
		}
	}

	if c.collectionLister != nil {
		if crashErr := c.StartCrashStream(ctx); crashErr != nil {
			c.logger.Error().Err(crashErr).Msg("Failed to restart crash diagnostics stream; will retry")
			c.scheduleReconnect(ctx)
			return
		}
	}

	if c.profileSource != nil {
		if pressureErr := c.StartPressureProfilesStream(ctx); pressureErr != nil {
			c.logger.Error().Err(pressureErr).Msg("Failed to restart pressure profiles stream; will retry")
			c.scheduleReconnect(ctx)
			return
		}
	}

	c.logger.Info().Msg("Successfully reconnected to Proxy")
}

// scheduleReconnect arms a delayed reconnect attempt after a transient failure during reconnect
// (e.g. the proxy pod not yet ready), so the agent keeps retrying instead of stalling forever.
// reconnect() is single-flight and re-checks ctx and c.disconnected, so a timer that fires after
// shutdown is a harmless no-op. Each failed reconnect arms exactly one timer, so retries proceed at
// roughly one per interval rather than in a tight loop.
func (c *Client) scheduleReconnect(ctx context.Context) {
	c.streamsMu.Lock()
	disconnected := c.disconnected
	interval := c.reconnectInterval
	c.streamsMu.Unlock()

	if disconnected {
		return
	}
	if interval <= 0 {
		interval = defaultReconnectRetryInterval
	}

	select {
	case <-ctx.Done():
		return
	default:
	}

	reconnectFn := c.reconnect
	if c.reconnectFn != nil {
		reconnectFn = c.reconnectFn
	}

	c.logger.Info().Dur("retry_in", interval).Msg("Scheduling proxy reconnect retry")
	time.AfterFunc(interval, func() {
		reconnectFn(ctx)
	})
}

// startHeartbeat starts the heartbeat ticker.
func (c *Client) startHeartbeat(ctx context.Context) {
	c.streamsMu.Lock()
	if c.heartbeatTicker != nil {
		c.heartbeatTicker.Stop()
	}

	c.heartbeatTicker = time.NewTicker(c.heartbeatInterval)
	tickerCh := c.heartbeatTicker.C
	stopCh := c.stopCh
	c.heartbeatWg.Add(1)
	c.streamsMu.Unlock()

	panicdiag.GoWithRecovery(ctx, panicdiag.RecoveryOptions{
		Component: "fodc-agent-heartbeat",
		Logger:    c.logger,
	}, nil, func(_ *context.Context) {
		defer c.heartbeatWg.Done()
		for {
			select {
			case <-ctx.Done():
				return
			case <-stopCh:
				return
			case <-tickerCh:
				if heartbeatErr := c.SendHeartbeat(ctx); heartbeatErr != nil {
					c.logger.Error().Err(heartbeatErr).Msg("Failed to send heartbeat")
				}
			}
		}
	})
}

// parseMetricKey parses a metric key string into a MetricKey struct.
func (c *Client) parseMetricKey(key string) (metrics.MetricKey, error) {
	if !strings.Contains(key, "{") {
		return metrics.MetricKey{
			Name:   key,
			Labels: []metrics.Label{},
		}, nil
	}

	parts := strings.SplitN(key, "{", 2)
	if len(parts) != 2 {
		return metrics.MetricKey{}, fmt.Errorf("invalid metric key format: %s", key)
	}

	name := parts[0]
	labelStr := strings.TrimSuffix(parts[1], "}")
	if !strings.HasSuffix(parts[1], "}") {
		return metrics.MetricKey{}, fmt.Errorf("invalid metric key format: %s", key)
	}

	labelParts := strings.Split(labelStr, ",")
	labels := make([]metrics.Label, 0, len(labelParts))

	for _, part := range labelParts {
		part = strings.TrimSpace(part)
		if part == "" {
			continue
		}

		kvParts := strings.SplitN(part, "=", 2)
		if len(kvParts) != 2 {
			continue
		}

		keyName := strings.TrimSpace(kvParts[0])
		keyValue := strings.Trim(strings.TrimSpace(kvParts[1]), `"`)

		labels = append(labels, metrics.Label{
			Name:  keyName,
			Value: keyValue,
		})
	}

	return metrics.MetricKey{
		Name:   name,
		Labels: labels,
	}, nil
}
