package main

import (
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"io/fs"
	"os"
	"path/filepath"
	"strings"

	sitter "github.com/smacker/go-tree-sitter"
	"github.com/smacker/go-tree-sitter/golang"
	"github.com/smacker/go-tree-sitter/javascript"
	"github.com/smacker/go-tree-sitter/python"
	"github.com/smacker/go-tree-sitter/typescript/typescript"
)

type FileInfo struct {
	Language string   `json:"language"`
	Exports  []string `json:"exports,omitempty"`
	Defines  []Symbol `json:"defines,omitempty"`
}

type Symbol struct {
	Name string `json:"name"`
	Type string `json:"type"`
	Line uint32 `json:"line"`
}

type Index struct {
	Files   map[string]FileInfo  `json:"files"`
	Symbols map[string]SymbolRef `json:"symbols"`
	Graph   map[string][]string  `json:"dependencyGraph"`
}

type SymbolRef struct {
	File string `json:"file"`
	Line uint32 `json:"line"`
	Type string `json:"type"`
}

var (
	ignoreDirs = map[string]bool{
		"node_modules": true,
		".git":         true,
		"target":       true,
		"dist":         true,
		"build":        true,
		".clj-kondo":   true,
		".shadow-cljs": true,
		".cpcache":     true,
	}

	extToLang = map[string]string{
		".js":    "javascript",
		".jsx":   "javascript",
		".ts":    "typescript",
		".tsx":   "typescript",
		".astro": "astro",
		".cljs":  "clojurescript",
		".clj":   "clojure",
		".cljc":  "clojure",
		".go":    "go",
		".py":    "python",
	}
)

// parseResult is the internal result before we trim for output.
type parseResult struct {
	Imports []string
	Exports []string
	Defines []Symbol
}

func main() {
	root := flag.String("root", ".", "Project root directory")
	output := flag.String("out", ".context/relationships.json", "Output file")
	flag.Parse()

	index := &Index{
		Files:   make(map[string]FileInfo),
		Symbols: make(map[string]SymbolRef),
		Graph:   make(map[string][]string),
	}

	parser := sitter.NewParser()

	err := filepath.WalkDir(*root, func(path string, d fs.DirEntry, err error) error {
		if err != nil {
			return nil
		}

		if d.IsDir() {
			if ignoreDirs[d.Name()] {
				return filepath.SkipDir
			}
			return nil
		}

		ext := filepath.Ext(path)
		lang, ok := extToLang[ext]
		if !ok {
			return nil
		}

		relPath, _ := filepath.Rel(*root, path)
		content, err := os.ReadFile(path)
		if err != nil {
			return nil
		}

		result := parseFile(parser, content, lang)

		// Deduplicate
		result.Imports = unique(result.Imports)
		result.Exports = unique(result.Exports)
		result.Defines = uniqueSymbols(result.Defines)

		// Build file entry — only include defines if file has exports
		fileInfo := FileInfo{Language: lang}
		fileInfo.Exports = result.Exports
		if len(result.Exports) > 0 {
			fileInfo.Defines = result.Defines
		}
		index.Files[relPath] = fileInfo

		// Global symbol index — only exported/module-level symbols
		exportSet := toSet(result.Exports)
		for _, sym := range result.Defines {
			if exportSet[sym.Name] || sym.Type == "class" || sym.Type == "struct" || sym.Type == "interface" || sym.Type == "type" || sym.Type == "protocol" || sym.Type == "multimethod" {
				index.Symbols[sym.Name] = SymbolRef{
					File: relPath,
					Line: sym.Line,
					Type: sym.Type,
				}
			}
		}

		// Dependency graph
		if len(result.Imports) > 0 {
			index.Graph[relPath] = result.Imports
		}

		return nil
	})

	if err != nil {
		fmt.Fprintf(os.Stderr, "Error walking directory: %v\n", err)
		os.Exit(1)
	}

	// Ensure output directory exists
	outDir := filepath.Dir(*output)
	if outDir != "." {
		os.MkdirAll(outDir, 0755)
	}

	jsonOut, _ := json.MarshalIndent(index, "", "  ")
	err = os.WriteFile(*output, jsonOut, 0644)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error writing output: %v\n", err)
		os.Exit(1)
	}

	fmt.Printf("Indexed %d files, %d symbols\n", len(index.Files), len(index.Symbols))
}

