package main

import (
	"encoding/json"
	"fmt"
	"io"
	"os"
	"sort"

	"gopkg.in/yaml.v3"
)

// exit codes
const (
	exitOK        = 0
	exitZeroMatch = 1
	exitError     = 2
)

// queryResult is the stable JSON shape for selector queries.
type queryResult struct {
	Env      string  `json:"env"`
	Kind     string  `json:"kind"`
	Selector string  `json:"selector"`
	Matches  []match `json:"matches"`
}

// envProjection is the stable JSON shape when no selector is given.
type envProjection struct {
	Env           string      `json:"test_env"`
	Profile       string      `json:"profile"`
	LastUpdatedAt string      `json:"last_updated_at"`
	Source        interface{} `json:"source,omitempty"`
	Edge          interface{} `json:"edge,omitempty"`
	Build         interface{} `json:"build,omitempty"`
}

// match is a single path/value pair in a selector query result.
type match struct {
	Path  string      `json:"path"`
	Value interface{} `json:"value"`
}

// flagState holds parsed CLI flags.
type flagState struct {
	env      string
	model    string
	node     string
	provider string
	help     bool
}

func main() {
	os.Exit(run(os.Args[1:], os.Stdout, os.Stderr, loadInventory))
}

// run executes the query flow and returns the process exit code.
func run(args []string, stdout io.Writer, stderr io.Writer, loader func(string) (map[string]interface{}, error)) int {
	fs, err := parseFlags(args)
	if err != nil {
		fmt.Fprintln(stderr, err)
		return exitError
	}

	if fs.help {
		printUsage(stdout)
		return exitOK
	}

	data, err := loader(fs.env)
	if err != nil {
		fmt.Fprintln(stderr, err)
		return exitError
	}

	var selector string
	var kind string
	switch {
	case fs.node != "":
		selector = fs.node
		kind = "node"
	case fs.provider != "":
		selector = fs.provider
		kind = "provider"
	case fs.model != "":
		selector = fs.model
		kind = "model"
	}

	if err := validateInventory(data, fs.env, kind); err != nil {
		fmt.Fprintln(stderr, err)
		return exitError
	}

	if selector == "" {
		proj := buildEnvProjection(data)
		if err := encodeJSON(stdout, proj); err != nil {
			fmt.Fprintf(stderr, "json encode error: %v\n", err)
			return exitError
		}
		return exitOK
	}

	matches := queryInventory(data, kind, selector)
	if len(matches) == 0 {
		fmt.Fprintf(stderr, "no match for %s=%q\n", kind, selector)
		return exitZeroMatch
	}

	result := queryResult{
		Env:      fs.env,
		Kind:     kind,
		Selector: selector,
		Matches:  matches,
	}
	if err := encodeJSON(stdout, result); err != nil {
		fmt.Fprintf(stderr, "json encode error: %v\n", err)
		return exitError
	}
	return exitOK
}

// parseFlags parses CLI flags. --env is required (dev|dev-corp). At most one of --model, --node, --provider is allowed.
func parseFlags(args []string) (flagState, error) {
	var fs flagState
	var selectors int

	i := 0
	for i < len(args) {
		switch args[i] {
		case "--env":
			i++
			if i >= len(args) {
				return fs, fmt.Errorf("--env requires a value")
			}
			fs.env = args[i]
		case "--model":
			i++
			if i >= len(args) {
				return fs, fmt.Errorf("--model requires a value")
			}
			fs.model = args[i]
			selectors++
		case "--node":
			i++
			if i >= len(args) {
				return fs, fmt.Errorf("--node requires a value")
			}
			fs.node = args[i]
			selectors++
		case "--provider":
			i++
			if i >= len(args) {
				return fs, fmt.Errorf("--provider requires a value")
			}
			fs.provider = args[i]
			selectors++
		case "--help", "-h":
			fs.help = true
		default:
			return fs, fmt.Errorf("unknown flag: %s", args[i])
		}
		i++
	}

	if !fs.help {
		if fs.env == "" {
			return fs, fmt.Errorf("--env is required (dev|dev-corp)")
		}
		if fs.env != "dev" && fs.env != "dev-corp" {
			return fs, fmt.Errorf("invalid environment %q: --env must be dev or dev-corp", fs.env)
		}
		if selectors > 1 {
			return fs, fmt.Errorf("at most one of --model, --node, --provider is allowed")
		}
	}
	return fs, nil
}

