package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"go/parser"
	"go/token"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"sync"
)

// DependencyGraph 表示依赖关系图
type DependencyGraph struct {
	Packages map[string]*PackageInfo
	Mutex    sync.RWMutex
}

// PackageInfo 表示包的信息
type PackageInfo struct {
	Name       string
	Path       string
	Imports    []string
	ImportedBy []string
	IsExternal bool
	IsStandard bool
}

// DependencyEdge 表示依赖边
type DependencyEdge struct {
	From string `json:"from"`
	To   string `json:"to"`
}

// DependencyOutput 输出结构
type DependencyOutput struct {
	Packages []PackageOutput  `json:"packages"`
	Edges    []DependencyEdge `json:"edges"`
}

// PackageOutput 包输出结构
type PackageOutput struct {
	Name       string   `json:"name"`
	Path       string   `json:"path"`
	Imports    []string `json:"imports"`
	ImportedBy []string `json:"imported_by,omitempty"`
	IsExternal bool     `json:"is_external"`
	IsStandard bool     `json:"is_standard"`
}

func main() {
	// 解析命令行参数
	rootDir := flag.String("dir", ".", "要扫描的Go工程目录")
	outputFormat := flag.String("format", "text", "输出格式: text, json, dot, csv")
	outputFile := flag.String("output", "", "输出文件路径（默认输出到控制台）")
	ignoreStdLib := flag.Bool("ignore-std", false, "忽略标准库依赖")
	includeTests := flag.Bool("include-tests", false, "包含测试文件")
	flag.Parse()

	// 检查目录是否存在
	if _, err := os.Stat(*rootDir); os.IsNotExist(err) {
		fmt.Printf("错误: 目录不存在: %s\n", *rootDir)
		os.Exit(1)
	}

	// 构建依赖图
	graph := &DependencyGraph{
		Packages: make(map[string]*PackageInfo),
	}

	fmt.Printf("正在扫描目录: %s\n", *rootDir)

	// 扫描Go文件
	err := scanDirectory(*rootDir, graph, *includeTests)
	if err != nil {
		fmt.Printf("扫描错误: %v\n", err)
		os.Exit(1)
	}

	// 计算导入关系
	calculateImportedBy(graph)

	// 生成输出
	output := generateOutput(graph, *ignoreStdLib)

	// 输出结果
	err = outputResult(output, *outputFormat, *outputFile)
	if err != nil {
		fmt.Printf("输出错误: %v\n", err)
		os.Exit(1)
	}

	fmt.Printf("\n扫描完成! 共发现 %d 个包，%d 个依赖关系\n",
		len(output.Packages), len(output.Edges))
}

// scanDirectory 递归扫描目录
func scanDirectory(dir string, graph *DependencyGraph, includeTests bool) error {
	return filepath.Walk(dir, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}

		// 跳过vendor目录和隐藏目录
		if info.IsDir() {
			if strings.Contains(path, "vendor") || strings.HasPrefix(filepath.Base(path), ".") {
				return filepath.SkipDir
			}
			return nil
		}

		// 只处理.go文件
		if filepath.Ext(path) != ".go" {
			return nil
		}

		// 可选跳过测试文件
		if !includeTests && strings.HasSuffix(path, "_test.go") {
			return nil
		}

		// 解析Go文件
		analyzeGoFile(path, graph)
		return nil
	})
}

// analyzeGoFile 分析单个Go文件
func analyzeGoFile(filePath string, graph *DependencyGraph) {
	fset := token.NewFileSet()
	node, err := parser.ParseFile(fset, filePath, nil, parser.ImportsOnly)
	if err != nil {
		fmt.Printf("警告: 无法解析文件 %s: %v\n", filePath, err)
		return
	}

	// 获取包名和相对路径
	relPath, _ := filepath.Rel(".", filepath.Dir(filePath))
	if relPath == "." {
		relPath = ""
	}
	packagePath := relPath
	if packagePath == "" {
		packagePath = "."
	}

	// 创建或获取包信息
	graph.Mutex.Lock()
	packageInfo, exists := graph.Packages[packagePath]
	if !exists {
		packageInfo = &PackageInfo{
			Name:       node.Name.Name,
			Path:       packagePath,
			Imports:    []string{},
			ImportedBy: []string{},
		}
		graph.Packages[packagePath] = packageInfo
	}
	graph.Mutex.Unlock()

	// 收集导入
	var imports []string
	for _, imp := range node.Imports {
		importPath := strings.Trim(imp.Path.Value, `"`)
		imports = append(imports, importPath)
	}

	// 更新导入列表（去重）
	graph.Mutex.Lock()
	existingImports := make(map[string]bool)
	for _, imp := range packageInfo.Imports {
		existingImports[imp] = true
	}
	for _, imp := range imports {
		if !existingImports[imp] {
			packageInfo.Imports = append(packageInfo.Imports, imp)
		}
	}
	graph.Mutex.Unlock()
}

