package exec

import (
	"encoding/json"
	"fmt"
	"sort"
)

const minProjectionDefinitionBytes = 24

type projectionSchemaCandidate struct {
	key   string
	name  string
	node  map[string]any
	size  int
	count int
}

// compactProjectionSchema preserves the validated action schema language while
// replacing repeated object subschemas with action-local $defs. Catalog
// validation continues to use the fully expanded descriptor schema.
func compactProjectionSchema(raw json.RawMessage) (json.RawMessage, error) {
	document, err := decodeStrict(raw)
	if err != nil {
		return nil, err
	}
	root, ok := document.(map[string]any)
	if !ok {
		return nil, fmt.Errorf("schema root must be an object")
	}

	candidatesByKey := make(map[string]*projectionSchemaCandidate)
	collectProjectionSchemaCandidates(root, true, candidatesByKey)
	candidates := make([]*projectionSchemaCandidate, 0, len(candidatesByKey))
	for _, candidate := range candidatesByKey {
		refBytes := len(`{"$ref":"#/$defs/S000"}`)
		savings := (candidate.count-1)*candidate.size - candidate.count*refBytes
		if candidate.count > 1 && candidate.size >= minProjectionDefinitionBytes && savings > 0 {
			candidates = append(candidates, candidate)
		}
	}
	sort.Slice(candidates, func(i, j int) bool {
		if candidates[i].size != candidates[j].size {
			return candidates[i].size > candidates[j].size
		}
		return candidates[i].key < candidates[j].key
	})
	selected := make(map[string]*projectionSchemaCandidate, len(candidates))
	for index, candidate := range candidates {
		candidate.name = fmt.Sprintf("S%03d", index+1)
		selected[candidate.key] = candidate
	}

	projected, _ := projectSchemaNode(root, "", selected).(map[string]any)
	definitions := make(map[string]any, len(candidates))
	for _, candidate := range candidates {
		definitions[candidate.name] = projectSchemaNode(candidate.node, candidate.key, selected)
	}
	reachable := reachableProjectionDefinitions(projected, definitions)
	if len(reachable) > 0 {
		used := make(map[string]any, len(reachable))
		for name := range reachable {
			used[name] = definitions[name]
		}
		projected["$defs"] = used
	}
	encoded, err := json.Marshal(projected)
	if err != nil {
		return nil, err
	}
	return encoded, nil
}

func collectProjectionSchemaCandidates(
	node any,
	root bool,
	candidates map[string]*projectionSchemaCandidate,
) {
	switch value := node.(type) {
	case map[string]any:
		if !root {
			key := canonicalProjectionSchemaNode(value)
			candidate := candidates[key]
			if candidate == nil {
				candidate = &projectionSchemaCandidate{
					key: key, node: value, size: len(key),
				}
				candidates[key] = candidate
			}
			candidate.count++
		}
		for _, keyword := range namedSchemaMapKeywords {
			children, _ := value[keyword].(map[string]any)
			for _, child := range children {
				collectProjectionSchemaCandidates(child, false, candidates)
			}
		}
		for _, keyword := range singleSchemaKeywords {
			if child, exists := value[keyword]; exists {
				collectProjectionSchemaCandidates(child, false, candidates)
			}
		}
		for _, keyword := range schemaArrayKeywords {
			children, _ := value[keyword].([]any)
			for _, child := range children {
				collectProjectionSchemaCandidates(child, false, candidates)
			}
		}
	}
}

func projectSchemaNode(node any, skipKey string, selected map[string]*projectionSchemaCandidate) any {
	switch value := node.(type) {
	case map[string]any:
		key := canonicalProjectionSchemaNode(value)
		if key != skipKey {
			if candidate := selected[key]; candidate != nil {
				return map[string]any{"$ref": "#/$defs/" + candidate.name}
			}
		}
		out := make(map[string]any, len(value))
		for name, child := range value {
			out[name] = child
		}
		for _, keyword := range namedSchemaMapKeywords {
			children, _ := value[keyword].(map[string]any)
			if len(children) == 0 {
				continue
			}
			projected := make(map[string]any, len(children))
			for name, child := range children {
				projected[name] = projectSchemaNode(child, skipKey, selected)
			}
			out[keyword] = projected
		}
		for _, keyword := range singleSchemaKeywords {
			if child, exists := value[keyword]; exists {
				out[keyword] = projectSchemaNode(child, skipKey, selected)
			}
		}
		for _, keyword := range schemaArrayKeywords {
			children, _ := value[keyword].([]any)
			if len(children) == 0 {
				continue
			}
			projected := make([]any, len(children))
			for index, child := range children {
				projected[index] = projectSchemaNode(child, skipKey, selected)
			}
			out[keyword] = projected
		}
		return out
	default:
		return value
	}
}

func reachableProjectionDefinitions(root map[string]any, definitions map[string]any) map[string]bool {
	reachable := make(map[string]bool)
	var visit func(any)
	visit = func(node any) {
		switch value := node.(type) {
		case map[string]any:
			if ref, ok := value["$ref"].(string); ok {
				const prefix = "#/$defs/"
				if len(ref) > len(prefix) && ref[:len(prefix)] == prefix {
					name := ref[len(prefix):]
					if !reachable[name] {
						reachable[name] = true
						visit(definitions[name])
					}
				}
			}
			for _, child := range value {
				visit(child)
			}
		case []any:
			for _, child := range value {
				visit(child)
			}
		}
	}
	visit(root)
	return reachable
}

func canonicalProjectionSchemaNode(node map[string]any) string {
	raw, _ := json.Marshal(node)
	return string(raw)
}