package tools

import (
	"bufio"
	"bytes"
	"fmt"
	"go/ast"
	"go/parser"
	"go/printer"
	"go/token"
	"regexp"
	"sort"
	"strings"
)

// symbolInfo holds extracted symbol metadata.
type symbolInfo struct {
	Name      string
	Kind      string // "import", "function", "method", "class", "type", "interface", "trait", "enum", "struct"
	StartLine int
}

// sourceFileExtensions lists extensions that we can extract symbols from.
var sourceFileExtensions = map[string]bool{
	".go":  true,
	".ts":  true,
	".tsx": true,
	".js":  true,
	".jsx": true,
	".mjs": true,
	".py":  true,
	".rs":  true,
}

// extractSymbols extracts symbols from a source file using language-specific
// parsing (go/ast for Go, regex for other languages).
func extractSymbols(path string, data []byte) ([]symbolInfo, error) {
	ext := ""
	if idx := strings.LastIndex(path, "."); idx >= 0 {
		ext = path[idx:]
	}

	switch ext {
	case ".go":
		return extractGoSymbols(path, data)
	case ".py":
		return extractPythonsymbols(data)
	case ".ts", ".tsx", ".js", ".jsx", ".mjs":
		return extractTSSymbols(data)
	case ".rs":
		return extractRustSymbols(data)
	default:
		return nil, nil
	}
}

// ────────────────────────────────────────────────────────────────────────────
// Go symbol extraction (using go/ast)
// ────────────────────────────────────────────────────────────────────────────

func extractGoSymbols(path string, data []byte) ([]symbolInfo, error) {
	fset := token.NewFileSet()
	// Parse imports only first for the import list.
	importFile, err := parser.ParseFile(fset, path, data, parser.ImportsOnly|parser.ParseComments)
	if err != nil {
		return nil, err
	}
	full, err := parser.ParseFile(fset, path, data, 0)
	if err != nil {
		return nil, err
	}

	var symbols []symbolInfo

	// Package
	symbols = append(symbols, symbolInfo{
		Name:      full.Name.Name,
		Kind:      "package",
		StartLine: fset.Position(full.Name.Pos()).Line,
	})

	// Imports
	for _, imp := range importFile.Imports {
		symbols = append(symbols, symbolInfo{
			Name:      strings.Trim(imp.Path.Value, "\""),
			Kind:      "import",
			StartLine: fset.Position(imp.Pos()).Line,
		})
	}

	// Declarations
	for _, decl := range full.Decls {
		switch d := decl.(type) {
		case *ast.GenDecl:
			if d.Tok != token.TYPE {
				continue
			}
			for _, spec := range d.Specs {
				if ts, ok := spec.(*ast.TypeSpec); ok {
					kind := goTypeKind(ts.Type)
					symbols = append(symbols, symbolInfo{
						Name:      ts.Name.Name,
						Kind:      kind,
						StartLine: fset.Position(ts.Pos()).Line,
					})
				}
			}
		case *ast.FuncDecl:
			line := fset.Position(d.Pos()).Line
			if d.Recv != nil {
				symbols = append(symbols, symbolInfo{
					Name:      fmt.Sprintf("%s.%s", goReceiverName(d.Recv), d.Name.Name),
					Kind:      "method",
					StartLine: line,
				})
			} else {
				symbols = append(symbols, symbolInfo{
					Name:      d.Name.Name,
					Kind:      "function",
					StartLine: line,
				})
			}
		}
	}

	sortSymbols(symbols)
	return symbols, nil
}

func goTypeKind(expr ast.Expr) string {
	switch expr.(type) {
	case *ast.StructType:
		return "struct"
	case *ast.InterfaceType:
		return "interface"
	case *ast.FuncType:
		return "func"
	default:
		return "alias"
	}
}

func goReceiverName(recv *ast.FieldList) string {
	if recv == nil || len(recv.List) == 0 {
		return "?"
	}
	var b bytes.Buffer
	if err := printer.Fprint(&b, token.NewFileSet(), recv.List[0].Type); err != nil {
		return "?"
	}
	return b.String()
}

