/*
 * Copyright (c) 2021. Ant Group. All rights reserved.
 * Copyright (c) 2022. Nydus Developers. All rights reserved.
 *
 * SPDX-License-Identifier: Apache-2.0
 */

package metrics

import (
	"context"
	"fmt"
	"os"
	"time"

	"github.com/pkg/errors"

	"github.com/containerd/log"
	"github.com/containerd/nydus-snapshotter/config"
	"github.com/containerd/nydus-snapshotter/pkg/daemon/types"
	"github.com/containerd/nydus-snapshotter/pkg/manager"
	"github.com/containerd/nydus-snapshotter/pkg/metrics/collector"
	"github.com/containerd/nydus-snapshotter/pkg/metrics/tool"
)

type ServerOpt func(*Server) error

type Server struct {
	managers          []*manager.Manager
	snCollectors      []*collector.SnapshotterMetricsCollector
	fsCollector       *collector.FsMetricsVecCollector
	cacheCollector    *collector.CacheMetricsVecCollector
	inflightCollector *collector.InflightMetricsVecCollector
	hungIOInterval    time.Duration
	collectInterval   time.Duration
}

func WithProcessManagers(managers []*manager.Manager) ServerOpt {
	return func(s *Server) error {
		s.managers = append(s.managers, managers...)
		return nil
	}
}

func WithCollectInterval(interval time.Duration) ServerOpt {
	return func(s *Server) error {
		if interval < 0 {
			return fmt.Errorf("collect interval (%v) must be positive", interval)
		}
		s.collectInterval = interval
		return nil
	}
}

func WithHungIOInterval(hungIOInterval time.Duration) ServerOpt {
	return func(s *Server) error {
		if hungIOInterval < 0 {
			return fmt.Errorf("hung IO interval (%v) must be positive", hungIOInterval)
		}
		s.hungIOInterval = hungIOInterval
		return nil
	}
}

func NewServer(ctx context.Context, opts ...ServerOpt) (*Server, error) {
	var s Server
	for _, o := range opts {
		if err := o(&s); err != nil {
			return nil, err
		}
	}

	s.fsCollector = collector.NewFsMetricsVecCollector()
	s.inflightCollector = collector.NewInflightMetricsVecCollector(s.hungIOInterval)
	s.cacheCollector = collector.NewCacheMetricsVecCollector()
	for _, pm := range s.managers {
		snCollector, err := collector.NewSnapshotterMetricsCollector(ctx, pm.CacheDir(), os.Getpid())
		if err != nil {
			return nil, errors.Wrap(err, "new snapshotter metrics collector failed")
		}
		s.snCollectors = append(s.snCollectors, snCollector)
	}

	return &s, nil
}

func (s *Server) CollectDaemonResourceMetrics(_ context.Context) {
	var daemonResource collector.DaemonResourceCollector
	for _, pm := range s.managers {
		// Collect daemon resource usage metrics.
		daemons := pm.ListDaemons()
		for _, d := range daemons {
			memRSS, err := tool.GetProcessMemoryRSSKiloBytes(d.Pid())
			if err != nil {
				log.L.Warnf("Failed to get daemon %s RSS memory", d.ID())
			}

			daemonResource.DaemonID = d.ID()
			daemonResource.Value = memRSS
			daemonResource.Collect()
		}
	}
}

func (s *Server) CollectFsMetrics(ctx context.Context) {
	var fsMetricsVec []collector.FsMetricsCollector

	for _, pm := range s.managers {
		// Collect FS metrics from fusedev daemons.
		if pm.FsDriver != config.FsDriverFusedev {
			continue
		}

		daemons := pm.ListDaemons()
		for _, d := range daemons {
			// Skip daemons that are not serving
			if d.State() != types.DaemonStateRunning {
				continue
			}

			for _, i := range d.RafsCache.List() {
				var sid string

				if d.IsSharedDaemon() {
					sid = i.SnapshotID
				} else {
					sid = ""
				}

				fsMetrics, err := d.GetFsMetrics(sid)
				if err != nil {
					log.G(ctx).Errorf("failed to get fs metric: %v", err)
					continue
				}

				fsMetricsVec = append(fsMetricsVec, collector.FsMetricsCollector{
					Metrics:  fsMetrics,
					ImageRef: i.ImageID,
				})
			}
		}
	}

	if fsMetricsVec != nil {
		s.fsCollector.MetricsVec = fsMetricsVec
		s.fsCollector.Collect()
	}
}

func (s *Server) CollectCacheMetrics(ctx context.Context) {
	var cacheMetricsVec []collector.CacheMetricsCollector

	for _, pm := range s.managers {
		daemons := pm.ListDaemons()
		for _, d := range daemons {
			// Skip daemons that are not serving
			if d.State() != types.DaemonStateRunning {
				continue
			}

			for _, i := range d.RafsCache.List() {
				var sid string

				if d.IsSharedDaemon() {
					sid = i.SnapshotID
				} else {
					sid = ""
				}

				cacheMetrics, err := d.GetCacheMetrics(sid)
				if err != nil {
					log.G(ctx).Errorf("failed to get cache metric: %v", err)
					continue
				}

				cacheMetricsVec = append(cacheMetricsVec, collector.CacheMetricsCollector{
					Metrics:  cacheMetrics,
					ImageRef: i.ImageID,
					DaemonID: d.ID(),
				})
			}
		}
	}

	if cacheMetricsVec != nil {
		s.cacheCollector.MetricsVec = cacheMetricsVec
		s.cacheCollector.Collect()
	}
}

func (s *Server) CollectInflightMetrics(ctx context.Context) {
	inflightMetricsVec := make([]*types.InflightMetrics, 0, 16)
	for _, pm := range s.managers {
		// Collect inflight metrics from fusedev daemons.
		if pm.FsDriver != config.FsDriverFusedev {
			continue
		}

		daemons := pm.ListDaemons()
		for _, d := range daemons {

			// Only count for daemon that is serving
			if d.State() != types.DaemonStateRunning {
				continue
			}

			inflightMetrics, err := d.GetInflightMetrics()
			if err != nil {
				log.G(ctx).Errorf("failed to get inflight metric: %v", err)
				continue
			}
			inflightMetricsVec = append(inflightMetricsVec, inflightMetrics)
		}
	}

	if inflightMetricsVec != nil {
		s.inflightCollector.MetricsVec = inflightMetricsVec
		s.inflightCollector.Collect()
	}
}

func (s *Server) StartCollectMetrics(ctx context.Context) error {
	timer := time.NewTicker(s.collectInterval)
	// The timer period is the same as the interval for determining hung IOs.
	//
	// Since the elapsed time of hung IO is configuration dependent,
	// e.g. timeout * retry times when the backend is a registry.
	// Therefore, we cannot get complete hung IO data after 1 minute.
	InflightTimer := time.NewTicker(s.hungIOInterval)

outer:
	for {
		select {
		case <-timer.C:
			s.CollectFsMetrics(ctx)
			s.CollectCacheMetrics(ctx)
			s.CollectDaemonResourceMetrics(ctx)
			// Collect snapshotter metrics.
			for _, snCollector := range s.snCollectors {
				snCollector.Collect()
			}
		case <-InflightTimer.C:
			s.CollectInflightMetrics(ctx)
		case <-ctx.Done():
			log.G(ctx).Infof("cancel metrics collecting")
			break outer
		}
	}

	return nil
}
