import os
import ast
from pathlib import Path

def get_docstring(node):
    """Extract docstring from a node."""
    return ast.get_docstring(node) or ""

def analyze_file(filepath):
    """Analyze a Python file and extract structure."""
    try:
        with open(filepath, 'r', encoding='utf-8') as f:
            tree = ast.parse(f.read())
    except Exception as e:
        return {'error': str(e)}
    
    result = {
        'classes': [],
        'functions': [],
        'imports': []
    }
    
    for node in ast.walk(tree):
        if isinstance(node, ast.ClassDef):
            methods = []
            for item in node.body:
                if isinstance(item, ast.FunctionDef):
                    methods.append({
                        'name': item.name,
                        'docstring': get_docstring(item)
                    })
            
            result['classes'].append({
                'name': node.name,
                'docstring': get_docstring(node),
                'methods': methods
            })
        
        elif isinstance(node, ast.FunctionDef) and node.col_offset == 0:
            # Only top-level functions
            result['functions'].append({
                'name': node.name,
                'docstring': get_docstring(node)
            })
    
    return result

def generate_project_map(root_dir, output_file='project_map.txt', exclude_dirs=None):
    """Generate a project map for LLM consumption."""
    if exclude_dirs is None:
        exclude_dirs = {'__pycache__', '.git', 'venv', 'env', '.venv', 'node_modules', '.pytest_cache'}
    
    root_path = Path(root_dir)
    lines = []
    
    lines.append("=" * 80)
    lines.append(f"PROJECT MAP: {root_path.name}")
    lines.append("=" * 80)
    lines.append("")
    
    # Collect all Python files
    py_files = []
    for root, dirs, files in os.walk(root_path):
        # Remove excluded directories
        dirs[:] = [d for d in dirs if d not in exclude_dirs]
        
        for file in sorted(files):
            if file.endswith('.py'):
                py_files.append(Path(root) / file)
    
    # Analyze each file
    for filepath in sorted(py_files):
        rel_path = filepath.relative_to(root_path)
        lines.append(f"\n{'─' * 80}")
        lines.append(f"FILE: {rel_path}")
        lines.append('─' * 80)
        
        analysis = analyze_file(filepath)
        
        if 'error' in analysis:
            lines.append(f"  ⚠ Error parsing file: {analysis['error']}")
            continue
        
        # Classes
        if analysis['classes']:
            lines.append("\n  CLASSES:")
            for cls in analysis['classes']:
                lines.append(f"    • {cls['name']}")
                if cls['docstring']:
                    doc_preview = cls['docstring'].split('\n')[0][:60]
                    lines.append(f"      └─ {doc_preview}")
                
                if cls['methods']:
                    lines.append(f"      Methods:")
                    for method in cls['methods']:
                        lines.append(f"        - {method['name']}()")
        
        # Functions
        if analysis['functions']:
            lines.append("\n  FUNCTIONS:")
            for func in analysis['functions']:
                lines.append(f"    • {func['name']}()")
                if func['docstring']:
                    doc_preview = func['docstring'].split('\n')[0][:60]
                    lines.append(f"      └─ {doc_preview}")
        
        if not analysis['classes'] and not analysis['functions']:
            lines.append("  (No classes or functions found)")
    
    # Write to file
    output = '\n'.join(lines)
    with open(output_file, 'w', encoding='utf-8') as f:
        f.write(output)
    
    print(f"Project map generated: {output_file}")
    print(f"Total files analyzed: {len(py_files)}")
    return output

# Usage
if __name__ == "__main__":
    # Change this to your project directory
    project_dir = "."
    
    # Generate the map
    generate_project_map(project_dir, output_file="project_map.txt")
    
    # Also print to console
    with open("project_map.txt", 'r') as f:
        print(f.read())