// Copyright 2022 Google LLC
//
// 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.

// Implementation of imgstatus which returns the update status of the instance

package imgstatus

import (
	"bytes"
	"fmt"
	"github.com/golang/protobuf/proto"
	"io"
	"io/ioutil"
	"policy-manager/pkg/configfetcher"
	"policy-manager/pkg/devicepolicy"
	"policy-manager/pkg/updateengine"
	"policy-manager/protos"
	"regexp"
	"strconv"
	"strings"
)

// Constants related to lsb-release file.
const (
	lsbChromeOSReleaseKey          = "CHROMEOS_RELEASE_VERSION"
	lsbChromeOSReleaseChannelKey   = "CHROMEOS_RELEASE_TRACK"
	lsbChromeOSReleaseMilestoneKey = "CHROMEOS_RELEASE_CHROME_MILESTONE"
	updateEngineCmd                = "update_engine_client"
)

// parseKeyValuePairs parses text with the form
//      key1=value1
//      key2=value2
//      ...
// into a map.
// All keys must have the format: "[a-zA-Z_]*".
// Any non-conforming lines are ignored.
func parseKeyValuePairs(reader io.Reader) (map[string]string, error) {
	content, err := ioutil.ReadAll(reader)
	if err != nil {
		return nil, err
	}

	pairs := make(map[string]string)

	// Regexp for matching key-value pairs.
	matchPair := regexp.MustCompile(`[a-zA-Z_]*=.*$`)

	// Extract release version string.
	lines := strings.Split(string(content), "\n")
	for _, line := range lines {
		line = strings.TrimSpace(line)
		if matchPair.MatchString(line) {
			pair := strings.Split(line, "=")
			key := strings.TrimSpace(pair[0])
			val := strings.TrimSpace(pair[1])
			pairs[key] = val
		}
	}

	return pairs, nil
}

// convertOSReleaseChannel takes in a release channel string from the
// lsb-release file and converts it into a ReleaseChannel type.
// If the release channel string cannot be parsed,
// RELEASE_CHANNEL_UNSPECIFIED will be returned.
func convertOSReleaseChannel(channelString string) protos.ReleaseChannel {
	// Regexp for matching release channel strings.
	// Release channel strings have the format "<name>-channel". For
	// example, "dev-channel".
	pattern := `^([a-z]{1,}-channel)$`
	matchVersionString := regexp.MustCompile(pattern)

	// Check we have a valid release channel string.
	if matchVersionString.MatchString(channelString) {
		ch := strings.ToUpper(strings.Split(channelString, "-")[0])
		if c, ok := protos.ReleaseChannel_value[ch]; ok {
			return protos.ReleaseChannel(c)
		}
	}

	// Unable to parse release channel.
	return protos.ReleaseChannel_RELEASE_CHANNEL_UNSPECIFIED
}

// parseOSInfoFromLSB returns the current OS version and release channel as
// provided by the lsb-release file. An error is returned only if the lsb file
// cannot be parsed.
func parseOSInfoFromLSB(lsbFile string) (*protos.OSVersion, error) {
	content, err := ioutil.ReadFile(lsbFile)
	if err != nil {
		return nil, err
	}

	// Parse lsb-release file as key value pairs.
	pairs, err := parseKeyValuePairs(bytes.NewReader(content))
	if err != nil {
		return nil, fmt.Errorf("unable to parse '%s': %s", lsbFile, err)
	}

	osVersion := new(protos.OSVersion)

	// Get release version string.
	if version, ok := pairs[lsbChromeOSReleaseKey]; ok {
		osVersion.VersionString = proto.String(version)
	}

	// Get the release channel.
	if channel, ok := pairs[lsbChromeOSReleaseChannelKey]; ok {
		ch := convertOSReleaseChannel(channel)
		osVersion.Channel = &ch
	}

	// Get release milestone.
	if milestone, ok := pairs[lsbChromeOSReleaseMilestoneKey]; ok {
		if n, err := strconv.ParseUint(milestone, 10, 32); err == nil {
			osVersion.Milestone = proto.Uint32(uint32(n))
		}
	}
	return osVersion, nil
}

// getStatusFromUpdateEngine returns the protos.InstanceStatus protobuf
// with the update status filled in.
func getStatusFromUpdateEngine(instanceConfigFile string, updateEngineCmd string) (*protos.InstanceStatus, error) {
	status := new(protos.InstanceStatus)

	// Check if update engine is disabled.
	config, err := devicepolicy.GetInstanceConfig(instanceConfigFile)
	if err != nil {
		return nil, fmt.Errorf("error fetching instance config: %v", err)
	}

	// Do not check update engine status if it is disabled.
	if *config.UpdateStrategy == "update_disabled" {
		status.UpdateStatus = proto.String("UPDATE_ENGINE_DISABLED")
		return status, nil
	}

	// Get response from update engine.
	response, err := updateengine.UpdateEngineStatus(updateEngineCmd)
	if err != nil {
		return nil, err
	}

	pairs, err := parseKeyValuePairs(bytes.NewReader(response))
	if err != nil {
		return nil, fmt.Errorf("unable to parse: %s", err)
	}

	// Set update status.
	if pairs["CURRENT_OP"] == "" {
		return nil, fmt.Errorf("failed to fetch update status. It is possible that update_engine has not fetched the status yet.")
	}
	status.UpdateStatus = proto.String(pairs["CURRENT_OP"])

	// Set new OS version.
	status.NewOsVersion = new(protos.OSVersion)
	status.NewOsVersion.VersionString = proto.String(pairs["NEW_VERSION"])

	// Set last update check timestamp.
	lastTime, err := strconv.ParseInt(pairs["LAST_CHECKED_TIME"], 10, 64)
	if err != nil {
		return nil, fmt.Errorf("unable to parse LAST_CHECKED_TIME: %s", err)
	}
	status.UpdateCheckTimestamp = proto.Int64(lastTime)

	return status, nil
}

// GetStatus() returns the current status of the instance as given by
// /etc/lsb-release file and update-engine.
// If an error occured, the error will be returned along with a status protobuf
// where the error field is set to the message describing the error. If the
// error field is set, all other fields other than the InstanceId are invalid.
// Note that this function must always return a status with InstanceId field
// populated regardless of whether an error occurred or not.
func GetStatus(lsbFile string, instanceConfigFile string) (*protos.InstanceStatus, error) {
	status := new(protos.InstanceStatus)
	var instanceID uint64
	id, err := configfetcher.GetInstanceID()
	if err == nil {
		instanceID = id
	}
	status.InstanceId = proto.Uint64(instanceID)

	// Get OS version from lsb-release file.
	osVersion, err := parseOSInfoFromLSB(lsbFile)
	if err != nil {
		status.Error = proto.String(err.Error())
		return status, err
	}

	status.OsVersion = osVersion

	// Get update status from update-engine via update_engine_client.
	updateStatus, err := getStatusFromUpdateEngine(instanceConfigFile, updateEngineCmd)
	if err != nil {
		status.UpdateStatus = nil
		status.Error = proto.String(err.Error())
		return status, err
	}

	proto.Merge(status, updateStatus)
	return status, nil
}