func toSet(slice []string) map[string]bool {
	s := make(map[string]bool, len(slice))
	for _, v := range slice {
		s[v] = true
	}
	return s
}

func parseFile(parser *sitter.Parser, content []byte, lang string) parseResult {
	var treeSitterLang *sitter.Language

	switch lang {
	case "javascript":
		treeSitterLang = javascript.GetLanguage()
	case "typescript":
		treeSitterLang = typescript.GetLanguage()
	case "astro":
		return parseAstro(parser, content)
	case "clojure", "clojurescript":
		return parseResult{}
	case "go":
		treeSitterLang = golang.GetLanguage()
	case "python":
		treeSitterLang = python.GetLanguage()
	default:
		return parseResult{}
	}

	parser.SetLanguage(treeSitterLang)
	tree, err := parser.ParseCtx(context.Background(), nil, content)
	if err != nil {
		return parseResult{}
	}
	defer tree.Close()

	root := tree.RootNode()

	switch lang {
	case "javascript", "typescript":
		return parseJSTS(root, content)
	case "clojure", "clojurescript":
		return parseClojure(root, content)
	case "go":
		return parseGo(root, content)
	case "python":
		return parsePython(root, content)
	}

	return parseResult{}
}

// isTopLevel checks if a node is a direct child of the program/module root.
func isTopLevel(node *sitter.Node) bool {
	parent := node.Parent()
	if parent == nil {
		return true
	}
	t := parent.Type()
	return t == "program" || t == "module" || t == "source_file" || t == "export_statement"
}

