// Copyright 2025 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

// Package model contains our data model for actions and functions to convert from/to REAPI types.
package model

import (
	"bytes"
	"errors"
	"fmt"
	"io/fs"
	"path/filepath"
	"strings"
	"time"

	repb "github.com/bazelbuild/remote-apis/build/bazel/remote/execution/v2"
	"github.com/google/uuid"
	iradix "github.com/hashicorp/go-immutable-radix/v2"

	"go.chromium.org/build/kajiya/blobstore"
	"go.chromium.org/build/kajiya/digest"
)

type EnvVar struct {
	Name  string
	Value string
}

// KajiyaFile represents an input file used by an Action.
type KajiyaFile struct {
	Name     string
	Digest   digest.Digest
	Mtime    time.Time
	UnixMode fs.FileMode
}

// KajiyaSymlink represents an input symlink used by an Action.
type KajiyaSymlink struct {
	Name   string
	Target string
}

// FileType specifies whether an output file is a file, directory, or unknown type.
type FileType int

const (
	Unknown FileType = iota
	File
	Directory
)

// KajiyaOutput represents an output generated by an action. It can be either a file or a directory.
type KajiyaOutput struct {
	Name string
	Type FileType
}

// KajiyaDirectory represents a directory in the input root of an Action.
type KajiyaDirectory struct {
	Digest   digest.Digest
	Files    []KajiyaFile
	Symlinks []KajiyaSymlink
	Dirs     []string
	Outputs  []KajiyaOutput

	Mtime    time.Time
	UnixMode fs.FileMode
}

// DirectoryTrie is an immutable Radix Trie that contains a KajiyaDirectory for each input present
// in the input root of an Action.
type DirectoryTrie = iradix.Tree[*KajiyaDirectory]

// Action contains all the information needed to execute an action.
type Action struct {
	// A unique identifier for this execute operation.
	OpName uuid.UUID

	// Digests for the action, command, and input root.
	ActionDigest    digest.Digest
	CommandDigest   digest.Digest
	InputRootDigest digest.Digest

	// Parameters influencing the use of caching.
	DoNotCache bool

	// Parameters influencing the execution.
	Args           []string
	EnvVars        []EnvVar
	WorkingDir     string
	InputTrie      *DirectoryTrie
	Timeout        time.Duration
	Platform       map[string][]string
	ContainerImage string

	// Parameters influencing the handling of outputs.
	CaptureWholeTree bool
}

// LoadAction loads an Action from the CAS given its digest.
func LoadAction(actionDigest digest.Digest, cas *blobstore.ContentAddressableStorage) (ka *Action, err error) {
	// Fetch the Action from the CAS.
	action, err := cas.Action(actionDigest)
	if err != nil {
		return nil, err
	}

	ka = &Action{
		ActionDigest: actionDigest,
		DoNotCache:   action.DoNotCache,
	}

	// Verify that the timeout is well-formed.
	if action.Timeout != nil {
		err := action.Timeout.CheckValid()
		if err != nil {
			return nil, fmt.Errorf("timeout is not a valid duration: %w", err)
		}
		ka.Timeout = action.Timeout.AsDuration()
	}

	// Fetch the Command from the CAS.
	ka.CommandDigest, err = digest.NewFromProto(action.CommandDigest)
	if err != nil {
		return nil, err
	}
	cmd, err := cas.Command(ka.CommandDigest)
	if err != nil {
		return nil, err
	}
	ka.Args = cmd.Arguments

	// Build a map of environment variables.
	prevName := ""
	for _, v := range cmd.EnvironmentVariables {
		if v.Name <= prevName {
			return nil, fmt.Errorf("environment variable names must be sorted and unique, but %q <= %q", v.Name, prevName)
		}
		prevName = v.Name

		ka.EnvVars = append(ka.EnvVars, EnvVar{
			Name:  v.Name,
			Value: v.Value,
		})
	}

	// Build a map of platform properties. REAPI v2.2+ says that the action's platform
	// properties take precedence over the command's, if both are present.
	p := action.Platform
	if p == nil {
		p = cmd.GetPlatform() //nolint:staticcheck
	}
	ka.Platform = make(map[string][]string)
	if p != nil {
		prevName = ""
		for _, prop := range p.Properties {
			// Note: Duplicate property names are explicitly allowed in the spec.
			if prop.Name < prevName {
				return nil, fmt.Errorf("platform properties must be sorted by name, but %q < %q", prop.Name, prevName)
			}
			prevName = prop.Name

			ka.Platform[prop.Name] = append(ka.Platform[prop.Name], prop.Value)
		}
	}

	if containerImages, ok := ka.Platform["container-image"]; ok {
		if len(containerImages) != 1 {
			return nil, fmt.Errorf("platform property container-image must have exactly one value, but got %d", len(containerImages))
		}
		ka.ContainerImage = containerImages[0]
	}

	ka.InputRootDigest, err = digest.NewFromProto(action.InputRootDigest)
	if err != nil {
		return nil, err
	}
	// Convert the input root to a radix trie.
	ka.InputTrie, err = treeToTrie(cas, action.InputRootDigest)
	if err != nil {
		return nil, err
	}

	// Validate the working directory.
	ka.WorkingDir = cmd.WorkingDirectory
	if ka.WorkingDir == "" {
		ka.WorkingDir = "."
	}
	if filepath.FromSlash(ka.WorkingDir) != filepath.Clean(ka.WorkingDir) {
		return nil, fmt.Errorf("working directory is not a clean path, wanted %q, got %q", filepath.Clean(ka.WorkingDir), filepath.FromSlash(cmd.WorkingDirectory))
	}
	if _, ok := ka.InputTrie.Get([]byte(ka.WorkingDir + "/")); !ok {
		return nil, fmt.Errorf("working directory %q not found in input root", ka.WorkingDir)
	}

	// Add the output paths to our trie.
	if len(cmd.OutputPaths) > 0 { // REAPI v2.1+
		if len(cmd.OutputPaths) == 1 && cmd.OutputPaths[0] == "" {
			ka.CaptureWholeTree = true
		} else {
			ka.addOutputsToTrie(cmd.OutputPaths, Unknown)
		}
	} else { // REAPI v2.0 (deprecated)
		if len(cmd.OutputDirectories) == 1 && cmd.OutputDirectories[0] == "" { //nolint:staticcheck
			ka.CaptureWholeTree = true
		} else {
			ka.addOutputsToTrie(cmd.OutputDirectories, Directory) //nolint:staticcheck
			ka.addOutputsToTrie(cmd.OutputFiles, File)            //nolint:staticcheck
		}
	}

	// TODO: not implemented yet
	if ka.CaptureWholeTree {
		return nil, fmt.Errorf("capturing the whole output tree is not supported yet")
	}

	return ka, nil
}

