// Copyright 2023 The envd Authors
//
// Licensed 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 config

import (
	"strconv"

	"github.com/sirupsen/logrus"

	"github.com/tensorchord/envd/pkg/util/osutil"
)

type EntryOptions struct {
	Name           string
	IFace          string
	Port           int
	PrivateKeyPath string

	EnableHostKeyCheck bool
	EnableAgentForward bool
	User               string
}

// AddEntry adds an entry to the user's sshconfig
func AddEntry(eo EntryOptions) error {
	eo.Name = BuildHostname(eo.Name)
	err := add(getSSHConfigPath(), eo)
	if err != nil {
		return err
	}

	// It helps user to directly use remote ssh plugin of VSCode on the host system.
	// VSCode will only check the host ssh files.
	if osutil.IsWsl() {
		err := addEntryToWSLHost(eo)
		if err != nil {
			logrus.Warnf("Failed to add entry to WSL host %s\n", err.Error())
		}
	}
	return nil
}

func add(path string, eo EntryOptions) error {
	cfg, err := getConfig(path)
	if err != nil {
		return err
	}

	_ = removeHost(cfg, eo.Name)

	host := newHost([]string{eo.Name}, []string{"entry generated by envd"})
	host.params = []*param{
		newParam(pubkeyAcceptedKeyTypesKeyword, []string{"+ssh-rsa"}, nil),
		newParam(hostKeyAlgorithms, []string{"+ssh-rsa"}, nil),
		newParam(hostNameKeyword, []string{eo.IFace}, nil),
		newParam(portKeyword, []string{strconv.Itoa(eo.Port)}, nil),
		newParam(userKnownHostsFileKeyword, []string{"/dev/null"}, nil),
		newParam(identityFile, []string{"\"" + eo.PrivateKeyPath + "\""}, nil),
	}
	if !eo.EnableHostKeyCheck {
		host.params = append(host.params,
			newParam(strictHostKeyCheckingKeyword, []string{"no"}, nil))
	}
	if eo.EnableAgentForward {
		host.params = append(host.params,
			newParam(forwardAgentKeyword, []string{"yes"}, nil))
	}
	if eo.User != "" {
		host.params = append(host.params,
			newParam(userKeyword, []string{eo.User}, nil))
	}

	cfg.hosts = append(cfg.hosts, host)
	return save(cfg, path)
}

// RemoveEntry removes the entry to the user's sshconfig if found
func RemoveEntry(name string) error {
	err := remove(getSSHConfigPath(), BuildHostname(name))
	if err != nil {
		return err
	}
	if osutil.IsWsl() {
		logrus.Debug("Try removing entry from WSL's ssh-agent")
		winSshConfig, err := osutil.GetWslHostSshConfig()
		if err != nil {
			return err
		}
		err = remove(winSshConfig, BuildHostname(name))
		if err != nil {
			return err
		}
	}
	return nil
}

func addEntryToWSLHost(eo EntryOptions) error {
	logrus.Debug("Try adding entry to WSL's ssh-agent")
	winSshConfig, err := osutil.GetWslHostSshConfig()
	if err != nil {
		return err
	}
	winKeyPath, err := osutil.CopyToWinEnvdHome(eo.PrivateKeyPath, 0600)
	if err != nil {
		return err
	}
	// Add the entry to the WSL host SSH config
	logrus.Debugf("Adding entry to WSL's ssh-agent: %s", winSshConfig)
	eo.PrivateKeyPath = winKeyPath
	err = add(winSshConfig, eo)
	if err != nil {
		return err
	}
	return nil
}
