// Copyright 2020 Nokia
// Licensed under the BSD 3-Clause License.
// SPDX-License-Identifier: BSD-3-Clause

package vr_sros

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"os"
	"path"
	"path/filepath"
	"regexp"
	"strings"
	"time"

	"github.com/charmbracelet/log"
	"github.com/scrapli/scrapligo/driver/network"
	"github.com/scrapli/scrapligo/driver/options"
	scraplilogging "github.com/scrapli/scrapligo/logging"
	"github.com/scrapli/scrapligo/platform"
	"github.com/scrapli/scrapligo/transport"
	"github.com/scrapli/scrapligo/util"
	clabconstants "github.com/srl-labs/containerlab/constants"
	clabexec "github.com/srl-labs/containerlab/exec"
	clabnetconf "github.com/srl-labs/containerlab/netconf"
	clabnodes "github.com/srl-labs/containerlab/nodes"
	clabtypes "github.com/srl-labs/containerlab/types"
	clabutils "github.com/srl-labs/containerlab/utils"
	"golang.org/x/crypto/ssh"
)

var (
	kindNames          = []string{"nokia_sros", "vr-sros", "vr-nokia_sros"}
	defaultCredentials = clabnodes.NewCredentials("admin", "admin")

	InterfaceRegexp = regexp.MustCompile(`1/1/(?P<port>\d+)`)
	InterfaceOffset = 1
	InterfaceHelp   = "1/1/X (where X >= 1) or ethX (where X >= 1)"
)

const (
	generateable     = true
	generateIfFormat = "1/1/%d"

	vrsrosDefaultType          = "sr-1"
	scrapliPlatformName        = "nokia_sros"
	scrapliPlatformNameClassic = "nokia_sros_classic"
	// envSrosConfigMode is the env var that controls the CLI mode used by SR OS.
	// When set to "classic" or "mixed", the classic CLI scrapligo platform is used.
	// Default (unset or "model-driven") uses the MD-CLI platform.
	envSrosConfigMode = "CLAB_SROS_CONFIG_MODE"
	configDirName     = "tftpboot"
	startupCfgFName   = "config.txt"
	licenseFName      = "license.txt"

	// OCI image title label used to detect SR-SIM container image (must use kind nokia_srsim).
	ociImageTitleLabel = "org.opencontainers.image.title"
	srsimImageTitle    = "srsim"
)

// SROSTemplateData holds ssh keys for template generation.
type SROSTemplateData struct {
	SSHPubKeysRSA   []string
	SSHPubKeysECDSA []string
}

// Register registers the node in the NodeRegistry.
func Register(r *clabnodes.NodeRegistry) {
	generateNodeAttributes := clabnodes.NewGenerateNodeAttributes(generateable, generateIfFormat)
	platformAttrs := &clabnodes.PlatformAttrs{
		ScrapliPlatformName: scrapliPlatformName,
	}

	nrea := clabnodes.NewNodeRegistryEntryAttributes(
		defaultCredentials,
		generateNodeAttributes,
		platformAttrs,
	)

	r.Register(kindNames, func() clabnodes.Node {
		return new(vrSROS)
	}, nrea)
}

type vrSROS struct {
	clabnodes.VRNode
	// SSH public keys extracted from the clab host
	sshPubKeys []ssh.PublicKey
}

