package main

import (
	"flag"
	"fmt"
	"os"
	"path/filepath"
	"strings"

	"github.com/google/go-jsonnet"
	"forgejo.csbx.dev/acmcarther/yesod/tools/jsonnet_compiler/helm_support"
)

// stringList handles repeated string flags (e.g. -J)
type stringList []string

func (s *stringList) String() string {
	return strings.Join(*s, ",")
}

func (s *stringList) Set(value string) error {
	*s = append(*s, value)
	return nil
}

// keyValueMap handles repeated key=value flags (e.g. --ext-str)
type keyValueMap map[string]string

func (m *keyValueMap) String() string {
	return fmt.Sprintf("%v", *m)
}

func (m *keyValueMap) Set(value string) error {
	parts := strings.SplitN(value, "=", 2)
	if len(parts) != 2 {
		return fmt.Errorf("expected key=value, got %s", value)
	}
	(*m)[parts[0]] = parts[1]
	return nil
}

func main() {
	vm := jsonnet.MakeVM()
	vm.NativeFunction(helm_support.NativeFunc())

	// Flags compatible with rules_jsonnet
	var (
		importPaths stringList
		extStr      = make(keyValueMap)
		extCode     = make(keyValueMap)
		tlaStr      = make(keyValueMap)
		tlaCode     = make(keyValueMap)
		outputFile  string
	)

	fs := flag.NewFlagSet("jsonnet", flag.ExitOnError)
	fs.Var(&importPaths, "J", "Add to the library search path")
	fs.Var(&importPaths, "jpath", "Add to the library search path")
	fs.Var(&extStr, "ext-str", "Provide external variable as string")
	fs.Var(&extStr, "V", "Provide external variable as string")
	fs.Var(&extCode, "ext-code", "Provide external variable as code")
	fs.Var(&tlaStr, "tla-str", "Provide top-level argument as string")
	fs.Var(&tlaStr, "A", "Provide top-level argument as string")
	fs.Var(&tlaCode, "tla-code", "Provide top-level argument as code")
	fs.StringVar(&outputFile, "o", "", "Output to the named file")
	fs.StringVar(&outputFile, "output-file", "", "Output to the named file")

	// Parse flags
	if err := fs.Parse(os.Args[1:]); err != nil {
		fmt.Fprintf(os.Stderr, "Error parsing flags: %v\n", err)
		os.Exit(1)
	}

	// Apply configuration to VM
	vm.Importer(&jsonnet.FileImporter{
		JPaths: importPaths,
	})

	for k, v := range extStr {
		vm.ExtVar(k, v)
	}
	for k, v := range extCode {
		vm.ExtCode(k, v)
	}
	for k, v := range tlaStr {
		vm.TLAVar(k, v)
	}
	for k, v := range tlaCode {
		vm.TLACode(k, v)
	}

	// Get input file and handle remaining flags (like -o after input file)
	args := fs.Args()
	var inputFile string
	for i := 0; i < len(args); i++ {
		arg := args[i]
		if arg == "-o" || arg == "--output-file" || arg == "-output-file" {
			if i+1 < len(args) {
				outputFile = args[i+1]
				i++
				continue
			}
		}
		if inputFile == "" && !strings.HasPrefix(arg, "-") {
			inputFile = arg
		}
	}

	if inputFile == "" {
		fmt.Fprintln(os.Stderr, "Usage: jsonnet_compiler [options] <filename>")
		os.Exit(1)
	}

	// Evaluate
	jsonOutput, err := vm.EvaluateFile(inputFile)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error evaluating snippet: %v\n", err)
		os.Exit(1)
	}

	// Write output
	if outputFile != "" {
		if err := os.MkdirAll(filepath.Dir(outputFile), 0755); err != nil {
            fmt.Fprintf(os.Stderr, "Error creating output directory: %v\n", err)
			os.Exit(1)
        }
		if err := os.WriteFile(outputFile, []byte(jsonOutput), 0644); err != nil {
			fmt.Fprintf(os.Stderr, "Error writing output file: %v\n", err)
			os.Exit(1)
		}
	} else {
		fmt.Println(jsonOutput)
	}
}