package main

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"

	"gopkg.in/yaml.v3"
)

func envReplace(b []byte, env map[string]string) []byte {
	for k, v := range env {
		b = bytes.ReplaceAll(b, []byte("${"+k+"}"), []byte(v))
	}
	return b
}

type YAMLFragment struct {
	content *yaml.Node
	baseDir string
	env     map[string]string
}

func (f *YAMLFragment) UnmarshalYAML(value *yaml.Node) error {
	var err error
	// process includes in fragments
	f.content, err = yamlResolveIncludes(value, f.baseDir, f.env)
	return err
}

type YAMLIncludeProcessor struct {
	target  interface{}
	baseDir string
	env     map[string]string
}

func (i *YAMLIncludeProcessor) UnmarshalYAML(value *yaml.Node) error {
	resolved, err := yamlResolveIncludes(value, i.baseDir, i.env)
	if err != nil {
		return err
	}
	return resolved.Decode(i.target)
}

func yamlResolveIncludes(node *yaml.Node, baseDir string, env map[string]string) (*yaml.Node, error) {
	if node.Tag == "!include" {
		if node.Kind != yaml.ScalarNode {
			return nil, errors.New("!include on a non-scalar node")
		}

		filename := baseDir + "/" + node.Value
		file, err := ioutil.ReadFile(filename)
		if err != nil {
			return nil, err
		}
		var f YAMLFragment
		f.baseDir = filepath.Dir(filename)
		f.env = env
		file = envReplace(file, env)
		err = yaml.Unmarshal(file, &f)
		if err != nil {
			err = fmt.Errorf("%s (included file: %s)", err, filename)
		}
		return f.content, err
	}
	if node.Kind == yaml.SequenceNode || node.Kind == yaml.MappingNode {
		var err error
		for i := range node.Content {
			node.Content[i], err = yamlResolveIncludes(node.Content[i], baseDir, env)
			if err != nil {
				return nil, err
			}
		}
	}
	return node, nil
}

func main() {

	var output interface{}

	argsWithoutProg := os.Args[1:]

	if len(argsWithoutProg) < 2 {
		panic("require input and output")
	}

	inputFilename := argsWithoutProg[0]
	outputFilename := argsWithoutProg[1]

	log.Println("processing " + inputFilename + " -> " + outputFilename)

	bIn, err := ioutil.ReadFile(inputFilename)
	if err != nil {
		log.Fatalln(err)
	}

	err = yaml.Unmarshal(bIn, &YAMLIncludeProcessor{
		target:  &output,
		baseDir: filepath.Dir(inputFilename),
		env:     map[string]string{},
	})
	if err != nil {
		log.Fatalln(err)
	}

	bOut, err := json.MarshalIndent(output, "", "    ")
	if err != nil {
		log.Fatalln(err)
	}
	err = ioutil.WriteFile(outputFilename, bOut, 0644)
	if err != nil {
		log.Fatalln(err)
	}
}