#!/usr/bin/env python3
import os
import argparse
from collections import defaultdict
from datetime import datetime

IGNORED_DIRS = {
    '.git',
    '.git_backup',
    '__pycache__',
    'node_modules',
    'dist',
    'build',
    '.next',
    '.cache',
    '.venv',
    'venv',
}

ALLOWED_EXTS = {
    '.tsx',
    '.css',
    '.go',
    '.ts',
    '.js',
    '.html',
    '.sql',
    '.py',
}


def walk_project(root: str):
    for current_dir, dirs, files in os.walk(root):
        dirs[:] = [d for d in dirs if d not in IGNORED_DIRS]
        yield current_dir, dirs, files


def count_stats(root: str):
    total_dirs = 0
    total_files = 0
    total_lines = 0
    by_ext = defaultdict(lambda: [0, 0])  # ext -> [file_count, line_count]

    for current_dir, dirs, files in walk_project(root):
        if current_dir != root:
            total_dirs += 1

        for name in files:
            ext = os.path.splitext(name)[1]
            if ext not in ALLOWED_EXTS:
                continue

            total_files += 1
            path = os.path.join(current_dir, name)
            ext_key = ext or '<no_ext>'
            line_count = 0
            try:
                with open(path, 'r', encoding='utf-8', errors='ignore') as f:
                    for _ in f:
                        line_count += 1
            except Exception:
                line_count = 0

            total_lines += line_count
            by_ext[ext_key][0] += 1
            by_ext[ext_key][1] += line_count

    return total_dirs, total_files, total_lines, by_ext


def print_tree(root: str, max_depth: int | None = None):
    tree_output = []
    root = os.path.abspath(root)
    for current_dir, dirs, files in walk_project(root):
        rel = os.path.relpath(current_dir, root)
        if rel == '.':
            depth = 0
            name = os.path.basename(root.rstrip(os.sep)) or root
        else:
            depth = rel.count(os.sep) + 1
            name = os.path.basename(current_dir)

        if max_depth is not None and depth > max_depth:
            dirs[:] = []
            continue

        indent = '    ' * depth
        tree_output.append(f"{indent}{name}/")

        file_indent = '    ' * (depth + 1)
        for filename in sorted(files):
            tree_output.append(f"{file_indent}{filename}")
    
    return tree_output


def main():
    parser = argparse.ArgumentParser(description='Project statistics: files, folders, lines, and structure.')
    parser.add_argument('path', nargs='?', default='.', help='Root path (default: current directory)')
    parser.add_argument('--max-tree-depth', type=int, default=3, help='Max depth for printed tree (default: 3)')
    parser.add_argument('--output-md', action='store_true', help='Output to stats.md file instead of console')
    args = parser.parse_args()

    root = os.path.abspath(args.path)

    total_dirs, total_files, total_lines, by_ext = count_stats(root)
    tree_lines = print_tree(root, max_depth=args.max_tree_depth)

    if args.output_md:
        # Get script directory for output file
        script_dir = os.path.dirname(os.path.abspath(__file__))
        output_file = os.path.join(script_dir, 'stats.md')
        
        with open(output_file, 'w', encoding='utf-8') as f:
            f.write(f"# Project Statistics\n\n")
            f.write(f"**Generated:** {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
            f.write(f"**Path:** `{root}`\n\n")
            
            f.write("## Summary\n\n")
            f.write(f"- **Total directories:** {total_dirs}\n")
            f.write(f"- **Total files:** {total_files}\n")
            f.write(f"- **Total lines (approximate):** {total_lines}\n\n")
            
            f.write("## Lines by Extension\n\n")
            f.write("| Extension | Files | Lines |\n")
            f.write("|-----------|-------|-------|\n")
            
            for ext, (file_count, line_count) in sorted(by_ext.items(), key=lambda kv: kv[1][1], reverse=True):
                label = ext if ext else '<no_ext>'
                f.write(f"| {label} | {file_count} | {line_count} |\n")
            
            f.write(f"\n## Directory Tree (max depth {args.max_tree_depth})\n\n")
            f.write("```\n")
            for line in tree_lines:
                f.write(line + "\n")
            f.write("```\n")
        
        print(f"Statistics saved to: {output_file}")
    else:
        # Original console output
        print(f"Analyzing project at: {root}")
        print()

        print('=== SUMMARY ===')
        print(f"Total directories: {total_dirs}")
        print(f"Total files:       {total_files}")
        print(f"Total lines (approximate): {total_lines}")
        print()

        print('=== LINES BY EXTENSION ===')
        for ext, (file_count, line_count) in sorted(by_ext.items(), key=lambda kv: kv[1][1], reverse=True):
            label = ext if ext else '<no_ext>'
            print(f"{label:10}  files={file_count:6}  lines={line_count:10}")

        print()
        print(f"=== DIRECTORY TREE (max depth {args.max_tree_depth}) ===")
        for line in tree_lines:
            print(line)


if __name__ == '__main__':
    main()