import os
from pathlib import Path

def generate_tree(path, prefix=""):
    """Generate tree structure"""
    items = []
    try:
        entries = sorted(path.iterdir(), key=lambda x: (x.is_file(), x.name))
        # Filter out ignored entries first
        filtered_entries = [e for e in entries if not should_ignore(e)]
        
        for i, entry in enumerate(filtered_entries):
            is_last = i == len(filtered_entries) - 1
            current_prefix = "└── " if is_last else "├── "
            
            items.append(f"{prefix}{current_prefix}{entry.name}")
            
            if entry.is_dir():
                next_prefix = prefix + ("    " if is_last else "│   ")
                items.extend(generate_tree(entry, next_prefix))
    except PermissionError:
        pass
    return items

def should_ignore(path):
    """Check if file/folder should be ignored"""
    # Explicitly check for virtual environments and common ignored folders
    if path.name in {'.venv', 'venv', '__pycache__', '.git', 'node_modules', '.idea', '.vscode'}:
        return True
    
    # Check if file is inside .venv or venv folder
    if '.venv' in path.parts or 'venv' in path.parts:
        return True
        
    # Ignore all hidden files/folders (starting with .)
    if path.name.startswith('.'):
        return True
    
    # Ignore specific files
    ignore_files = {
        'CODE_DOCUMENTATION.md', 'CODE_DOCUMENTATION.ipynb',
        'CODE_DOCUMENTATION.html', 'CODE_DOCUMENTATION.pdf',
        'CODE_DOCUMENTATION.docx', 'CODE_DOCUMENTATION.txt',
        'CODE_DOCUMENTATION.csv', 'CODE_DOCUMENTATION.xlsx',
        'CODE_DOCUMENTATION.pptx', 'CODE_DOCUMENTATION.ods',
        'CODE_DOCUMENTATION.odp', 'CODE_DOCUMENTATION.odt',
        'uv.lock', 'poetry.lock', 'Pipfile.lock',
        '.DS_Store'
    }
    
    # Ignore by file extension
    ignore_extensions = {'.pyc', '.pyo', '.pyd', '.so', '.egg-info'}
    
    return (path.name in ignore_files or 
            path.suffix in ignore_extensions or
            path.name.endswith('.egg-info'))

def get_code_files(directory):
    """Get all relevant code files"""
    code_extensions = {'.py', '.js', '.ts', '.html', '.css', '.sql', '.yaml', '.yml', '.json', '.toml', '.cfg', '.ini'}
    code_files = []
    
    for file_path in directory.rglob("*"):
        # Skip if it's a directory
        if file_path.is_dir():
            continue
            
        # Skip ignored files/folders
        if should_ignore(file_path):
            continue
            
        # Only include files with relevant extensions
        if file_path.suffix.lower() in code_extensions:
            code_files.append(file_path)
    
    return sorted(code_files)

def main():
    current_dir = Path.cwd()
    output_file = current_dir / "CODE_DOCUMENTATION.md"
    
    # Generate markdown
    markdown = f"# {current_dir.name}\n\n"
    markdown += f"Generated on: {Path.cwd()}\n\n"
    
    # Add tree structure
    markdown += "## Project Structure\n\n```\n"
    markdown += f"{current_dir.name}/\n"
    tree_items = generate_tree(current_dir)
    for item in tree_items:
        markdown += f"{item}\n"
    markdown += "```\n\n"
    
    # Get all code files
    code_files = get_code_files(current_dir)
    
    if code_files:
        markdown += "## Source Code\n\n"
        
        for file_path in code_files:
            try:
                with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
                    content = f.read()
                
                rel_path = file_path.relative_to(current_dir)
                file_extension = file_path.suffix.lstrip('.')
                
                # Use appropriate syntax highlighting
                if file_extension == 'py':
                    lang = 'python'
                elif file_extension in ['js', 'ts']:
                    lang = 'javascript'
                elif file_extension in ['html']:
                    lang = 'html'
                elif file_extension in ['css']:
                    lang = 'css'
                elif file_extension in ['sql']:
                    lang = 'sql'
                elif file_extension in ['yaml', 'yml']:
                    lang = 'yaml'
                elif file_extension in ['json']:
                    lang = 'json'
                else:
                    lang = file_extension
                
                markdown += f"### {rel_path}\n\n"
                markdown += f"```{lang}\n{content}\n```\n\n"
                
            except Exception as e:
                markdown += f"### {rel_path}\n\n"
                markdown += f"*Could not read file: {str(e)}*\n\n"
                continue
    else:
        markdown += "## Source Code\n\n*No code files found.*\n\n"
    
    # Write output
    try:
        with open(output_file, 'w', encoding='utf-8') as f:
            f.write(markdown)
        print(f"✅ Documentation generated successfully: {output_file}")
        print(f"📁 Total files documented: {len(code_files)}")
    except Exception as e:
        print(f"❌ Error writing documentation: {str(e)}")

if __name__ == "__main__":
    main()