// calculateImportedBy 计算每个包被哪些包导入
func calculateImportedBy(graph *DependencyGraph) {
	graph.Mutex.Lock()
	defer graph.Mutex.Unlock()

	// 清空现有的importedBy
	for _, pkg := range graph.Packages {
		pkg.ImportedBy = []string{}
		pkg.IsExternal = false
		pkg.IsStandard = true // 先假设是标准库，后续会修正
	}

	// 遍历所有包，构建导入关系
	for pkgPath, pkg := range graph.Packages {
		for _, imp := range pkg.Imports {
			// 标记外部包和标准库
			if isStandardPackage(imp) {
				// 如果是标准库，检查是否在已有包中
				if _, exists := graph.Packages[imp]; !exists {
					graph.Packages[imp] = &PackageInfo{
						Name:       filepath.Base(imp),
						Path:       imp,
						IsStandard: true,
					}
				}
			} else {
				// 外部包
				if importedPkg, exists := graph.Packages[imp]; exists {
					importedPkg.IsExternal = true
					importedPkg.IsStandard = false
				}
			}

			// 添加importedBy关系
			if importedPkg, exists := graph.Packages[imp]; exists {
				importedPkg.ImportedBy = append(importedPkg.ImportedBy, pkgPath)
			}
		}
	}

	// 标记项目内部的包
	for pkgPath, pkg := range graph.Packages {
		if _, exists := graph.Packages[pkgPath]; exists && !pkg.IsExternal {
			// 如果路径中包含.或者没有/，可能是本地包
			if strings.Contains(pkgPath, ".") || !strings.Contains(pkgPath, "/") {
				pkg.IsStandard = true
			} else {
				pkg.IsStandard = false
			}
		}
	}
}

// isStandardPackage 检查是否是标准库包
func isStandardPackage(pkgPath string) bool {
	// 标准库通常不包含域名
	if strings.Contains(pkgPath, ".") {
		return false
	}

	// 常见的一级标准库目录
	stdLibs := map[string]bool{
		"fmt":       true,
		"io":        true,
		"net":       true,
		"http":      true,
		"os":        true,
		"strings":   true,
		"strconv":   true,
		"encoding":  true,
		"json":      true,
		"xml":       true,
		"time":      true,
		"sync":      true,
		"math":      true,
		"sort":      true,
		"container": true,
		"crypto":    true,
		"database":  true,
		"debug":     true,
		"embed":     true,
		"errors":    true,
		"expvar":    true,
		"flag":      true,
		"go":        true,
		"hash":      true,
		"html":      true,
		"image":     true,
		"index":     true,
		"log":       true,
		"mime":      true,
		"path":      true,
		"reflect":   true,
		"regexp":    true,
		"runtime":   true,
		"testing":   true,
		"text":      true,
		"unicode":   true,
		"unsafe":    true,
	}

	// 获取第一级目录
	firstPart := pkgPath
	if idx := strings.Index(pkgPath, "/"); idx != -1 {
		firstPart = pkgPath[:idx]
	}

	return stdLibs[firstPart]
}

// generateOutput 生成输出数据
func generateOutput(graph *DependencyGraph, ignoreStdLib bool) *DependencyOutput {
	graph.Mutex.RLock()
	defer graph.Mutex.RUnlock()

	output := &DependencyOutput{
		Packages: []PackageOutput{},
		Edges:    []DependencyEdge{},
	}

	// 收集包信息
	var packagePaths []string
	for pkgPath := range graph.Packages {
		packagePaths = append(packagePaths, pkgPath)
	}
	sort.Strings(packagePaths)

	for _, pkgPath := range packagePaths {
		pkg := graph.Packages[pkgPath]

		// 如果忽略标准库且当前包是标准库，则跳过
		if ignoreStdLib && pkg.IsStandard {
			continue
		}

		// 添加包信息
		pkgOutput := PackageOutput{
			Name:       pkg.Name,
			Path:       pkg.Path,
			Imports:    make([]string, len(pkg.Imports)),
			ImportedBy: make([]string, len(pkg.ImportedBy)),
			IsExternal: pkg.IsExternal,
			IsStandard: pkg.IsStandard,
		}
		copy(pkgOutput.Imports, pkg.Imports)
		copy(pkgOutput.ImportedBy, pkg.ImportedBy)
		output.Packages = append(output.Packages, pkgOutput)

		// 添加依赖边（排除指向被忽略的标准库的边）
		for _, imp := range pkg.Imports {
			if importedPkg, exists := graph.Packages[imp]; exists {
				if ignoreStdLib && importedPkg.IsStandard {
					continue
				}
				output.Edges = append(output.Edges, DependencyEdge{
					From: pkgPath,
					To:   imp,
				})
			}
		}
	}

	return output
}

// outputResult 输出结果
func outputResult(output *DependencyOutput, format string, outputFile string) error {
	var result string
	var err error

	switch format {
	case "json":
		result, err = outputJSON(output)
	case "dot":
		result, err = outputDOT(output)
	case "csv":
		result, err = outputCSV(output)
	default: // text
		result, err = outputText(output)
	}

	if err != nil {
		return err
	}

	// 输出到文件或控制台
	if outputFile != "" {
		err = os.WriteFile(outputFile, []byte(result), 0644)
		if err != nil {
			return err
		}
		fmt.Printf("结果已保存到: %s\n", outputFile)
	} else {
		fmt.Println(result)
	}

	return nil
}

