import os
import sys
import argparse
from pathlib import Path

# Enable ANSI colors for Windows
if os.name == 'nt':
    os.system("")

class C:
    CYAN = '\033[96m'
    GREEN = '\033[92m'
    YELLOW = '\033[93m'
    RED = '\033[91m'
    RESET = '\033[0m'
    BOLD = '\033[1m'

SKIP_DIRS = {
    ".git", "venv", ".venv", "env", ".env", "envs", "node_modules", 
    "__pycache__", "build", "dist", ".idea", ".vscode", "site-packages"
}

def print_banner():
    banner = fr"""{C.CYAN}{C.BOLD}
  _   _           _    _____                 
 | \ | |         | |  / ____|                
 |  \| | ___  ___| |_| (___   ___ __ _ _ __  
 | . ` |/ _ \/ __| __|\___ \ / __/ _` | '_ \ 
 | |\  |  __/\__ \ |_ ____) | (_| (_| | | | |
 |_| \_|\___||___/\__|_____/ \___\__,_|_| |_|   
                                      
 | Developed by : turnt ducky 🦆 |{C.RESET}"""
    print(banner)

def generate_tree_lines(dir_path: Path, prefix: str = "", show_all: bool = False) -> list:
    """
    Recursively builds the file structure and returns a list of strings.
    """
    lines = []
    try:
        # Filter and sort contents
        contents = []
        for p in dir_path.iterdir():
            if not show_all and p.name in SKIP_DIRS:
                continue
            contents.append(p)
            
        contents.sort(key=lambda x: (not x.is_dir(), x.name.lower()))

        pointers = [("├── ", "│   "), ("└── ", "    ")]

        for i, path in enumerate(contents):
            is_last = (i == len(contents) - 1)
            pointer, indent = pointers[1] if is_last else pointers[0]

            lines.append(f"{prefix}{pointer}{path.name}")

            if path.is_dir():
                lines.extend(generate_tree_lines(path, prefix + indent, show_all))

    except PermissionError:
        lines.append(f"{prefix}└── [Permission Denied]")
    except Exception as e:
        lines.append(f"{prefix}└── [Error: {str(e)}]")
    
    return lines

def main():
    parser = argparse.ArgumentParser(description="Generate and display a clean directory tree.")
    parser.add_argument("path", nargs="?", default=".", help="The folder path to scan (defaults to current directory)")
    parser.add_argument("-o", "--output", dest="output_file", help="Path to save the output text file (e.g., -o tree.txt)")
    parser.add_argument("--all", action="store_true", help="Include hidden/junk directories like node_modules and .git")
    
    args = parser.parse_args()

    # Clean up quotes if the user pasted them
    user_input = args.path.strip()
    if user_input.startswith('"') and user_input.endswith('"'):
        user_input = user_input[1:-1]

    target_dir = Path(user_input).resolve()

    print_banner()

    if not target_dir.exists():
        print(f" {C.RED}[!]{C.RESET} Error: The path '{target_dir}' does not exist.")
        sys.exit(1)
    if not target_dir.is_dir():
        print(f" {C.RED}[!]{C.RESET} Error: The path '{target_dir}' is not a directory.")
        sys.exit(1)

    print(f" {C.CYAN}[*]{C.RESET} Scanning: {C.YELLOW}{target_dir}{C.RESET}")

    # Build the output list
    output_lines = []
    output_lines.append(f"Directory Tree for: {target_dir}")
    output_lines.append("=" * 50)
    output_lines.append(target_dir.name or str(target_dir))
    
    tree_lines = generate_tree_lines(target_dir, show_all=args.all)
    
    if not tree_lines:
        print(f" {C.RED}[!]{C.RESET} Directory is empty or all contents were skipped.")
        sys.exit(0)
        
    output_lines.extend(tree_lines)
    final_output = "\n".join(output_lines)

    # Print to console (skip if it's massive, or just print it)
    print("\n" + final_output + "\n")

    # Save to file if requested
    if args.output_file:
        out_path = Path(args.output_file).resolve()
        try:
            with open(out_path, "w", encoding="utf-8") as f:
                f.write(final_output + "\n")
            print(f" {C.GREEN}[✓]{C.RESET} Success! Tree saved to {C.YELLOW}{out_path}{C.RESET}")
        except Exception as e:
            print(f" {C.RED}[!]{C.RESET} Error saving to file: {e}")
    else:
        print(f" {C.GREEN}[✓]{C.RESET} Scan complete.")

if __name__ == "__main__":
    main()