import subprocess
from collections import defaultdict

def list_installed_packages():
    """获取所有已安装的包列表。"""
    result = subprocess.run(
        ["dpkg-query", "-f", "${binary:Package}\n", "-W"],
        stdout=subprocess.PIPE,
        text=True
    )
    return result.stdout.strip().split("\n")

def get_package_dependencies(package):
    """查询单个包的直接依赖项。"""
    result = subprocess.run(
        ["apt-cache", "depends", package],
        stdout=subprocess.PIPE,
        text=True
    )
    dependencies = []
    for line in result.stdout.strip().split("\n"):
        if line.strip().startswith("Depends:"):
            dep = line.split(":", 1)[1].strip()
            dep = dep.split(":")[0]  # 去掉冒号后的部分
            dependencies.append(dep)
    return dependencies

def build_dependency_graph(packages):
    """构建包的依赖关系图。"""
    totalPackages = len(packages)
    processedPackages = 0
    graph = defaultdict(list)
    for package in packages:
        dependencies = get_package_dependencies(package)
        for dep in dependencies:
            print(f"{package} -> {dep}")
            graph[package].append(dep)
        processedPackages += 1
        print(f"已处理 {processedPackages}/{totalPackages} 个包")
    return graph

def remove_redundant_edges(graph):
    """去除冗余的边。"""
    def dfs(node, visited):
        if node in visited:
            return visited[node]
        visited[node] = set()
        for neighbor in graph[node]:
            visited[node].update(dfs(neighbor, visited))
        visited[node].add(node)
        return visited[node]

    # 创建 graph 的静态副本来避免动态修改引发问题
    nodes = list(graph.keys())
    reachable = {}
    for node in nodes:  # 这里使用静态副本
        dfs(node, reachable)

    minimal_graph = defaultdict(list)
    for node in nodes:  # 再次使用静态副本
        direct_deps = set(graph[node])
        for dep in graph[node]:
            direct_deps -= reachable[dep]
        minimal_graph[node] = list(direct_deps)
    return minimal_graph

def generate_mermaid_graph(graph):
    """生成 Mermaid 图表的语法。"""
    lines = ["stateDiagram-v2"]
    for package, dependencies in graph.items():
        for dep in dependencies:
            lines.append(f"    {package} --> {dep}")
    return "\n".join(lines)

def main():
    print("正在获取已安装的包...")
    packages = list_installed_packages()

    print("正在构建依赖图...")
    graph = build_dependency_graph(packages)

    print("正在去除冗余边...")
    minimal_graph = remove_redundant_edges(graph)

    print("正在生成 Mermaid 图表语法...")
    mermaid_graph = generate_mermaid_graph(minimal_graph)

    with open("dependency_graph.mmd", "w") as file:
        file.write("---\n")
        file.write("title: APT Dependency Graph\n")
        file.write("---\n\n")
        file.write(mermaid_graph)

    print("Mermaid 图表已生成并保存为 dependency_graph.mmd")

if __name__ == "__main__":
    main()