func (s *vrSROS) Init(cfg *clabtypes.NodeConfig, opts ...clabnodes.NodeOption) error {
	// Init DefaultNode
	s.VRNode = *clabnodes.NewVRNode(s, defaultCredentials, scrapliPlatformName)
	// set virtualization requirement
	s.HostRequirements.VirtRequired = true
	s.LicensePolicy = clabtypes.LicensePolicyWarn
	// SR OS requires unbound pubkey authentication mode until this is
	// gets fixed in later SR OS release.
	s.SSHConfig.PubkeyAuthentication = clabtypes.PubkeyAuthValueUnbound

	s.Cfg = cfg
	for _, o := range opts {
		o(s)
	}
	// vr-sros type sets the vrnetlab/sros variant (https://github.com/hellt/vrnetlab/sros)
	if s.Cfg.NodeType == "" {
		s.Cfg.NodeType = vrsrosDefaultType
	}
	// env vars are used to set launch.py arguments in vrnetlab container
	defEnv := map[string]string{
		"CONNECTION_MODE":    clabnodes.VrDefConnMode,
		"DOCKER_NET_V4_ADDR": s.Mgmt.IPv4Subnet,
		"DOCKER_NET_V6_ADDR": s.Mgmt.IPv6Subnet,
	}
	s.Cfg.Env = clabutils.MergeStringMaps(defEnv, s.Cfg.Env)

	// mount tftpboot dir
	s.Cfg.Binds = append(s.Cfg.Binds, fmt.Sprint(path.Join(s.Cfg.LabDir, "tftpboot"), ":/tftpboot"))
	if s.Cfg.Env["CONNECTION_MODE"] == "macvtap" {
		// mount dev dir to enable macvtap
		s.Cfg.Binds = append(s.Cfg.Binds, "/dev:/dev")
	}

	s.Cfg.Cmd = fmt.Sprintf(
		"--trace --connection-mode %s --hostname %s --variant %q",
		s.Cfg.Env["CONNECTION_MODE"],
		s.Cfg.ShortName,
		s.Cfg.NodeType,
	)

	s.InterfaceRegexp = InterfaceRegexp
	s.InterfaceOffset = InterfaceOffset
	s.InterfaceHelp = InterfaceHelp

	if len(s.Cfg.Components) > 0 {
		log.Warnf(
			"node %q: kind nokia_sros (vrnetlab) does not support components; components are ignored. Use kind nokia_srsim for distributed/chassis topologies with components",
			s.Cfg.ShortName,
		)
	}

	return nil
}

func (s *vrSROS) PreDeploy(ctx context.Context, params *clabnodes.PreDeployParams) error {
	clabutils.CreateDirectory(s.Cfg.LabDir, clabconstants.PermissionsOpen)
	_, err := s.LoadOrGenerateCertificate(params.Cert, params.TopologyName)
	if err != nil {
		return err
	}

	// store public keys extracted from clab host
	s.sshPubKeys = params.SSHPubKeys

	if err := s.verifyNokiaSrosImage(ctx); err != nil {
		return err
	}

	return createVrSROSFiles(s)
}

// verifyNokiaSrosImage ensures the image used with kind nokia_sros is not the SRsim
// container image (use kind nokia_srsim for that). It inspects the image label
// org.opencontainers.image.title and returns an error if it is "srsim".
func (s *vrSROS) verifyNokiaSrosImage(ctx context.Context) error {
	if s.GetRuntime() == nil {
		return nil
	}
	insp, err := s.GetRuntime().InspectImage(ctx, s.Cfg.Image)
	if err != nil {
		// Skip check when runtime does not support image inspection (e.g. Podman).
		if strings.Contains(err.Error(), "not implemented") {
			log.Debug(
				"Skipping nokia_sros image kind check: runtime does not support image inspection",
			)
			return nil
		}
		return err
	}
	if insp != nil && insp.Config.Labels != nil {
		if title, ok := insp.Config.Labels[ociImageTitleLabel]; ok && title == srsimImageTitle {
			return fmt.Errorf(
				"node %q: kind is nokia_sros (vrnetlab) but the provided image is for the nokia_srsim container image; use kind: nokia_srsim with this image, or use a vrnetlab nokia_sros image for kind: nokia_sros",
				s.Cfg.ShortName,
			)
		}
	}
	return nil
}