// ────────────────────────────────────────────────────────────────────────────
// Python symbol extraction (regex-based)
// ────────────────────────────────────────────────────────────────────────────

var (
	pyFuncRe      = regexp.MustCompile(`^\s*def\s+(\w+)\s*\(`)
	pyClassRe     = regexp.MustCompile(`^\s*class\s+(\w+)\s*[(:]`)
	pyImportRe    = regexp.MustCompile(`^\s*(?:from\s+\S+\s+)?import\s+`)
	pyDecoratorRe = regexp.MustCompile(`^\s*@`)
)

func extractPythonsymbols(data []byte) ([]symbolInfo, error) {
	var symbols []symbolInfo
	scanner := bufio.NewScanner(bytes.NewReader(data))
	lineNum := 0
	inDecorator := false

	for scanner.Scan() {
		lineNum++
		line := scanner.Text()

		// Track decorator stacks
		if pyDecoratorRe.MatchString(line) {
			inDecorator = true
			continue
		}

		// Import lines
		if pyImportRe.MatchString(line) {
			// Extract the module name for a simple "import foo" or "from foo import bar"
			importPath := strings.TrimSpace(line)
			symbols = append(symbols, symbolInfo{
				Name:      importPath,
				Kind:      "import",
				StartLine: lineNum,
			})
			inDecorator = false
			continue
		}

		matches := pyClassRe.FindStringSubmatch(line)
		if matches != nil {
			symbols = append(symbols, symbolInfo{
				Name:      matches[1],
				Kind:      "class",
				StartLine: lineNum,
			})
			inDecorator = false
			continue
		}

		matches = pyFuncRe.FindStringSubmatch(line)
		if matches != nil {
			kind := "function"
			if inDecorator {
				// Could be a decorated function or method — still mark as function
				// (class-method detection requires indentation tracking)
			}
			symbols = append(symbols, symbolInfo{
				Name:      matches[1],
				Kind:      kind,
				StartLine: lineNum,
			})
			inDecorator = false
			continue
		}

		inDecorator = false
	}

	sortSymbols(symbols)
	return symbols, scanner.Err()
}

// ────────────────────────────────────────────────────────────────────────────
// TypeScript / JavaScript symbol extraction (regex-based)
// ────────────────────────────────────────────────────────────────────────────

var (
	tsFuncRe      = regexp.MustCompile(`(?:export\s+(?:default\s+)?)?(?:async\s+)?function\s+(\w+)\s*[<(]`)
	tsArrowRe     = regexp.MustCompile(`(?:export\s+(?:default\s+)?)?(?:const|let|var)\s+(\w+)\s*=\s*(?:async\s*)?\s*\(`)
	tsClassRe     = regexp.MustCompile(`(?:export\s+(?:default\s+)?)?class\s+(\w+)`)
	tsInterfaceRe = regexp.MustCompile(`(?:export\s+(?:default\s+)?)?interface\s+(\w+)`)
	tsTypeRe      = regexp.MustCompile(`(?:export\s+(?:default\s+)?)?type\s+(\w+)\s*=`)
	tsEnumRe      = regexp.MustCompile(`(?:export\s+(?:default\s+)?)?enum\s+(\w+)`)
	tsImportRe    = regexp.MustCompile(`^\s*import\s+`)
	tsMethodRe    = regexp.MustCompile(`^\s*(?:public\s+|private\s+|protected\s+|static\s+|async\s+)*(\w+)\s*[<(][^)]*\)\s*(?::\s*\S+\s*)?\{`)
)