func printUsage(w io.Writer) {
	fmt.Fprintln(w, "usage: inventory-query --env <env> [--model <model>] | [--node <node>] | [--provider <provider>]")
	fmt.Fprintln(w, "  --env       required. inventory environment: dev or dev-corp")
	fmt.Fprintln(w, "  --model     match a model alias/id")
	fmt.Fprintln(w, "  --node      match a node id or alias")
	fmt.Fprintln(w, "  --provider  match a provider id")
}

// loadInventory reads the inventory.yaml for the given env from agent-test/<env>/.
func loadInventory(env string) (map[string]interface{}, error) {
	path := fmt.Sprintf("agent-test/%s/inventory.yaml", env)
	data, err := os.ReadFile(path)
	if err != nil {
		return nil, fmt.Errorf("cannot read inventory for %s: %w", env, err)
	}
	var doc map[string]interface{}
	if err := yaml.Unmarshal(data, &doc); err != nil {
		return nil, fmt.Errorf("invalid YAML in %s: %w", path, err)
	}
	return doc, nil
}

// validateInventory validates the inventory structure.
func validateInventory(data map[string]interface{}, requestedEnv string, kind string) error {
	testEnv, ok := data["test_env"].(string)
	if !ok {
		return fmt.Errorf("missing or invalid type for test_env")
	}
	if testEnv != requestedEnv {
		return fmt.Errorf("env identity mismatch: requested %s, found %s", requestedEnv, testEnv)
	}

	if _, ok := data["profile"].(string); !ok {
		return fmt.Errorf("missing or invalid type for profile")
	}
	if _, ok := data["last_updated_at"].(string); !ok {
		return fmt.Errorf("missing or invalid type for last_updated_at")
	}
	if _, ok := data["source"].(map[string]interface{}); !ok {
		return fmt.Errorf("missing or invalid type for source")
	}
	if _, ok := data["edge"].(map[string]interface{}); !ok {
		return fmt.Errorf("missing or invalid type for edge")
	}
	if _, ok := data["build"].(map[string]interface{}); !ok {
		return fmt.Errorf("missing or invalid type for build")
	}

	if kind == "model" {
		m, ok := data["model"].(map[string]interface{})
		if !ok {
			return fmt.Errorf("missing or invalid type for model")
		}
		if _, ok := m["aliases"].(map[string]interface{}); !ok {
			return fmt.Errorf("missing or invalid type for model.aliases")
		}
	}

	if kind == "provider" {
		if _, ok := data["model"].(map[string]interface{}); !ok {
			return fmt.Errorf("missing or invalid type for model")
		}
		if _, ok := data["nodes"].([]interface{}); !ok {
			return fmt.Errorf("missing or invalid type for nodes")
		}
	}

	if kind == "node" {
		if _, ok := data["nodes"].([]interface{}); !ok {
			return fmt.Errorf("missing or invalid type for nodes")
		}
	}

	return nil
}

// buildEnvProjection returns the bounded env projection without model/nodes.
func buildEnvProjection(data map[string]interface{}) envProjection {
	proj := envProjection{}
	if v, ok := data["test_env"].(string); ok {
		proj.Env = v
	}
	if v, ok := data["profile"].(string); ok {
		proj.Profile = v
	}
	if v, ok := data["last_updated_at"].(string); ok {
		proj.LastUpdatedAt = v
	}
	if v, ok := data["source"]; ok {
		proj.Source = v
	}
	if v, ok := data["edge"]; ok {
		proj.Edge = v
	}
	if v, ok := data["build"]; ok {
		proj.Build = v
	}
	return proj
}