func parseJSTS(root *sitter.Node, content []byte) parseResult {
	var result parseResult

	var walk func(*sitter.Node)
	walk = func(node *sitter.Node) {
		switch node.Type() {
		case "import_statement":
			for i := 0; i < int(node.ChildCount()); i++ {
				child := node.Child(i)
				if child.Type() == "string" {
					importPath := strings.Trim(child.Content(content), "\"'")
					result.Imports = append(result.Imports, importPath)
				}
			}

		case "export_statement":
			hasDefault := false
			for i := 0; i < int(node.ChildCount()); i++ {
				child := node.Child(i)
				switch child.Type() {
				case "identifier":
					// export default SomeName or export SomeName
					result.Exports = append(result.Exports, child.Content(content))
				case "function_declaration", "class_declaration":
					// export function foo / export class C
					walk(child)
					nameNode := child.ChildByFieldName("name")
					if nameNode != nil {
						result.Exports = append(result.Exports, nameNode.Content(content))
					}
				case "lexical_declaration", "variable_declaration":
					// export const x = ... / export let y = ...
					walk(child)
					for j := 0; j < int(child.ChildCount()); j++ {
						decl := child.Child(j)
						if decl.Type() == "variable_declarator" {
							nameNode := decl.ChildByFieldName("name")
							if nameNode != nil && nameNode.Type() == "identifier" {
								result.Exports = append(result.Exports, nameNode.Content(content))
							}
						}
					}
				case "export_clause":
					// export { name1, name2 } or export { name1 as alias }
					extractExportClause(child, content, &result)
				case "type_alias_declaration", "interface_declaration":
					// export type Foo = ... / export interface Bar { ... }
					nameNode := child.ChildByFieldName("name")
					if nameNode != nil {
						name := nameNode.Content(content)
						result.Exports = append(result.Exports, name)
						result.Defines = append(result.Defines, Symbol{
							Name: name,
							Type: child.Type()[:len(child.Type())-len("_declaration")],
							Line: child.StartPoint().Row + 1,
						})
					}
				case "default":
					hasDefault = true
				}
			}
			// For "export default function name" / "export default class name",
			// the function/class name was already captured by the walk(child) above.
			// For "export default <expression>", mark the file's default export.
			if hasDefault {
				// Check if a named declaration was already captured
				foundNamed := false
				for i := 0; i < int(node.ChildCount()); i++ {
					child := node.Child(i)
					if child.Type() == "function_declaration" || child.Type() == "class_declaration" {
						nameNode := child.ChildByFieldName("name")
						if nameNode != nil {
							result.Exports = append(result.Exports, nameNode.Content(content))
							foundNamed = true
						}
					}
				}
				if !foundNamed {
					// export default <identifier> already handled above
					// export default <expression> — nothing useful to name
				}
			}

		case "function_declaration":
			if !isTopLevel(node) {
				break
			}
			nameNode := node.ChildByFieldName("name")
			if nameNode != nil {
				result.Defines = append(result.Defines, Symbol{
					Name: nameNode.Content(content),
					Type: "function",
					Line: node.StartPoint().Row + 1,
				})
			}

		case "class_declaration":
			if !isTopLevel(node) {
				break
			}
			nameNode := node.ChildByFieldName("name")
			if nameNode != nil {
				result.Defines = append(result.Defines, Symbol{
					Name: nameNode.Content(content),
					Type: "class",
					Line: node.StartPoint().Row + 1,
				})
			}

		case "lexical_declaration", "variable_declaration":
			if !isTopLevel(node) {
				break
			}
			for i := 0; i < int(node.ChildCount()); i++ {
				child := node.Child(i)
				if child.Type() == "variable_declarator" {
					nameNode := child.ChildByFieldName("name")
					if nameNode == nil {
						continue
					}
					// Skip destructuring patterns
					if nameNode.Type() != "identifier" {
						continue
					}
					symType := "variable"
					valueNode := child.ChildByFieldName("value")
					if valueNode != nil && valueNode.Type() == "arrow_function" {
						symType = "function"
					}
					result.Defines = append(result.Defines, Symbol{
						Name: nameNode.Content(content),
						Type: symType,
						Line: child.StartPoint().Row + 1,
					})
				}
			}
		}

		for i := 0; i < int(node.ChildCount()); i++ {
			walk(node.Child(i))
		}
	}

	walk(root)
	return result
}

// parseAstro handles .astro files by extracting the frontmatter (TypeScript)
// and parsing it for imports. Astro components are always exported by their file.
func parseAstro(parser *sitter.Parser, content []byte) parseResult {
	var result parseResult

	// Extract frontmatter between --- markers
	frontmatter := extractAstroFrontmatter(content)
	if len(frontmatter) > 0 {
		// Parse frontmatter as TypeScript
		parser.SetLanguage(typescript.GetLanguage())
		tree, err := parser.ParseCtx(context.Background(), nil, frontmatter)
		if err == nil {
			defer tree.Close()
			fmResult := parseJSTS(tree.RootNode(), frontmatter)
			result.Imports = fmResult.Imports
			// Include defines from frontmatter (exported functions, etc.)
			result.Defines = fmResult.Defines
		}
	}

	// Astro components are always exported as the default export
	result.Exports = append(result.Exports, "default")
	result.Defines = append(result.Defines, Symbol{
		Name: "default",
		Type: "component",
		Line: 1,
	})

	return result
}

// extractAstroFrontmatter extracts content between --- markers at the start of an Astro file
func extractAstroFrontmatter(content []byte) []byte {
	str := string(content)

	// Find the first ---
	firstIdx := strings.Index(str, "---")
	if firstIdx == -1 {
		return nil
	}

	// Check that --- is at the start (allowing leading whitespace/newlines)
	prefix := strings.TrimSpace(str[:firstIdx])
	if prefix != "" {
		return nil
	}

	// Find the closing ---
	rest := str[firstIdx+3:]
	secondIdx := strings.Index(rest, "---")
	if secondIdx == -1 {
		return nil
	}

	return []byte(rest[:secondIdx])
}