// outputText 文本格式输出
func outputText(output *DependencyOutput) (string, error) {
	var builder strings.Builder

	builder.WriteString("=== 包依赖分析报告 ===\n\n")

	// 按类型分组
	var stdPackages, localPackages, extPackages []PackageOutput
	for _, pkg := range output.Packages {
		if pkg.IsStandard {
			stdPackages = append(stdPackages, pkg)
		} else if pkg.IsExternal {
			extPackages = append(extPackages, pkg)
		} else {
			localPackages = append(localPackages, pkg)
		}
	}

	// 输出本地包
	if len(localPackages) > 0 {
		builder.WriteString("本地包:\n")
		for _, pkg := range localPackages {
			builder.WriteString(fmt.Sprintf("  %s (%s)\n", pkg.Path, pkg.Name))
			if len(pkg.Imports) > 0 {
				builder.WriteString("    导入:\n")
				for _, imp := range pkg.Imports {
					builder.WriteString(fmt.Sprintf("      - %s\n", imp))
				}
			}
			if len(pkg.ImportedBy) > 0 {
				builder.WriteString("    被以下包导入:\n")
				for _, by := range pkg.ImportedBy {
					builder.WriteString(fmt.Sprintf("      - %s\n", by))
				}
			}
			builder.WriteString("\n")
		}
	}

	// 输出外部包
	if len(extPackages) > 0 {
		builder.WriteString("外部依赖:\n")
		for _, pkg := range extPackages {
			builder.WriteString(fmt.Sprintf("  %s\n", pkg.Path))
			if len(pkg.ImportedBy) > 0 {
				builder.WriteString("    被以下包导入:\n")
				for _, by := range pkg.ImportedBy {
					builder.WriteString(fmt.Sprintf("      - %s\n", by))
				}
			}
			builder.WriteString("\n")
		}
	}

	// 输出标准库（简略）
	if len(stdPackages) > 0 {
		builder.WriteString("标准库依赖:\n")
		libs := make(map[string]bool)
		for _, pkg := range stdPackages {
			libs[pkg.Path] = true
		}
		var libList []string
		for lib := range libs {
			libList = append(libList, lib)
		}
		sort.Strings(libList)
		for _, lib := range libList {
			builder.WriteString(fmt.Sprintf("  %s\n", lib))
		}
	}

	// 依赖统计
	builder.WriteString("\n=== 依赖统计 ===\n")
	builder.WriteString(fmt.Sprintf("总包数: %d\n", len(output.Packages)))
	builder.WriteString(fmt.Sprintf("本地包: %d\n", len(localPackages)))
	builder.WriteString(fmt.Sprintf("外部包: %d\n", len(extPackages)))
	builder.WriteString(fmt.Sprintf("标准库: %d\n", len(stdPackages)))
	builder.WriteString(fmt.Sprintf("依赖关系数: %d\n", len(output.Edges)))

	return builder.String(), nil
}

// outputJSON JSON格式输出
func outputJSON(output *DependencyOutput) (string, error) {
	data, err := json.MarshalIndent(output, "", "  ")
	if err != nil {
		return "", err
	}
	return string(data), nil
}

// outputDOT DOT格式输出（用于Graphviz）
func outputDOT(output *DependencyOutput) (string, error) {
	var builder strings.Builder

	builder.WriteString("digraph GoDependencies {\n")
	builder.WriteString("  rankdir=LR;\n")
	builder.WriteString("  node [shape=box, style=filled];\n\n")

	// 定义节点
	for _, pkg := range output.Packages {
		color := "lightblue"
		if pkg.IsExternal {
			color = "lightcoral"
		} else if pkg.IsStandard {
			color = "lightgrey"
		}

		builder.WriteString(fmt.Sprintf("  \"%s\" [label=\"%s\", fillcolor=\"%s\"];\n",
			pkg.Path, pkg.Path, color))
	}

	builder.WriteString("\n")

	// 定义边
	for _, edge := range output.Edges {
		builder.WriteString(fmt.Sprintf("  \"%s\" -> \"%s\";\n", edge.From, edge.To))
	}

	builder.WriteString("}\n")
	return builder.String(), nil
}

// outputCSV CSV格式输出
func outputCSV(output *DependencyOutput) (string, error) {
	var builder strings.Builder

	// 写表头
	builder.WriteString("From,To,Type\n")

	// 写数据
	for _, edge := range output.Edges {
		toPkg := findPackage(output.Packages, edge.To)

		depType := "local"
		if toPkg != nil {
			if toPkg.IsExternal {
				depType = "external"
			} else if toPkg.IsStandard {
				depType = "standard"
			}
		}

		builder.WriteString(fmt.Sprintf("%s,%s,%s\n", edge.From, edge.To, depType))
	}

	return builder.String(), nil
}

func findPackage(packages []PackageOutput, path string) *PackageOutput {
	for _, pkg := range packages {
		if pkg.Path == path {
			return &pkg
		}
	}
	return nil
}