// queryInventory searches the inventory for matches of the given kind and selector.
func queryInventory(data map[string]interface{}, kind, selector string) []match {
	var matches []match
	switch kind {
	case "node":
		matches = queryNodes(data, selector)
	case "provider":
		matches = queryProviders(data, selector)
	case "model":
		matches = queryModels(data, selector)
	}
	sortMatches(matches)
	return matches
}

// queryNodes finds nodes by id or alias at top-level nodes[].
func queryNodes(data map[string]interface{}, selector string) []match {
	var matches []match
	nodes, ok := data["nodes"].([]interface{})
	if !ok {
		return matches
	}
	for i, n := range nodes {
		obj, ok := n.(map[string]interface{})
		if !ok {
			continue
		}
		path := fmt.Sprintf("nodes[%d]", i)
		if matchString(obj, "id", selector) || matchString(obj, "alias", selector) {
			matches = append(matches, match{Path: path, Value: obj})
		}
	}
	sortMatches(matches)
	return matches
}

// queryProviders finds providers by id in any canonical provider shapes.
func queryProviders(data map[string]interface{}, selector string) []match {
	collector := newProviderCollector(selector)
	collector.collect(data, "")
	sortMatches(collector.matches)
	return collector.matches
}

func joinPath(base, elem string) string {
	if base == "" {
		return elem
	}
	return base + "." + elem
}

// queryModels finds models by alias key in model.aliases or by id/alias/model_id in model subtree entities.
func queryModels(data map[string]interface{}, selector string) []match {
	var matches []match
	seen := make(map[string]bool)

	model, ok := data["model"].(map[string]interface{})
	if !ok {
		return matches
	}

	// Check model.aliases keys.
	if aliases, ok := model["aliases"].(map[string]interface{}); ok {
		for key, val := range aliases {
			if key == selector && !seen["model.aliases."+key] {
				matches = append(matches, match{Path: "model.aliases." + key, Value: val})
				seen["model.aliases."+key] = true
			}
		}
	}

	// Check model.subtree entities recursively for id/alias/model_id.
	entityMatches := collectEntityMatches(model, "model", selector, seen)
	matches = append(matches, entityMatches...)

	sortMatches(matches)
	return matches
}

// collectEntityMatches recursively searches a subtree for entities matching id/alias/model_id.
func collectEntityMatches(node interface{}, basePath, selector string, seen map[string]bool) []match {
	var matches []match

	switch n := node.(type) {
	case map[string]interface{}:
		if id, ok := n["id"].(string); ok && id == selector {
			key := basePath + ".id"
			if !seen[key] {
				matches = append(matches, match{Path: key, Value: id})
				seen[key] = true
			}
		}
		if a, ok := n["alias"].(string); ok && a == selector {
			key := basePath + ".alias"
			if !seen[key] {
				matches = append(matches, match{Path: key, Value: a})
				seen[key] = true
			}
		}
		if mid, ok := n["model_id"].(string); ok && mid == selector {
			key := basePath + ".model_id"
			if !seen[key] {
				matches = append(matches, match{Path: key, Value: mid})
				seen[key] = true
			}
		}

		for k, v := range n {
			childPath := basePath + "." + k
			matches = append(matches, collectEntityMatches(v, childPath, selector, seen)...)
		}

	case []interface{}:
		for i, v := range n {
			childPath := fmt.Sprintf("%s[%d]", basePath, i)
			matches = append(matches, collectEntityMatches(v, childPath, selector, seen)...)
		}
	}

	return matches
}

// matchString checks if the given key in obj has the exact string value.
func matchString(obj map[string]interface{}, key, value string) bool {
	if v, ok := obj[key].(string); ok && v == value {
		return true
	}
	return false
}

// sortMatches sorts matches by path ascending.
func sortMatches(matches []match) {
	sort.Slice(matches, func(i, j int) bool {
		return matches[i].Path < matches[j].Path
	})
}

// encodeJSON encodes v as indented JSON to w.
func encodeJSON(w io.Writer, v interface{}) error {
	enc := json.NewEncoder(w)
	enc.SetIndent("", "  ")
	enc.SetEscapeHTML(false)
	return enc.Encode(v)
}