// stella-callgraph-go
// Call graph extraction tool for Go projects using SSA analysis.
package main

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

	"golang.org/x/tools/go/callgraph"
	"golang.org/x/tools/go/callgraph/cha"
	"golang.org/x/tools/go/callgraph/rta"
	"golang.org/x/tools/go/packages"
	"golang.org/x/tools/go/ssa"
	"golang.org/x/tools/go/ssa/ssautil"
)

// CallGraphResult is the output structure
type CallGraphResult struct {
	Module      string       `json:"module"`
	Nodes       []Node       `json:"nodes"`
	Edges       []Edge       `json:"edges"`
	Entrypoints []Entrypoint `json:"entrypoints"`
}

// Node represents a function in the call graph
type Node struct {
	ID          string   `json:"id"`
	Package     string   `json:"package"`
	Name        string   `json:"name"`
	Signature   string   `json:"signature"`
	Position    Position `json:"position"`
	Visibility  string   `json:"visibility"`
	Annotations []string `json:"annotations"`
}

// Edge represents a call between functions
type Edge struct {
	From string   `json:"from"`
	To   string   `json:"to"`
	Kind string   `json:"kind"`
	Site Position `json:"site"`
}

// Position in source code
type Position struct {
	File   string `json:"file"`
	Line   int    `json:"line"`
	Column int    `json:"column"`
}

// Entrypoint represents an entry point function
type Entrypoint struct {
	ID     string `json:"id"`
	Type   string `json:"type"`
	Route  string `json:"route,omitempty"`
	Method string `json:"method,omitempty"`
}

func main() {
	var (
		projectPath string
		algorithm   string
		jsonFormat  bool
	)

	flag.StringVar(&projectPath, "path", ".", "Path to Go project")
	flag.StringVar(&algorithm, "algo", "cha", "Call graph algorithm: cha, rta, or static")
	flag.BoolVar(&jsonFormat, "json", false, "Output formatted JSON")
	flag.Parse()

	if len(flag.Args()) > 0 {
		projectPath = flag.Args()[0]
	}

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

	var output []byte
	if jsonFormat {
		output, err = json.MarshalIndent(result, "", "  ")
	} else {
		output, err = json.Marshal(result)
	}

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

	fmt.Println(string(output))
}

func analyzeProject(projectPath string, algorithm string) (*CallGraphResult, error) {
	absPath, err := filepath.Abs(projectPath)
	if err != nil {
		return nil, fmt.Errorf("invalid path: %w", err)
	}

	// Load packages
	cfg := &packages.Config{
		Mode: packages.LoadAllSyntax,
		Dir:  absPath,
	}

	pkgs, err := packages.Load(cfg, "./...")
	if err != nil {
		return nil, fmt.Errorf("failed to load packages: %w", err)
	}

	if len(pkgs) == 0 {
		return nil, fmt.Errorf("no packages found")
	}

	// Check for errors
	for _, pkg := range pkgs {
		if len(pkg.Errors) > 0 {
			// Log but continue
			for _, e := range pkg.Errors {
				fmt.Fprintf(os.Stderr, "Warning: %v\n", e)
			}
		}
	}

	// Build SSA
	prog, _ := ssautil.AllPackages(pkgs, ssa.SanityCheckFunctions)
	prog.Build()

	// Extract module name
	moduleName := extractModuleName(absPath, pkgs)

	// Build call graph using the specified algorithm
	var cg *callgraph.Graph
	switch algorithm {
	case "rta":
		// RTA (Rapid Type Analysis) - more precise for programs with main
		mains := ssautil.MainPackages(prog.AllPackages())
		if len(mains) > 0 {
			var roots []*ssa.Function
			for _, main := range mains {
				if mainFn := main.Func("main"); mainFn != nil {
					roots = append(roots, mainFn)
				}
				if initFn := main.Func("init"); initFn != nil {
					roots = append(roots, initFn)
				}
			}
			if len(roots) > 0 {
				rtaResult := rta.Analyze(roots, true)
				cg = rtaResult.CallGraph
			}
		}
		if cg == nil {
			// Fall back to CHA if no main packages
			cg = cha.CallGraph(prog)
		}
	case "cha":
		// CHA (Class Hierarchy Analysis) - sound but less precise
		cg = cha.CallGraph(prog)
	default:
		// Default to CHA
		cg = cha.CallGraph(prog)
	}

	// Collect nodes and edges from call graph
	nodes := make([]Node, 0)
	edges := make([]Edge, 0)
	entrypoints := make([]Entrypoint, 0)
	seenNodes := make(map[string]bool)
	seenEdges := make(map[string]bool)

	// If we have a call graph, use it for edges
	if cg != nil {
		callgraph.GraphVisitEdges(cg, func(edge *callgraph.Edge) error {
			if edge.Caller.Func == nil || edge.Callee.Func == nil {
				return nil
			}

			callerID := makeSymbolID(edge.Caller.Func)
			calleeID := makeSymbolID(edge.Callee.Func)

			// Add caller node if not seen
			if !seenNodes[callerID] {
				seenNodes[callerID] = true
				nodes = append(nodes, makeNodeFromFunction(prog, edge.Caller.Func))
			}

			// Add callee node if not seen
			if !seenNodes[calleeID] {
				seenNodes[calleeID] = true
				nodes = append(nodes, makeNodeFromFunction(prog, edge.Callee.Func))
			}

			// Add edge
			edgeKey := fmt.Sprintf("%s|%s", callerID, calleeID)
			if !seenEdges[edgeKey] {
				seenEdges[edgeKey] = true

				kind := "direct"
				if edge.Site != nil {
					if _, ok := edge.Site.(*ssa.Go); ok {
						kind = "goroutine"
					} else if _, ok := edge.Site.(*ssa.Defer); ok {
						kind = "defer"
					}
				}

				var site Position
				if edge.Site != nil {
					pos := prog.Fset.Position(edge.Site.Pos())
					site = Position{
						File: pos.Filename,
						Line: pos.Line,
					}
				}

				edges = append(edges, Edge{
					From: callerID,
					To:   calleeID,
					Kind: kind,
					Site: site,
				})
			}

			return nil
		})
	}

	// Also scan all functions to find any missing nodes and entrypoints
	for _, pkg := range prog.AllPackages() {
		if pkg == nil {
			continue
		}

		for _, member := range pkg.Members {
			fn, ok := member.(*ssa.Function)
			if !ok {
				continue
			}

			nodeID := makeSymbolID(fn)
			if !seenNodes[nodeID] {
				seenNodes[nodeID] = true
				nodes = append(nodes, makeNodeFromFunction(prog, fn))
			}

			// Check for entrypoints
			if ep := detectEntrypoint(fn); ep != nil {
				entrypoints = append(entrypoints, *ep)
			}
		}
	}

	return &CallGraphResult{
		Module:      moduleName,
		Nodes:       nodes,
		Edges:       edges,
		Entrypoints: entrypoints,
	}, nil
}