func extractTSSymbols(data []byte) ([]symbolInfo, error) {
	var symbols []symbolInfo
	scanner := bufio.NewScanner(bytes.NewReader(data))
	lineNum := 0

	// Track brace depth to detect class body methods.
	braceDepth := 0
	inClass := false

	for scanner.Scan() {
		lineNum++
		line := scanner.Text()
		trimmed := strings.TrimSpace(line)

		// Count braces on this line (applied at end of iteration).
		opens := strings.Count(line, "{")
		closes := strings.Count(line, "}")

		// Import lines — skip everything else, just track braces.
		if tsImportRe.MatchString(line) {
			importPath := trimmed
			symbols = append(symbols, symbolInfo{
				Name:      importPath,
				Kind:      "import",
				StartLine: lineNum,
			})
			goto applyBraces
		}

		// Skip empty and comment-only lines.
		if trimmed == "" || strings.HasPrefix(trimmed, "//") || strings.HasPrefix(trimmed, "/*") {
			goto applyBraces
		}

		// Interface
		if matches := tsInterfaceRe.FindStringSubmatch(line); matches != nil {
			symbols = append(symbols, symbolInfo{
				Name:      matches[1],
				Kind:      "interface",
				StartLine: lineNum,
			})
			goto applyBraces
		}

		// Type alias
		if matches := tsTypeRe.FindStringSubmatch(line); matches != nil {
			symbols = append(symbols, symbolInfo{
				Name:      matches[1],
				Kind:      "type",
				StartLine: lineNum,
			})
			goto applyBraces
		}

		// Enum
		if matches := tsEnumRe.FindStringSubmatch(line); matches != nil {
			symbols = append(symbols, symbolInfo{
				Name:      matches[1],
				Kind:      "enum",
				StartLine: lineNum,
			})
			goto applyBraces
		}

		// Class
		if matches := tsClassRe.FindStringSubmatch(line); matches != nil {
			symbols = append(symbols, symbolInfo{
				Name:      matches[1],
				Kind:      "class",
				StartLine: lineNum,
			})
			// Set inClass BEFORE applying braces for this line, so that
			// methods on the same line (unlikely) or next line are detected.
			if opens > 0 {
				inClass = true
			}
			goto applyBraces
		}

		// Method inside class body (at depth 1 brace level).
		if inClass && braceDepth == 1 {
			if matches := tsMethodRe.FindStringSubmatch(line); matches != nil {
				name := matches[1]
				if !isKeyword(name) {
					symbols = append(symbols, symbolInfo{
						Name:      name,
						Kind:      "method",
						StartLine: lineNum,
					})
				}
			}
		}

		// Function declaration (top-level only, or nested function).
		if !inClass || braceDepth > 1 {
			if matches := tsFuncRe.FindStringSubmatch(line); matches != nil {
				symbols = append(symbols, symbolInfo{
					Name:      matches[1],
					Kind:      "function",
					StartLine: lineNum,
				})
			}
		}

		// Arrow function assigned to const/let/var.
		if !inClass || braceDepth > 1 {
			if matches := tsArrowRe.FindStringSubmatch(line); matches != nil {
				symbols = append(symbols, symbolInfo{
					Name:      matches[1],
					Kind:      "function",
					StartLine: lineNum,
				})
			}
		}

	applyBraces:
		braceDepth += opens - closes
		if braceDepth <= 0 {
			braceDepth = 0
			inClass = false
		}
	}

	sortSymbols(symbols)
	return symbols, scanner.Err()
}

func isKeyword(s string) bool {
	switch s {
	case "if", "else", "for", "while", "switch", "case", "return", "throw",
		"try", "catch", "finally", "new", "delete", "typeof", "instanceof",
		"constructor", "get", "set":
		return true
	}
	return false
}

// ────────────────────────────────────────────────────────────────────────────
// Rust symbol extraction (regex-based)
// ────────────────────────────────────────────────────────────────────────────

var (
	rsFnRe     = regexp.MustCompile(`^\s*(?:pub(?:\s*\(\s*crate\s*\))?\s+)?(?:async\s+)?(?:unsafe\s+)?fn\s+(\w+)\s*[<(]`)
	rsStructRe = regexp.MustCompile(`^\s*(?:pub(?:\s*\(\s*crate\s*\))?\s+)?struct\s+(\w+)`)
	rsEnumRe   = regexp.MustCompile(`^\s*(?:pub(?:\s*\(\s*crate\s*\))?\s+)?enum\s+(\w+)`)
	rsTraitRe  = regexp.MustCompile(`^\s*(?:pub(?:\s*\(\s*crate\s*\))?\s+)?(?:unsafe\s+)?trait\s+(\w+)`)
	rsImplRe   = regexp.MustCompile(`^\s*impl\s+`)
	rsUseRe    = regexp.MustCompile(`^\s*(?:pub\s+)?use\s+`)
)

