package main

import (
	"encoding/json"
	"fmt"
	"io"
	"os"
	"path/filepath"

	jsonnet "github.com/google/go-jsonnet"
	"github.com/google/go-jsonnet/formatter"
)

// Request supports two modes: "evaluate" (default) and "format".
type Request struct {
	Mode     string            `json:"mode,omitempty"`
	Filename string            `json:"filename"`
	Snippet  string            `json:"snippet"`
	ExtVars  map[string]string `json:"ext_vars,omitempty"`
	TLAVars  map[string]string `json:"tla_vars,omitempty"`
	LibFiles map[string]string `json:"lib_files,omitempty"`
	JPaths   []string          `json:"jpaths,omitempty"`

	// Format options (only used when mode == "format")
	FormatOptions *FormatOptions `json:"format_options,omitempty"`
}

type FormatOptions struct {
	Indent           *int    `json:"indent,omitempty"`
	MaxBlankLines    *int    `json:"max_blank_lines,omitempty"`
	StringStyle      *string `json:"string_style,omitempty"`  // "double", "single", "leave"
	CommentStyle     *string `json:"comment_style,omitempty"` // "hash", "slash", "leave"
	PrettyFieldNames *bool   `json:"pretty_field_names,omitempty"`
	PadArrays        *bool   `json:"pad_arrays,omitempty"`
	PadObjects       *bool   `json:"pad_objects,omitempty"`
	SortImports      *bool   `json:"sort_imports,omitempty"`
	UseImplicitPlus  *bool   `json:"use_implicit_plus,omitempty"`
}

type Response struct {
	Result string `json:"result,omitempty"`
	Error  string `json:"error,omitempty"`
}

// ── Memory Importer ──

type memoryImporter struct {
	files map[string]jsonnet.Contents
}

func newMemoryImporter(files map[string]string) *memoryImporter {
	m := &memoryImporter{files: make(map[string]jsonnet.Contents, len(files))}
	for path, content := range files {
		m.files[filepath.Clean(path)] = jsonnet.MakeContents(content)
	}
	return m
}

func (m *memoryImporter) Import(importedFrom, importedPath string) (jsonnet.Contents, string, error) {
	var resolved string
	if filepath.IsAbs(importedPath) {
		resolved = filepath.Clean(importedPath)
	} else {
		dir := filepath.Dir(importedFrom)
		if dir == "." || dir == "" {
			resolved = filepath.Clean(importedPath)
		} else {
			resolved = filepath.Clean(filepath.Join(dir, importedPath))
		}
	}
	if contents, ok := m.files[resolved]; ok {
		return contents, resolved, nil
	}
	clean := filepath.Clean(importedPath)
	if contents, ok := m.files[clean]; ok {
		return contents, clean, nil
	}
	return jsonnet.Contents{}, "", fmt.Errorf("import %q not found in lib_files (resolved: %q)", importedPath, resolved)
}

// ── Main ──

func main() {
	input, err := io.ReadAll(os.Stdin)
	if err != nil {
		writeResponse(Response{Error: fmt.Sprintf("failed to read stdin: %s", err)})
		return
	}

	var req Request
	if err := json.Unmarshal(input, &req); err != nil {
		req = Request{Filename: "input.jsonnet", Snippet: string(input)}
	}
	if req.Filename == "" {
		req.Filename = "input.jsonnet"
	}

	switch req.Mode {
	case "format":
		doFormat(req)
	default:
		doEvaluate(req)
	}
}

func doEvaluate(req Request) {
	vm := jsonnet.MakeVM()
	if len(req.LibFiles) > 0 {
		vm.Importer(newMemoryImporter(req.LibFiles))
	} else if len(req.JPaths) > 0 {
		vm.Importer(&jsonnet.FileImporter{JPaths: req.JPaths})
	}
	for k, v := range req.ExtVars {
		vm.ExtVar(k, v)
	}
	for k, v := range req.TLAVars {
		vm.TLAVar(k, v)
	}

	result, err := vm.EvaluateAnonymousSnippet(req.Filename, req.Snippet)
	if err != nil {
		writeResponse(Response{Error: err.Error()})
		return
	}
	writeResponse(Response{Result: result})
}

func doFormat(req Request) {
	opts := formatter.DefaultOptions()
	if fo := req.FormatOptions; fo != nil {
		if fo.Indent != nil {
			opts.Indent = *fo.Indent
		}
		if fo.MaxBlankLines != nil {
			opts.MaxBlankLines = *fo.MaxBlankLines
		}
		if fo.StringStyle != nil {
			switch *fo.StringStyle {
			case "double":
				opts.StringStyle = formatter.StringStyleDouble
			case "single":
				opts.StringStyle = formatter.StringStyleSingle
			case "leave":
				opts.StringStyle = formatter.StringStyleLeave
			}
		}
		if fo.CommentStyle != nil {
			switch *fo.CommentStyle {
			case "hash":
				opts.CommentStyle = formatter.CommentStyleHash
			case "slash":
				opts.CommentStyle = formatter.CommentStyleSlash
			case "leave":
				opts.CommentStyle = formatter.CommentStyleLeave
			}
		}
		if fo.PrettyFieldNames != nil {
			opts.PrettyFieldNames = *fo.PrettyFieldNames
		}
		if fo.PadArrays != nil {
			opts.PadArrays = *fo.PadArrays
		}
		if fo.PadObjects != nil {
			opts.PadObjects = *fo.PadObjects
		}
		if fo.SortImports != nil {
			opts.SortImports = *fo.SortImports
		}
		if fo.UseImplicitPlus != nil {
			opts.UseImplicitPlus = *fo.UseImplicitPlus
		}
	}

	result, err := formatter.Format(req.Filename, req.Snippet, opts)
	if err != nil {
		writeResponse(Response{Error: err.Error()})
		return
	}
	writeResponse(Response{Result: result})
}

func writeResponse(resp Response) {
	out, _ := json.Marshal(resp)
	os.Stdout.Write(out)
}