func makeNodeFromFunction(prog *ssa.Program, fn *ssa.Function) Node {
	pos := prog.Fset.Position(fn.Pos())
	pkgPath := ""
	if fn.Pkg != nil {
		pkgPath = fn.Pkg.Pkg.Path()
	}

	return Node{
		ID:         makeSymbolID(fn),
		Package:    pkgPath,
		Name:       fn.Name(),
		Signature:  fn.Signature.String(),
		Position: Position{
			File:   pos.Filename,
			Line:   pos.Line,
			Column: pos.Column,
		},
		Visibility:  getVisibility(fn.Name()),
		Annotations: detectAnnotations(fn),
	}
}

func extractModuleName(projectPath string, pkgs []*packages.Package) string {
	// Try to get from go.mod
	goModPath := filepath.Join(projectPath, "go.mod")
	if data, err := os.ReadFile(goModPath); err == nil {
		lines := strings.Split(string(data), "\n")
		for _, line := range lines {
			if strings.HasPrefix(line, "module ") {
				return strings.TrimSpace(strings.TrimPrefix(line, "module "))
			}
		}
	}

	// Fall back to first package path
	if len(pkgs) > 0 {
		return pkgs[0].PkgPath
	}

	return filepath.Base(projectPath)
}

func makeSymbolID(fn *ssa.Function) string {
	if fn.Pkg == nil {
		return fmt.Sprintf("go:external/%s", fn.Name())
	}

	pkg := fn.Pkg.Pkg.Path()
	if fn.Signature.Recv() != nil {
		// Method
		recv := fn.Signature.Recv().Type().String()
		recv = strings.TrimPrefix(recv, "*")
		if idx := strings.LastIndex(recv, "."); idx >= 0 {
			recv = recv[idx+1:]
		}
		return fmt.Sprintf("go:%s.%s.%s", pkg, recv, fn.Name())
	}
	return fmt.Sprintf("go:%s.%s", pkg, fn.Name())
}

func getVisibility(name string) string {
	if len(name) == 0 {
		return "private"
	}
	if name[0] >= 'A' && name[0] <= 'Z' {
		return "public"
	}
	return "private"
}

func detectAnnotations(fn *ssa.Function) []string {
	// Go doesn't have annotations, but we can detect patterns
	annotations := make([]string, 0)

	// Detect handler patterns from naming
	if strings.HasSuffix(fn.Name(), "Handler") {
		annotations = append(annotations, "handler")
	}
	if strings.HasSuffix(fn.Name(), "Middleware") {
		annotations = append(annotations, "middleware")
	}

	return annotations
}

func detectEntrypoint(fn *ssa.Function) *Entrypoint {
	name := fn.Name()
	pkg := ""
	if fn.Pkg != nil {
		pkg = fn.Pkg.Pkg.Path()
	}

	nodeID := makeSymbolID(fn)

	// main.main
	if name == "main" && strings.HasSuffix(pkg, "main") {
		return &Entrypoint{
			ID:   nodeID,
			Type: "cli_command",
		}
	}

	// init functions
	if name == "init" {
		return &Entrypoint{
			ID:   nodeID,
			Type: "background_job",
		}
	}

	// HTTP handler patterns
	if strings.HasSuffix(name, "Handler") || strings.Contains(name, "Handle") {
		return &Entrypoint{
			ID:   nodeID,
			Type: "http_handler",
		}
	}

	// gRPC patterns
	if strings.HasSuffix(name, "Server") && strings.HasPrefix(name, "Register") {
		return &Entrypoint{
			ID:   nodeID,
			Type: "grpc_method",
		}
	}

	return nil
}