func (s *vrSROS) PostDeploy(ctx context.Context, _ *clabnodes.PostDeployParams) error {
	// b holds the configuration to be applied to the node
	b := &bytes.Buffer{}

	// apply partial configs if partial config is used and existing node config does not exist
	if clabutils.IsPartialConfigFile(s.Cfg.StartupConfig) && !nodeConfigExists(s.Cfg.LabDir) {
		log.Info("Adding configuration",
			"node", s.Cfg.LongName,
			"type", "partial",
			"source", s.Cfg.StartupConfig)

		r, err := os.Open(s.Cfg.StartupConfig)
		if err != nil {
			return err
		}

		defer r.Close() // skipcq: GO-S2307

		_, err = io.Copy(b, r)
		if err != nil {
			return err
		}
	}

	// skip ssh key configuration if CLAB_SKIP_SROS_SSH_KEY_CONFIG env var is set
	// which is needed for SR OS nodes running in classic CLI mode, because our key
	// injection mechanism assumes MD-CLI mode.
	_, skipSSHKeyCfg := os.LookupEnv("CLAB_SKIP_SROS_SSH_KEY_CONFIG")

	if len(s.sshPubKeys) > 0 && !skipSSHKeyCfg {
		log.Info("Adding public keys configuration", "node", s.Cfg.LongName)

		sshConf, err := s.generateSSHPublicKeysConfig()
		if err != nil {
			return err
		}

		_, err = io.Copy(b, sshConf)
		if err != nil {
			return err
		}
	}

	// apply the aggregated config snippets
	if b.Len() > 0 {
		err := s.applyPartialConfig(ctx, s.Cfg.MgmtIPv4Address, s.scrapliPlatform(),
			s.Cfg.Credentials.Username, s.Cfg.Credentials.Password,
			b,
		)
		if err != nil {
			return err
		}
	}

	return nil
}

// scrapliPlatform returns the scrapligo platform name based on the configured CLI mode.
// When CLAB_SROS_CONFIG_MODE is "classic" or "mixed", the classic CLI platform is used.
// The default (unset or "model-driven") uses the MD-CLI platform.
func (s *vrSROS) scrapliPlatform() string {
	cfgMode := strings.ToLower(s.Cfg.Env[envSrosConfigMode])
	if cfgMode == "classic" || cfgMode == "mixed" {
		return scrapliPlatformNameClassic
	}
	return scrapliPlatformName
}

func (s *vrSROS) SaveConfig(_ context.Context) (*clabnodes.SaveConfigResult, error) {
	err := clabnetconf.SaveRunningConfig(s.Cfg.LongName,
		s.Cfg.Credentials.Username,
		s.Cfg.Credentials.Password,
		s.scrapliPlatform(),
	)
	if err != nil {
		return nil, err
	}

	log.Infof("saved %s running configuration to startup configuration file\n", s.Cfg.ShortName)
	return nil, nil
}

func createVrSROSFiles(node clabnodes.Node) error {
	nodeCfg := node.Config()

	// use default startup config load function if config in full form is provided
	if !clabutils.IsPartialConfigFile(nodeCfg.StartupConfig) {
		// do not create new config file if there's existing config file
		if nodeConfigExists(nodeCfg.LabDir) {
			log.Infof("Using existing config file (%s) instead of applying a new one",
				filepath.Join(nodeCfg.LabDir, configDirName, startupCfgFName))
		} else {
			clabnodes.LoadStartupConfigFileVr(node, configDirName, startupCfgFName)
		}
	}

	if nodeCfg.License != "" {
		// copy license file to node specific lab directory
		src := nodeCfg.License
		dst := filepath.Join(nodeCfg.LabDir, configDirName, licenseFName)
		if err := clabutils.CopyFile(context.Background(), src, dst,
			clabconstants.PermissionsFileDefault); err != nil {
			return fmt.Errorf("file copy [src %s -> dst %s] failed %v", src, dst, err)
		}
		log.Debugf("CopyFile src %s -> dst %s succeeded", src, dst)
	}

	return nil
}