func extractExportClause(node *sitter.Node, content []byte, result *parseResult) {
	for i := 0; i < int(node.ChildCount()); i++ {
		child := node.Child(i)
		if child.Type() == "export_specifier" {
			// export { name } or export { name as alias }
			nameNode := child.ChildByFieldName("name")
			if nameNode != nil {
				result.Exports = append(result.Exports, nameNode.Content(content))
			}
		}
	}
}

func parseClojure(root *sitter.Node, content []byte) parseResult {
	var result parseResult

	var walk func(*sitter.Node)
	walk = func(node *sitter.Node) {
		if node.Type() == "list_lit" && node.ChildCount() > 0 {
			first := findFirstSymbol(node, content)

			switch first {
			case "ns":
				parseNsForm(node, content, &result)

			case "def", "defonce":
				name := findNthSymbol(node, content, 1)
				if name != "" {
					result.Defines = append(result.Defines, Symbol{
						Name: name,
						Type: "var",
						Line: node.StartPoint().Row + 1,
					})
					result.Exports = append(result.Exports, name)
				}

			case "defn", "defn-":
				name := findNthSymbol(node, content, 1)
				if name != "" {
					symType := "function"
					if first == "defn-" {
						symType = "private-function"
					} else {
						result.Exports = append(result.Exports, name)
					}
					result.Defines = append(result.Defines, Symbol{
						Name: name,
						Type: symType,
						Line: node.StartPoint().Row + 1,
					})
				}

			case "defmulti":
				name := findNthSymbol(node, content, 1)
				if name != "" {
					result.Defines = append(result.Defines, Symbol{
						Name: name,
						Type: "multimethod",
						Line: node.StartPoint().Row + 1,
					})
					result.Exports = append(result.Exports, name)
				}

			case "defprotocol":
				name := findNthSymbol(node, content, 1)
				if name != "" {
					result.Defines = append(result.Defines, Symbol{
						Name: name,
						Type: "protocol",
						Line: node.StartPoint().Row + 1,
					})
					result.Exports = append(result.Exports, name)
				}

			case "defrecord", "deftype":
				name := findNthSymbol(node, content, 1)
				if name != "" {
					result.Defines = append(result.Defines, Symbol{
						Name: name,
						Type: first,
						Line: node.StartPoint().Row + 1,
					})
					result.Exports = append(result.Exports, name)
				}
			}
		}

		for i := 0; i < int(node.ChildCount()); i++ {
			walk(node.Child(i))
		}
	}

	walk(root)
	return result
}

func parseNsForm(node *sitter.Node, content []byte, result *parseResult) {
	var walkNs func(*sitter.Node)
	walkNs = func(n *sitter.Node) {
		if n.Type() == "kwd_lit" {
			kwd := n.Content(content)
			if kwd == ":require" || kwd == ":require-macros" {
				parent := n.Parent()
				if parent != nil {
					foundRequire := false
					for i := 0; i < int(parent.ChildCount()); i++ {
						child := parent.Child(i)
						if child == n {
							foundRequire = true
							continue
						}
						if foundRequire {
							extractRequires(child, content, result)
						}
					}
				}
			}
		}
		for i := 0; i < int(n.ChildCount()); i++ {
			walkNs(n.Child(i))
		}
	}
	walkNs(node)
}

func extractRequires(node *sitter.Node, content []byte, result *parseResult) {
	if node.Type() == "vec_lit" {
		sym := findFirstSymbol(node, content)
		if sym != "" {
			result.Imports = append(result.Imports, sym)
		}
	} else if node.Type() == "sym_lit" {
		sym := node.Content(content)
		result.Imports = append(result.Imports, sym)
	}
}

func findFirstSymbol(node *sitter.Node, content []byte) string {
	for i := 0; i < int(node.ChildCount()); i++ {
		child := node.Child(i)
		if child.Type() == "sym_lit" {
			return child.Content(content)
		}
	}
	return ""
}

func findNthSymbol(node *sitter.Node, content []byte, n int) string {
	count := 0
	for i := 0; i < int(node.ChildCount()); i++ {
		child := node.Child(i)
		if child.Type() == "sym_lit" {
			if count == n {
				return child.Content(content)
			}
			count++
		}
	}
	return ""
}