func (ka *Action) addOutputsToTrie(outputs []string, ft FileType) error {
	if len(outputs) == 0 {
		return nil
	}

	prevName := ""
	for _, path := range outputs {
		if path <= prevName {
			return fmt.Errorf("output paths must be sorted and unique, but %q <= %q", path, prevName)
		}
		prevName = path

		if strings.HasPrefix(path, "/") || strings.HasSuffix(path, "/") || path != filepath.Clean(path) || path == "." {
			return fmt.Errorf("invalid output path %q (has a leading or trailing slash, or is not clean)", path)
		}

		path = filepath.Join(ka.WorkingDir, path)
		if !filepath.IsLocal(path) {
			return fmt.Errorf("output path %q escapes the input root", path)
		}
		parentPath, parent, _ := ka.InputTrie.Root().LongestPrefix([]byte(path))
		if !bytes.HasSuffix(parentPath, []byte("/")) {
			return fmt.Errorf("parent path %q does not end with a slash", parentPath)
		}
		pathInParent := strings.TrimPrefix(path, string(parentPath))
		parent.Outputs = append(parent.Outputs, KajiyaOutput{Name: pathInParent, Type: ft})
	}
	return nil
}

func convertUnixMode(explicitMode uint32, isExecutable bool) (fs.FileMode, error) {
	if explicitMode != 0 {
		if explicitMode&^0777 != 0 {
			return 0, fmt.Errorf("invalid mode %o", explicitMode)
		}
		return fs.FileMode(explicitMode), nil
	}

	if isExecutable {
		return 0755, nil
	}
	return 0644, nil
}

