#!/usr/bin/env python3
import os
import sys
import argparse
from pathlib import Path

if sys.platform == "win32":
    sys.stdout.reconfigure(encoding="utf-8")
    sys.stderr.reconfigure(encoding="utf-8")


def count_lines(file_path: Path) -> int:
    try:
        with open(file_path, "r", encoding="utf-8", errors="ignore") as f:
            return sum(1 for _ in f)
    except Exception:
        return 0


def should_include(file_path: Path, extensions: list[str]) -> bool:
    if not extensions:
        return True
    return file_path.suffix in extensions


def walk_dir(root_path: Path, extensions: list[str], exclude_dirs: set[str]):
    results = []
    total_lines = 0

    for root, dirs, files in os.walk(root_path):
        dirs[:] = [d for d in dirs if d not in exclude_dirs]

        rel_root = Path(root).relative_to(root_path)
        indent = len(rel_root.parts) if str(rel_root) != "." else 0
        prefix = "  " * indent

        if indent > 0:
            print(f"{prefix}└── {rel_root.name}/")

        for i, file in enumerate(files):
            file_path = Path(root) / file

            if not should_include(file_path, extensions):
                continue

            lines = count_lines(file_path)
            total_lines += lines
            results.append((file_path, lines, indent + 1))

            is_last = (i == len(files) - 1) and not any(
                should_include(Path(root) / f, extensions) for f in dirs
            )
            connector = "└──" if is_last else "├──"

            print(f"{'  ' * (indent + 1)}{connector} {file} ({lines} lines)")

    return total_lines


def main_all():
    directories = {
        ".": [".cpp", ".h", ".hlsl"],
        "docs": [".md"],
    }

    total_all = 0
    results = []

    for directory, extensions in directories.items():
        root_path = Path(directory)
        if not root_path.exists():
            continue

        exclude_dirs = {
            ".git",
            "build",
            "Release",
            "bin",
            "__pycache__",
            ".ruff_cache",
            "stbi",
        }

        print(f"\n{'=' * 60}")
        print(f"项目文件统计: {root_path}")
        print(f"后缀过滤: {extensions}")
        print(f"{'=' * 60}\n")

        total = walk_dir(root_path, extensions, exclude_dirs)
        results.append((directory, total))
        total_all += total

    # 单独统计 stbi
    stbi_path = Path("stbi")
    if stbi_path.exists():
        exclude_dirs = {".git", "__pycache__"}
        print(f"\n{'=' * 60}")
        print(f"项目文件统计: stbi (第三方库)")
        print(f"后缀过滤: ['.cpp', '.h']")
        print(f"{'=' * 60}\n")

        total = walk_dir(stbi_path, [".cpp", ".h"], exclude_dirs)
        results.append(("stbi", total))
        total_all += total

    print(f"\n{'=' * 60}")
    print("汇总统计")
    print(f"{'=' * 60}")
    for name, lines in results:
        print(f"{name:15} {lines:>10,} 行")
    print(f"{'=' * 60}")
    print(f"{'总计':15} {total_all:>10,} 行")
    print(f"{'=' * 60}\n")


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="统计项目文件行数")
    parser.add_argument(
        "-e", "--extension", action="append", help="指定后缀名，如: .py .ts .js"
    )
    parser.add_argument("-d", "--directory", default=".", help="指定子文件夹路径")
    parser.add_argument(
        "-x", "--exclude", action="append", default=[], help="排除的文件夹"
    )
    parser.add_argument(
        "--all",
        action="store_true",
        help="统计所有源码目录 (./docs/stbi)",
    )

    args = parser.parse_args()

    if args.all:
        main_all()
    else:
        root_path = Path(args.directory).resolve()
        extensions = args.extension
        exclude_dirs = {
            ".git",
            "build",
            "Release",
            "bin",
            "__pycache__",
        }
        exclude_dirs.update(args.exclude)

        if not root_path.exists():
            print(f"错误: 目录 {root_path} 不存在")
            sys.exit(1)

        print(f"\n{'=' * 60}")
        print(f"项目文件统计: {root_path}")
        if extensions:
            print(f"后缀过滤: {extensions}")
        print(f"{'=' * 60}\n")

        total = walk_dir(root_path, extensions, exclude_dirs)

        print(f"\n{'=' * 60}")
        print(f"总行数: {total}")
        print(f"{'=' * 60}\n")