func unique(slice []string) []string {
	seen := make(map[string]bool)
	result := []string{}
	for _, s := range slice {
		if !seen[s] {
			seen[s] = true
			result = append(result, s)
		}
	}
	return result
}

func uniqueSymbols(syms []Symbol) []Symbol {
	seen := make(map[string]bool)
	result := []Symbol{}
	for _, s := range syms {
		if !seen[s.Name] {
			seen[s.Name] = true
			result = append(result, s)
		}
	}
	return result
}

func parseGo(root *sitter.Node, content []byte) parseResult {
	var result parseResult

	var walk func(*sitter.Node)
	walk = func(node *sitter.Node) {
		switch node.Type() {
		case "import_declaration":
			for i := 0; i < int(node.ChildCount()); i++ {
				child := node.Child(i)
				if child.Type() == "import_spec" || child.Type() == "import_spec_list" {
					extractGoImports(child, content, &result)
				}
			}

		case "function_declaration":
			nameNode := node.ChildByFieldName("name")
			if nameNode != nil {
				name := nameNode.Content(content)
				if len(name) > 0 && name[0] >= 'A' && name[0] <= 'Z' {
					result.Exports = append(result.Exports, name)
				}
				result.Defines = append(result.Defines, Symbol{
					Name: name,
					Type: "function",
					Line: node.StartPoint().Row + 1,
				})
			}

		case "method_declaration":
			nameNode := node.ChildByFieldName("name")
			if nameNode != nil {
				name := nameNode.Content(content)
				receiverNode := node.ChildByFieldName("receiver")
				receiver := ""
				if receiverNode != nil {
					receiver = extractReceiverType(receiverNode, content)
				}
				fullName := name
				if receiver != "" {
					fullName = receiver + "." + name
				}
				if len(name) > 0 && name[0] >= 'A' && name[0] <= 'Z' {
					result.Exports = append(result.Exports, fullName)
				}
				result.Defines = append(result.Defines, Symbol{
					Name: fullName,
					Type: "method",
					Line: node.StartPoint().Row + 1,
				})
			}

		case "type_declaration":
			for i := 0; i < int(node.ChildCount()); i++ {
				child := node.Child(i)
				if child.Type() == "type_spec" {
					nameNode := child.ChildByFieldName("name")
					if nameNode != nil {
						name := nameNode.Content(content)
						typeNode := child.ChildByFieldName("type")
						typeKind := "type"
						if typeNode != nil {
							switch typeNode.Type() {
							case "struct_type":
								typeKind = "struct"
							case "interface_type":
								typeKind = "interface"
							}
						}
						if len(name) > 0 && name[0] >= 'A' && name[0] <= 'Z' {
							result.Exports = append(result.Exports, name)
						}
						result.Defines = append(result.Defines, Symbol{
							Name: name,
							Type: typeKind,
							Line: child.StartPoint().Row + 1,
						})
					}
				}
			}

		case "var_declaration", "const_declaration":
			declType := "var"
			if node.Type() == "const_declaration" {
				declType = "const"
			}
			for i := 0; i < int(node.ChildCount()); i++ {
				child := node.Child(i)
				if child.Type() == "var_spec" || child.Type() == "const_spec" {
					nameNode := child.ChildByFieldName("name")
					if nameNode != nil {
						name := nameNode.Content(content)
						if len(name) > 0 && name[0] >= 'A' && name[0] <= 'Z' {
							result.Exports = append(result.Exports, name)
						}
						result.Defines = append(result.Defines, Symbol{
							Name: name,
							Type: declType,
							Line: child.StartPoint().Row + 1,
						})
					}
				}
			}
		}

		for i := 0; i < int(node.ChildCount()); i++ {
			walk(node.Child(i))
		}
	}

	walk(root)
	return result
}

