package main

import (
	"fmt"
	"go/ast"
	"go/parser"
	"go/token"
	"os"
	"path/filepath"
	"sort"
	"strings"
)

// Analyzer parses Go source code and extracts structural information.
type Analyzer struct {
	sourceDir string
	fset      *token.FileSet
}

// NewAnalyzer creates a new Analyzer for the given source directory.
func NewAnalyzer(sourceDir string) *Analyzer {
	return &Analyzer{
		sourceDir: sourceDir,
		fset:      token.NewFileSet(),
	}
}

// Analyze runs the full analysis pipeline.
func (a *Analyzer) Analyze() (*AnalysisResult, error) {
	serverDir := filepath.Join(a.sourceDir, "server")

	// 1. Discover all Go files grouped by directory
	fileGroups, err := a.discoverFiles(serverDir)
	if err != nil {
		return nil, fmt.Errorf("discovering files: %w", err)
	}

	// 2. Parse each group into modules
	result := &AnalysisResult{}
	allImports := make(map[string]*ImportInfo)

	for dir, files := range fileGroups {
		module, imports, err := a.parseModule(dir, files)
		if err != nil {
			return nil, fmt.Errorf("parsing module %s: %w", dir, err)
		}
		result.Modules = append(result.Modules, *module)
		for _, imp := range imports {
			if existing, ok := allImports[imp.ImportPath]; ok {
				existing.UsedInFiles = append(existing.UsedInFiles, imp.UsedInFiles...)
			} else {
				allImports[imp.ImportPath] = &imp
			}
		}
	}

	// 3. Build module-level dependencies from import analysis
	result.Dependencies = a.buildDependencies(result.Modules)

	// 4. Collect imports
	for _, imp := range allImports {
		result.Imports = append(result.Imports, *imp)
	}
	sort.Slice(result.Imports, func(i, j int) bool {
		return result.Imports[i].ImportPath < result.Imports[j].ImportPath
	})

	// Sort modules by name
	sort.Slice(result.Modules, func(i, j int) bool {
		return result.Modules[i].Name < result.Modules[j].Name
	})

	return result, nil
}

// discoverFiles walks the source tree and groups .go files by directory.
func (a *Analyzer) discoverFiles(root string) (map[string][]string, error) {
	groups := make(map[string][]string)
	err := filepath.Walk(root, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		if info.IsDir() {
			if info.Name() == "configs" || info.Name() == "testdata" {
				return filepath.SkipDir
			}
			return nil
		}
		if !strings.HasSuffix(info.Name(), ".go") {
			return nil
		}
		dir := filepath.Dir(path)
		groups[dir] = append(groups[dir], path)
		return nil
	})
	return groups, err
}

// parseModule parses all Go files in a directory into a Module.
func (a *Analyzer) parseModule(dir string, files []string) (*Module, []ImportInfo, error) {
	moduleName := a.moduleNameFromDir(dir)

	module := &Module{
		Name:      moduleName,
		GoPackage: moduleName,
		GoFile:    dir,
	}

	var sourceFiles []string
	var testFiles []string
	for _, f := range files {
		if strings.HasSuffix(f, "_test.go") {
			testFiles = append(testFiles, f)
		} else {
			sourceFiles = append(sourceFiles, f)
		}
	}

	var allImports []ImportInfo
	totalLines := 0

	for _, f := range sourceFiles {
		features, imports, lines, err := a.parseSourceFile(f)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Warning: skipping %s: %v\n", f, err)
			continue
		}
		module.Features = append(module.Features, features...)
		allImports = append(allImports, imports...)
		totalLines += lines
	}

	for _, f := range testFiles {
		tests, _, lines, err := a.parseTestFile(f)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Warning: skipping test %s: %v\n", f, err)
			continue
		}
		module.Tests = append(module.Tests, tests...)
		totalLines += lines
	}

	module.GoLineCount = totalLines
	return module, allImports, nil
}

// parseSourceFile extracts functions, methods, and imports from a Go source file.
func (a *Analyzer) parseSourceFile(filePath string) ([]Feature, []ImportInfo, int, error) {
	src, err := os.ReadFile(filePath)
	if err != nil {
		return nil, nil, 0, err
	}

	file, err := parser.ParseFile(a.fset, filePath, src, parser.ParseComments)
	if err != nil {
		return nil, nil, 0, err
	}

	lines := strings.Count(string(src), "\n") + 1
	relPath := a.relPath(filePath)

	var features []Feature
	var imports []ImportInfo

	for _, imp := range file.Imports {
		path := strings.Trim(imp.Path.Value, "\"")
		imports = append(imports, ImportInfo{
			ImportPath:  path,
			IsStdlib:    isStdlib(path),
			UsedInFiles: []string{relPath},
		})
	}

	for _, decl := range file.Decls {
		fn, ok := decl.(*ast.FuncDecl)
		if !ok {
			continue
		}

		feature := Feature{
			Name:         fn.Name.Name,
			GoFile:       relPath,
			GoMethod:     fn.Name.Name,
			GoLineNumber: a.fset.Position(fn.Pos()).Line,
		}

		startLine := a.fset.Position(fn.Pos()).Line
		endLine := a.fset.Position(fn.End()).Line
		feature.GoLineCount = endLine - startLine + 1

		if fn.Recv != nil && len(fn.Recv.List) > 0 {
			feature.GoClass = a.receiverTypeName(fn.Recv.List[0].Type)
			feature.Name = feature.GoClass + "." + fn.Name.Name
		}

		if fn.Doc != nil {
			feature.Description = strings.TrimSpace(fn.Doc.Text())
		}

		features = append(features, feature)
	}

	return features, imports, lines, nil
}