func treeToTrie(cas *blobstore.ContentAddressableStorage, rootDigest *repb.Digest) (*DirectoryTrie, error) {
	var dirQueue []*repb.DirectoryNode
	trieBuilder := iradix.New[*KajiyaDirectory]().Txn()
	var missingBlobs []digest.Digest

	// Iteratively add directories to the trie.
	dirQueue = append(dirQueue, &repb.DirectoryNode{Digest: rootDigest, Name: "./"})
	for len(dirQueue) > 0 {
		dirNode := dirQueue[0]
		dirQueue = dirQueue[1:]

		// Get the directory from the CAS.
		dirDigest, err := digest.NewFromProto(dirNode.Digest)
		if err != nil {
			return nil, err
		}
		dir, err := cas.Directory(dirDigest)
		if err != nil {
			var mberr *blobstore.MissingBlobsError
			if errors.As(err, &mberr) {
				missingBlobs = append(missingBlobs, mberr.Blobs...)
				continue
			}
			return nil, err
		}
		unixMode, err := convertUnixMode(dir.GetNodeProperties().GetUnixMode().GetValue(), true)
		if err != nil {
			return nil, fmt.Errorf("directory %q: %w", dirNode.Name, err)
		}
		kd := KajiyaDirectory{
			Digest:   dirDigest,
			Files:    make([]KajiyaFile, 0, len(dir.Files)),
			Symlinks: make([]KajiyaSymlink, 0, len(dir.Symlinks)),
			Dirs:     make([]string, 0, len(dir.Directories)),
			UnixMode: unixMode,
		}
		if dir.GetNodeProperties().GetMtime() != nil {
			kd.Mtime = dir.NodeProperties.Mtime.AsTime()
		}

		prevName := ""
		for _, file := range dir.Files {
			if file.Name <= prevName {
				return nil, fmt.Errorf("file names must be sorted and unique, but %q <= %q", file.Name, prevName)
			}
			prevName = file.Name

			// Validate the file name.
			if file.Name == "" {
				return nil, fmt.Errorf("file name must not be empty")
			}
			if strings.IndexByte(file.Name, '/') >= 0 {
				return nil, fmt.Errorf("file name %q contains path separators", file.Name)
			}

			// Check if the file is present in the CAS.
			fileDigest, err := digest.NewFromProto(file.Digest)
			if err != nil {
				return nil, err
			}
			if !cas.Has(fileDigest) {
				missingBlobs = append(missingBlobs, fileDigest)
			}

			// Construct a KajiyaFile for this file and add it to the directory.
			unixMode, err := convertUnixMode(file.GetNodeProperties().GetUnixMode().GetValue(), file.IsExecutable)
			if err != nil {
				return nil, fmt.Errorf("file %q: %w", file.Name, err)
			}
			kf := KajiyaFile{
				Name:     file.Name,
				Digest:   fileDigest,
				UnixMode: unixMode,
			}
			if file.GetNodeProperties().GetMtime() != nil {
				kf.Mtime = file.NodeProperties.Mtime.AsTime()
			}
			kd.Files = append(kd.Files, kf)
		}

		prevName = ""
		for _, symlink := range dir.Symlinks {
			if symlink.Name <= prevName {
				return nil, fmt.Errorf("symlink names must be sorted and unique, but %q <= %q", symlink.Name, prevName)
			}
			prevName = symlink.Name

			// Validate the symlink name and target.
			if symlink.Name == "" {
				return nil, fmt.Errorf("symlink name must not be empty")
			}
			if strings.IndexByte(symlink.Name, '/') >= 0 {
				return nil, fmt.Errorf("symlink name %q contains path separators", symlink.Name)
			}
			if filepath.IsAbs(symlink.Target) {
				return nil, fmt.Errorf("symlink %q points to absolute path %q", symlink.Name, symlink.Target)
			}
			if symlink.GetNodeProperties().GetMtime() != nil {
				return nil, fmt.Errorf("symlink %q specifies an mtime, which is unsupported", symlink.Name)
			}
			if symlink.GetNodeProperties().GetUnixMode() != nil {
				return nil, fmt.Errorf("symlink %q specifies a unix mode, which is unsupported", symlink.Name)
			}
			kd.Symlinks = append(kd.Symlinks, KajiyaSymlink{
				Name:   symlink.Name,
				Target: symlink.Target,
			})
		}

		prevName = ""
		for _, subDir := range dir.Directories {
			if subDir.Name <= prevName {
				return nil, fmt.Errorf("directory names must be sorted and unique, but %q <= %q", subDir.Name, prevName)
			}
			prevName = subDir.Name

			// Validate the subdirectory name and target.
			if subDir.Name == "" {
				return nil, fmt.Errorf("directory name must not be empty")
			}
			if strings.IndexByte(subDir.Name, '/') >= 0 {
				return nil, fmt.Errorf("directory name %q contains path separators", subDir.Name)
			}

			kd.Dirs = append(kd.Dirs, subDir.Name)

			// Prepend the parent directory's name so that we can insert the node with its full path into the trie.
			subDir.Name = filepath.ToSlash(filepath.Join(dirNode.Name, subDir.Name))
			dirQueue = append(dirQueue, subDir)
		}

		// Add the directory to the trie. Ensure that the directory name ends with a slash,
		// so that we can prefix-match paths safely in the trie later.
		if _, didUpdate := trieBuilder.Insert([]byte(filepath.ToSlash(dirNode.Name)+"/"), &kd); didUpdate {
			return nil, fmt.Errorf("duplicate directory name in tree %q", dirNode.Name)
		}
	}

	if len(missingBlobs) > 0 {
		return nil, &blobstore.MissingBlobsError{Blobs: missingBlobs}
	}

	return trieBuilder.Commit(), nil
}