func extractGoImports(node *sitter.Node, content []byte, result *parseResult) {
	if node.Type() == "import_spec_list" {
		for i := 0; i < int(node.ChildCount()); i++ {
			extractGoImports(node.Child(i), content, result)
		}
		return
	}

	if node.Type() == "import_spec" {
		pathNode := node.ChildByFieldName("path")
		if pathNode != nil {
			importPath := strings.Trim(pathNode.Content(content), "\"")
			result.Imports = append(result.Imports, importPath)
		}
	}
}

func extractReceiverType(node *sitter.Node, content []byte) string {
	var walk func(*sitter.Node) string
	walk = func(n *sitter.Node) string {
		if n.Type() == "type_identifier" {
			return n.Content(content)
		}
		for i := 0; i < int(n.ChildCount()); i++ {
			if result := walk(n.Child(i)); result != "" {
				return result
			}
		}
		return ""
	}
	return walk(node)
}

func parsePython(root *sitter.Node, content []byte) parseResult {
	var result parseResult

	var walk func(*sitter.Node, bool)
	walk = func(node *sitter.Node, topLevel bool) {
		switch node.Type() {
		case "import_statement":
			for i := 0; i < int(node.ChildCount()); i++ {
				child := node.Child(i)
				if child.Type() == "dotted_name" {
					result.Imports = append(result.Imports, child.Content(content))
				}
			}

		case "import_from_statement":
			moduleNode := node.ChildByFieldName("module_name")
			if moduleNode != nil {
				result.Imports = append(result.Imports, moduleNode.Content(content))
			} else {
				for i := 0; i < int(node.ChildCount()); i++ {
					child := node.Child(i)
					if child.Type() == "dotted_name" || child.Type() == "relative_import" {
						result.Imports = append(result.Imports, child.Content(content))
						break
					}
				}
			}

		case "function_definition":
			if !topLevel {
				break
			}
			nameNode := node.ChildByFieldName("name")
			if nameNode != nil {
				name := nameNode.Content(content)
				symType := "function"
				if hasPythonDecorator(node, content, "property") {
					symType = "property"
				} else if hasPythonDecorator(node, content, "staticmethod") {
					symType = "staticmethod"
				} else if hasPythonDecorator(node, content, "classmethod") {
					symType = "classmethod"
				}
				if !strings.HasPrefix(name, "_") {
					result.Exports = append(result.Exports, name)
				}
				result.Defines = append(result.Defines, Symbol{
					Name: name,
					Type: symType,
					Line: node.StartPoint().Row + 1,
				})
			}

		case "class_definition":
			if !topLevel {
				break
			}
			nameNode := node.ChildByFieldName("name")
			if nameNode != nil {
				name := nameNode.Content(content)
				if !strings.HasPrefix(name, "_") {
					result.Exports = append(result.Exports, name)
				}
				result.Defines = append(result.Defines, Symbol{
					Name: name,
					Type: "class",
					Line: node.StartPoint().Row + 1,
				})
			}

		case "assignment":
			if topLevel {
				leftNode := node.ChildByFieldName("left")
				if leftNode != nil && leftNode.Type() == "identifier" {
					name := leftNode.Content(content)
					if !strings.HasPrefix(name, "_") {
						result.Exports = append(result.Exports, name)
					}
					result.Defines = append(result.Defines, Symbol{
						Name: name,
						Type: "variable",
						Line: node.StartPoint().Row + 1,
					})
				}
			}
		}

		// Only recurse as top-level for module-level children
		isModule := node.Type() == "module"
		for i := 0; i < int(node.ChildCount()); i++ {
			child := node.Child(i)
			walk(child, isModule)
		}
	}

	walk(root, false)
	return result
}

func hasPythonDecorator(node *sitter.Node, content []byte, name string) bool {
	parent := node.Parent()
	if parent == nil {
		return false
	}

	for i := 0; i < int(parent.ChildCount()); i++ {
		child := parent.Child(i)
		if child == node {
			break
		}
		if child.Type() == "decorator" {
			decoratorContent := child.Content(content)
			if strings.Contains(decoratorContent, name) {
				return true
			}
		}
	}
	return false
}