from pathlib import Path

output_filename = "Folder_structure.txt"
folders_to_skip = ["__pycache__", "test", "venv", "build"]
  
def print_tree(directory_path, max_level=2, dirs_only=True, 
  ignore_hidden=True, exclude_dirs=None, current_level=1, 
  prefix="", output_file="Folder_tree.txt"):
  """
    Prints tree structure and optionally writes it to a file
  
    directory_path: The target root folder path
    max_level:      Maximum depth to traverse (int). None means unlimited
    dirs_only:      If True, only displays directory folders
    ignore_hidden:  If True, skips files and folders starting with '.'
    exclude_dirs:   A list of folder names to skip (e.g. ['__pycache__'])
    current_level:  Internal counter tracking current depth
    prefix:         Internal string builder for visual branch styling
    output_file:    File object to write the tree output to.
  """
  path = Path(directory_path)
  if exclude_dirs is None:
    exclude_dirs = set()
  else:
    exclude_dirs = set(exclude_dirs)
  
  # Base case: Stop traversal if we exceed depth limits
  if max_level is not None and current_level > max_level:
    return

  # Print and save the root folder on the first execution
  if current_level == 1:
    root_name = path.name or str(path)
    print(root_name)
    if output_file:
      output_file.write(root_name + "\n")

  try:
    entries = sorted(list(path.iterdir()), key=lambda e: e.name.lower())
  except PermissionError:
    return # Skip locked system folders gracefully

  # Filtering Logic
  filtered_entries = []
  for entry in entries:
    
    # Skip hidden assets if option is enabled (default True)
    if ignore_hidden and entry.name.startswith('.'):
      continue
      
    # kip specifically excluded directory names
    if entry.is_dir() and entry.name in exclude_dirs:
      continue
      
    # Skip individual files if dirs_only is active (default True)
    if dirs_only and not entry.is_dir():
      continue
      
    filtered_entries.append(entry)

  # Visual Tree Rendering Execution
  count = len(filtered_entries)
  for index, entry in enumerate(filtered_entries):
    is_last = (index == count - 1)
    connector = "└── " if is_last else "├── "
    
    # Build line text
    line = f"{prefix}{connector}{entry.name}"
    print(line)
    if output_file:
      output_file.write(line + "\n")

    # Recursively step into child folders
    if entry.is_dir():
      next_prefix = prefix + ("  " if is_last else "│   ")
      print_tree(
        entry, 
        max_level=max_level, 
        dirs_only=dirs_only, 
        ignore_hidden=ignore_hidden,
        exclude_dirs=exclude_dirs,
        current_level=current_level + 1, 
        prefix=next_prefix,
        output_file=output_file
      )

if __name__ == "__main__":
  target_folder = "."

  print(f"Generating tree and saving to '{output_filename}'...\n")
  
  # Open the file using UTF-8 encoding to support tree unicode lines safely
  with open(output_filename, "w", encoding="utf-8") as f:
    print_tree(
      target_folder, 
      max_level=3, 
      dirs_only=True, 
      ignore_hidden=True, 
      exclude_dirs=folders_to_skip,
      output_file=f
    )
    
  print(f"\nDone! File saved successfully.")