func extractRustSymbols(data []byte) ([]symbolInfo, error) {
	var symbols []symbolInfo
	scanner := bufio.NewScanner(bytes.NewReader(data))
	lineNum := 0

	for scanner.Scan() {
		lineNum++
		line := scanner.Text()
		trimmed := strings.TrimSpace(line)

		// Skip empty and comments
		if trimmed == "" || strings.HasPrefix(trimmed, "//") || strings.HasPrefix(trimmed, "/*") {
			continue
		}

		// Use
		if rsUseRe.MatchString(line) {
			symbols = append(symbols, symbolInfo{
				Name:      trimmed,
				Kind:      "import",
				StartLine: lineNum,
			})
			continue
		}

		// Struct
		if matches := rsStructRe.FindStringSubmatch(line); matches != nil {
			symbols = append(symbols, symbolInfo{
				Name:      matches[1],
				Kind:      "struct",
				StartLine: lineNum,
			})
			continue
		}

		// Enum
		if matches := rsEnumRe.FindStringSubmatch(line); matches != nil {
			symbols = append(symbols, symbolInfo{
				Name:      matches[1],
				Kind:      "enum",
				StartLine: lineNum,
			})
			continue
		}

		// Trait
		if matches := rsTraitRe.FindStringSubmatch(line); matches != nil {
			symbols = append(symbols, symbolInfo{
				Name:      matches[1],
				Kind:      "trait",
				StartLine: lineNum,
			})
			continue
		}

		// Impl
		if rsImplRe.MatchString(line) {
			// Extract type name from "impl Foo" or "impl<T> Foo" or "impl Trait for Foo"
			rest := strings.TrimPrefix(trimmed, "impl")
			rest = strings.TrimSpace(rest)
			// Remove generics
			if idx := strings.Index(rest, "<"); idx >= 0 {
				depth := 1
				for i := idx + 1; i < len(rest); i++ {
					if rest[i] == '<' {
						depth++
					} else if rest[i] == '>' {
						depth--
						if depth == 0 {
							rest = strings.TrimSpace(rest[:idx] + rest[i+1:])
							break
						}
					}
				}
			}
			// Handle "for Type"
			if idx := strings.Index(rest, " for "); idx >= 0 {
				rest = strings.TrimSpace(rest[idx+5:])
			}
			name := strings.Fields(rest)[0]
			symbols = append(symbols, symbolInfo{
				Name:      name,
				Kind:      "impl",
				StartLine: lineNum,
			})
			continue
		}

		// Function
		if matches := rsFnRe.FindStringSubmatch(line); matches != nil {
			symbols = append(symbols, symbolInfo{
				Name:      matches[1],
				Kind:      "function",
				StartLine: lineNum,
			})
		}
	}

	sortSymbols(symbols)
	return symbols, scanner.Err()
}

// ────────────────────────────────────────────────────────────────────────────
// Helpers
// ────────────────────────────────────────────────────────────────────────────

func sortSymbols(symbols []symbolInfo) {
	sort.Slice(symbols, func(i, j int) bool {
		return symbols[i].StartLine < symbols[j].StartLine
	})
}

// groupSymbols groups symbols by kind and returns formatted strings.
func groupSymbols(symbols []symbolInfo) map[string][]string {
	groups := make(map[string][]string)
	for _, s := range symbols {
		kind := s.Kind
		switch kind {
		case "package":
			groups["package"] = append(groups["package"], s.Name)
		case "import":
			groups["imports"] = append(groups["imports"], s.Name)
		case "function":
			groups["functions"] = append(groups["functions"], s.Name)
		case "method":
			groups["methods"] = append(groups["methods"], s.Name)
		case "class", "interface", "struct", "enum", "trait", "alias", "func":
			// All structural types go under "types" for consistency.
			groups["types"] = append(groups["types"], fmt.Sprintf("%s %s", s.Name, kind))
		case "type":
			// TS type alias
			groups["types"] = append(groups["types"], fmt.Sprintf("%s type", s.Name))
		case "impl":
			groups["types"] = append(groups["types"], fmt.Sprintf("impl %s", s.Name))
		}
	}
	return groups
}