// parseTestFile extracts test functions from a Go test file.
func (a *Analyzer) parseTestFile(filePath string) ([]TestFunc, []ImportInfo, int, error) {
	src, err := os.ReadFile(filePath)
	if err != nil {
		return nil, nil, 0, err
	}

	file, err := parser.ParseFile(a.fset, filePath, src, parser.ParseComments)
	if err != nil {
		return nil, nil, 0, err
	}

	lines := strings.Count(string(src), "\n") + 1
	relPath := a.relPath(filePath)

	var tests []TestFunc
	var imports []ImportInfo

	for _, imp := range file.Imports {
		path := strings.Trim(imp.Path.Value, "\"")
		imports = append(imports, ImportInfo{
			ImportPath:  path,
			IsStdlib:    isStdlib(path),
			UsedInFiles: []string{relPath},
		})
	}

	for _, decl := range file.Decls {
		fn, ok := decl.(*ast.FuncDecl)
		if !ok {
			continue
		}
		name := fn.Name.Name
		if !strings.HasPrefix(name, "Test") && !strings.HasPrefix(name, "Benchmark") {
			continue
		}

		startLine := a.fset.Position(fn.Pos()).Line
		endLine := a.fset.Position(fn.End()).Line

		test := TestFunc{
			Name:         name,
			GoFile:       relPath,
			GoMethod:     name,
			GoLineNumber: startLine,
			GoLineCount:  endLine - startLine + 1,
		}

		if fn.Doc != nil {
			test.Description = strings.TrimSpace(fn.Doc.Text())
		}

		test.FeatureName = a.inferFeatureName(name)
		test.BestFeatureIdx = -1
		if fn.Body != nil {
			test.Calls = a.extractCalls(fn.Body)
		}
		tests = append(tests, test)
	}

	return tests, imports, lines, nil
}

// buildDependencies creates module-level dependencies based on cross-package imports.
func (a *Analyzer) buildDependencies(modules []Module) []Dependency {
	pkgToModule := make(map[string]string)
	for _, m := range modules {
		pkgToModule[m.GoPackage] = m.Name
	}

	var deps []Dependency
	for _, m := range modules {
		if m.Name != "server" && m.GoPackage != "server" {
			deps = append(deps, Dependency{
				SourceModule:   "server",
				TargetModule:   m.Name,
				DependencyKind: "calls",
			})
		}
	}

	return deps
}

// moduleNameFromDir converts a directory path to a module name.
func (a *Analyzer) moduleNameFromDir(dir string) string {
	base := filepath.Base(dir)
	if base == "server" {
		return "server"
	}
	return base
}

// relPath returns a path relative to the analyzer's source directory.
func (a *Analyzer) relPath(absPath string) string {
	rel, err := filepath.Rel(a.sourceDir, absPath)
	if err != nil {
		return absPath
	}
	return rel
}

// receiverTypeName extracts the type name from a method receiver.
func (a *Analyzer) receiverTypeName(expr ast.Expr) string {
	switch t := expr.(type) {
	case *ast.StarExpr:
		return a.receiverTypeName(t.X)
	case *ast.Ident:
		return t.Name
	default:
		return ""
	}
}

// inferFeatureName attempts to derive a feature name from a test name.
func (a *Analyzer) inferFeatureName(testName string) string {
	name := testName
	for _, prefix := range []string{"Test", "Benchmark"} {
		if strings.HasPrefix(name, prefix) {
			name = strings.TrimPrefix(name, prefix)
			break
		}
	}
	if name == "" {
		return ""
	}
	if idx := strings.Index(name, "_"); idx > 0 {
		name = name[:idx] + "." + name[idx+1:]
	}
	return name
}

// extractCalls walks an AST block statement and extracts all function/method calls.
func (a *Analyzer) extractCalls(body *ast.BlockStmt) []CallInfo {
	seen := make(map[string]bool)
	var calls []CallInfo

	ast.Inspect(body, func(n ast.Node) bool {
		callExpr, ok := n.(*ast.CallExpr)
		if !ok {
			return true
		}

		var ci CallInfo
		switch fun := callExpr.Fun.(type) {
		case *ast.Ident:
			ci = CallInfo{FuncName: fun.Name}
		case *ast.SelectorExpr:
			ci = CallInfo{
				RecvOrPkg:  extractIdent(fun.X),
				MethodName: fun.Sel.Name,
				IsSelector: true,
			}
		default:
			return true
		}

		key := ci.callKey()
		if !seen[key] && !isFilteredCall(ci) {
			seen[key] = true
			calls = append(calls, ci)
		}
		return true
	})

	return calls
}

