package lattice

import (
	"fmt"
	"sort"
	"strings"

	"github.com/RoaringBitmap/roaring"
	"github.com/agentic-research/mache/api"
)

// friendlyTypeNames maps tree-sitter node types to human-friendly group names.
var friendlyTypeNames = map[string]string{
	// Go
	"function_declaration": "functions",
	"method_declaration":   "methods",
	"type_declaration":     "types",
	// Python
	"function_definition": "functions",
	"class_definition":    "classes",
	// JavaScript / TypeScript
	"class_declaration": "classes",
	"method_definition": "methods",
	// Rust
	"function_item": "functions",
	"struct_item":   "structs",
	"enum_item":     "enums",
	"impl_item":     "implementations",
	"trait_item":    "traits",
	// Elixir
	"call": "definitions",
	// HCL/Terraform
	"hcl_container": "blocks",
	// SQL
	"create_table": "tables",
	"create_view":  "views",
	// HTML / DOM
	"element":        "elements",
	"script_element": "scripts",
	"style_element":  "styles",
	// CSS
	"rule_set":            "rules",
	"media_statement":     "media",
	"keyframes_statement": "keyframes",
	"import_statement":    "imports",
}

// containmentRules defines which AST container types can nest inside others,
// per language. Keys are parent types; values are allowed child types.
// Languages not listed (or with empty maps) default to flat — no nesting.
var containmentRules = map[string]map[string][]string{
	"go":        {}, // flat — no nested named functions
	"hcl":       {}, // flat
	"terraform": {}, // flat
	"rust":      {}, // flat at top level
	"elixir": {
		"call": {"call"}, // modules contain function definitions
	},
	"python": {
		"class_definition": {"function_definition"},
	},
	"javascript": {
		"class_declaration": {"method_definition"},
	},
	"typescript": {
		"class_declaration": {"method_definition"},
	},
}

// ProjectAST converts Formal Concepts from a flattened AST into a recursive schema.
//
// Strategy:
// 1. Identify "Container" types: Nodes that have a 'name' (identifier) and a 'body' (block).
//   - e.g. function_definition, class_definition
//
// 2. Create a Schema Node for each Container type.
// 3. Make the schema recursive: Every container can contain every other container type.
//   - This covers classes inside functions, functions inside classes, etc.
//
// 4. Default to mapping "source" -> content of the node.
func ProjectAST(concepts []Concept, ctx *FormalContext, config ProjectConfig) *api.Topology {
	containerTypes := make(map[string]string) // type -> name_type (e.g. identifier, property_identifier)

	// Find all unique types present in the data
	types := make(map[string]bool)
	for _, attr := range ctx.Attributes {
		if attr.Name == "type" || strings.HasPrefix(attr.Name, "type=") {
			// Extract type name from "type=foo"
			t := strings.TrimPrefix(attr.Name, "type=")
			types[t] = true
		}
	}

	// Use lattice closure to check if type=X implies has_name+has_body.
	for t := range types {
		// Find the attribute index for "type=t"
		typeAttrIdx := -1
		for j, attr := range ctx.Attributes {
			if attr.Name == "type="+t {
				typeAttrIdx = j
				break
			}
		}
		if typeAttrIdx == -1 {
			continue
		}

		// Check if this type has a name and body
		// We look at the concept generated by this attribute
		// closure({type=t}) -> The set of all attributes that objects of this type share.
		closure := ctx.Closure(roaring.BitmapOf(uint32(typeAttrIdx)))

		hasName := false
		hasBody := false
		nameType := "identifier" // default

		iter := closure.Iterator()
		for iter.HasNext() {
			idx := iter.Next()
			attrName := ctx.Attributes[idx].Name
			if attrName == "has_name" {
				hasName = true
			}
			if after, ok := strings.CutPrefix(attrName, "field_name_type="); ok {
				nameType = after
			}
			if attrName == "has_body" {
				hasBody = true
			}
		}

		if hasName && hasBody {
			containerTypes[t] = nameType
		}
	}

	// Sort types for deterministic output
	var sortedTypes []string
	for t := range containerTypes {
		sortedTypes = append(sortedTypes, t)
	}
	sort.Strings(sortedTypes)

	// Build the recursive schema
	// All containers are children of the root AND children of each other.
	var children []api.Node

	// Create a node for each container type
	for _, t := range sortedTypes {
		nameType := containerTypes[t]

		// The Engine's backward scan (extractDocComments) walks @scope's preceding
		// siblings to find contiguous comment nodes. It injects the result as a
		// "doc" template variable and extends the "source" file content to include
		// them. No complex tree-sitter query needed — the simple selector works.
		selector := fmt.Sprintf("(%s name: (%s) @name) @scope", t, nameType)

		node := api.Node{
			Name:          "{{.name}}",
			Selector:      selector,
			SkipSelfMatch: true,
			Language:      config.Language,
			Files: []api.Leaf{
				{
					Name:            "source",
					ContentTemplate: "{{.scope}}",
				},
				{
					Name:            "ast.json",
					ContentTemplate: "{{. | json}}",
				},
				{
					Name:            "doc",
					ContentTemplate: "{{.doc}}",
				},
			},
		}
		children = append(children, node)
	}

	// 1 level of nesting via containmentRules. Go structs can't be cyclic and
	// JSON can't encode references, so we stop at depth 1.

	// Language-aware containment: only assign children that the language permits.
	// Unknown languages default to flat (safe — misses nesting but never generates garbage).
	lang := strings.ToLower(config.Language)
	rules := containmentRules[lang] // nil map is fine — all lookups return nil

	// Build a lookup from type name → index in children slice
	typeIndex := make(map[string]int, len(sortedTypes))
	for i, t := range sortedTypes {
		typeIndex[t] = i
	}

	// Assign permitted children (1 level of nesting — sufficient for real languages).
	// Root → [Func, Class], Class → [Method] (per rules), Func → [] (flat).
	result := make([]api.Node, len(children))
	copy(result, children)
	for i, t := range sortedTypes {
		allowed := rules[t]
		if len(allowed) == 0 {
			continue
		}
		var nested []api.Node
		for _, childType := range allowed {
			if j, ok := typeIndex[childType]; ok {
				nested = append(nested, children[j])
			}
		}
		if len(nested) > 0 {
			result[i].Children = nested
		}
	}

	// Wrap each container type in a friendly-name grouping directory.
	// e.g. function_definition → "functions/", class_definition → "classes/"
	var grouped []api.Node
	groupIndex := make(map[string]int) // friendly_name → index in grouped

	for i, t := range sortedTypes {
		friendly := friendlyTypeNames[t]
		if friendly == "" {
			friendly = t // fallback: use raw type name
		}
		if idx, ok := groupIndex[friendly]; ok {
			grouped[idx].Children = append(grouped[idx].Children, result[i])
		} else {
			group := api.Node{
				Name:     friendly,
				Selector: "$",
				Language: config.Language,
				Children: []api.Node{result[i]},
			}
			groupIndex[friendly] = len(grouped)
			grouped = append(grouped, group)
		}
	}

	return &api.Topology{
		Version: "v1",
		Nodes:   grouped,
	}
}