// nodeConfigExists returns true if a file at <labdir>/<node>/tftpboot/config.txt exists.
func nodeConfigExists(labDir string) bool {
	_, err := os.Stat(filepath.Join(labDir, configDirName, startupCfgFName))
	return err == nil
}

// isHealthy checks if the "/health" file created by vrnetlab exists and contains "0 running".
func (s *vrSROS) isHealthy(ctx context.Context) bool {
	ex := clabexec.NewExecCmdFromSlice([]string{"grep", "0 running", "/health"})

	res, err := s.RunExec(ctx, ex)
	if err != nil {
		return false
	}

	log.Debugf("Node %q health status: %v", s.Cfg.ShortName, res.ReturnCode == 0)

	return res.ReturnCode == 0
}

// applyPartialConfig applies partial configuration to the SR OS.
func (s *vrSROS) applyPartialConfig(ctx context.Context, addr, platformName,
	username, password string, config io.Reader,
) error { // skipcq: GO-R1005
	var err error
	var d *network.Driver

	configContent, err := clabutils.SubstituteEnvsAndTemplate(config, s.Cfg)
	if err != nil {
		return err
	}

	configContentStr := configContent.String()

	// check file contains content, otherwise exit early
	if strings.TrimSpace(configContentStr) == "" {
		return nil
	}

	log.Info("Waiting for node to be ready. This may take a while",
		"node", s.Cfg.LongName,
		"log", fmt.Sprintf("docker logs -f %[1]s", s.Cfg.LongName),
	)

	for loop := true; loop; {
		if !s.isHealthy(ctx) {
			time.Sleep(5 * time.Second) // cool-off period
			log.Debugf("Waiting for %s to become healthy", s.Cfg.ShortName)
			continue
		}

		select {
		case <-ctx.Done():
			return fmt.Errorf("%s: timed out waiting to accept configs", addr)
		default:
			sl := log.StandardLog(log.StandardLogOptions{
				ForceLevel: log.DebugLevel,
			})
			li, err := scraplilogging.NewInstance(
				scraplilogging.WithLevel("debug"),
				scraplilogging.WithLogger(sl.Print))
			if err != nil {
				return err
			}

			opts := []util.Option{
				options.WithAuthNoStrictKey(),
				options.WithAuthUsername(username),
				options.WithAuthPassword(password),
				options.WithTransportType(transport.StandardTransport),
				options.WithTimeoutOps(5 * time.Second),
				options.WithLogger(li),
			}

			p, err := platform.NewPlatform(platformName, addr, opts...)
			if err != nil {
				return fmt.Errorf("%s: failed to create platform: %+v", addr, err)
			}

			d, err = p.GetNetworkDriver()
			if err != nil {
				return fmt.Errorf("%s: could not create the driver: %+v", addr, err)
			}

			err = d.Open()
			if err == nil {
				// driver successfully opened, exit the loop
				loop = false
			} else {
				log.Debugf("%s: not yet ready - %v", addr, err)
				time.Sleep(5 * time.Second) // cool-off period
			}
		}
	}

	// Normalize character sequences to avoid interaction issues with CLI
	replacer := strings.NewReplacer(
		"\r\n", "\n", // replace EOL CRLF with LF
		"\t", "    ", // replace tabs with 4 spaces
	)
	configContentStr = replacer.Replace(configContentStr)

	// converting string to newline delimited string slice
	cfgs := strings.Split(configContentStr, "\n")

	// config snippets should not have commit command, so we need to commit manually
	// and quit from the config mode
	cfgs = append(cfgs, "commit", "/admin save", "/exit all", "quit-config")

	mr, err := d.SendConfigs(cfgs)
	if err != nil || (mr != nil && mr.Failed != nil) {
		if mr != nil {
			return fmt.Errorf("failed to apply config; error: %+v %+v", err, mr.Failed)
		} else {
			return fmt.Errorf("failed to apply config; error: %+v", err)
		}
	}

	return nil
}