// extractIdent extracts an identifier name from an expression (handles X in X.Y).
func extractIdent(expr ast.Expr) string {
	switch e := expr.(type) {
	case *ast.Ident:
		return e.Name
	case *ast.SelectorExpr:
		return extractIdent(e.X) + "." + e.Sel.Name
	default:
		return ""
	}
}

// isFilteredCall returns true if a call should be excluded from feature matching.
func isFilteredCall(c CallInfo) bool {
	if c.IsSelector {
		recv := c.RecvOrPkg
		// testing.T/B methods
		if recv == "t" || recv == "b" || recv == "tb" {
			return true
		}
		// stdlib packages
		if stdlibPkgs[recv] {
			return true
		}
		// NATS client libs
		if recv == "nats" || recv == "nuid" || recv == "nkeys" || recv == "jwt" {
			return true
		}
		return false
	}

	// Go builtins
	name := c.FuncName
	if builtinFuncs[name] {
		return true
	}

	// Test assertion helpers
	lower := strings.ToLower(name)
	if strings.HasPrefix(name, "require_") {
		return true
	}
	for _, prefix := range []string{"check", "verify", "assert", "expect"} {
		if strings.HasPrefix(lower, prefix) {
			return true
		}
	}

	return false
}

// featureRef identifies a feature within the analysis result.
type featureRef struct {
	moduleIdx  int
	featureIdx int
	goFile     string
	goClass    string
}

// resolveCallGraph matches test calls against known features across all modules.
func resolveCallGraph(result *AnalysisResult) {
	// Build method index: go_method name → list of feature refs
	methodIndex := make(map[string][]featureRef)
	for mi, mod := range result.Modules {
		for fi, feat := range mod.Features {
			ref := featureRef{
				moduleIdx:  mi,
				featureIdx: fi,
				goFile:     feat.GoFile,
				goClass:    feat.GoClass,
			}
			methodIndex[feat.GoMethod] = append(methodIndex[feat.GoMethod], ref)
		}
	}

	// For each test, resolve calls to features
	for mi := range result.Modules {
		mod := &result.Modules[mi]
		for ti := range mod.Tests {
			test := &mod.Tests[ti]
			seen := make(map[int]bool) // feature indices already linked
			var linked []int

			testFileBase := sourceFileBase(test.GoFile)

			for _, call := range test.Calls {
				// Look up the method name
				name := call.MethodName
				if !call.IsSelector {
					name = call.FuncName
				}

				candidates := methodIndex[name]
				if len(candidates) == 0 {
					continue
				}
				// Ambiguity threshold: skip very common method names
				if len(candidates) > 10 {
					continue
				}

				// Filter to same module
				var sameModule []featureRef
				for _, ref := range candidates {
					if ref.moduleIdx == mi {
						sameModule = append(sameModule, ref)
					}
				}
				if len(sameModule) == 0 {
					continue
				}

				for _, ref := range sameModule {
					if !seen[ref.featureIdx] {
						seen[ref.featureIdx] = true
						linked = append(linked, ref.featureIdx)
					}
				}
			}

			test.LinkedFeatures = linked

			// Set BestFeatureIdx using priority:
			// (a) existing inferFeatureName match
			// (b) same-file-base match
			// (c) first remaining candidate
			if test.BestFeatureIdx < 0 && len(linked) > 0 {
				// Try same-file-base match first
				for _, fi := range linked {
					featFileBase := sourceFileBase(mod.Features[fi].GoFile)
					if featFileBase == testFileBase {
						test.BestFeatureIdx = fi
						break
					}
				}
				// Fall back to first candidate
				if test.BestFeatureIdx < 0 {
					test.BestFeatureIdx = linked[0]
				}
			}
		}
	}
}

// sourceFileBase strips _test.go suffix and path to get the base file name.
func sourceFileBase(goFile string) string {
	base := filepath.Base(goFile)
	base = strings.TrimSuffix(base, "_test.go")
	base = strings.TrimSuffix(base, ".go")
	return base
}

var stdlibPkgs = map[string]bool{
	"fmt": true, "time": true, "strings": true, "bytes": true, "errors": true,
	"os": true, "math": true, "sort": true, "reflect": true, "sync": true,
	"context": true, "io": true, "filepath": true, "strconv": true,
	"encoding": true, "json": true, "binary": true, "hex": true, "rand": true,
	"runtime": true, "atomic": true, "slices": true, "testing": true,
	"net": true, "bufio": true, "crypto": true, "log": true, "regexp": true,
	"unicode": true, "http": true, "url": true,
}

var builtinFuncs = map[string]bool{
	"make": true, "append": true, "len": true, "cap": true, "close": true,
	"delete": true, "panic": true, "recover": true, "print": true,
	"println": true, "copy": true, "new": true,
}

// isStdlib checks if an import path is a Go standard library package.
func isStdlib(importPath string) bool {
	firstSlash := strings.Index(importPath, "/")
	var first string
	if firstSlash < 0 {
		first = importPath
	} else {
		first = importPath[:firstSlash]
	}
	return !strings.Contains(first, ".")
}