import os
import argparse
from tabulate import tabulate
import csv

def find_env_files(directory):
    """Recursively traverse the directory and find all .env files."""
    env_files = []
    for root, _, files in os.walk(directory):
        for file in files:
            if file.endswith('.env'):
                env_files.append(os.path.join(root, file))
    return env_files

def read_env_file(file_path):
    """Read the contents of an .env file and return a list of key-value pairs."""
    key_value_pairs = []
    with open(file_path, 'r') as file:
        for line in file:
            line = line.strip()
            if line and not line.startswith('#') and '=' in line:
                key, value = line.split('=', 1)
                key_value_pairs.append((key.strip(), value.strip()))
    return key_value_pairs

def create_table(env_files, root_directory):
    """Create a table with directory name, .env file name, environment key, and values."""
    table = []
    for env_file in env_files:
        key_value_pairs = read_env_file(env_file)
        directory_name = os.path.dirname(os.path.relpath(env_file, start=root_directory))
        file_name = os.path.basename(env_file)
        for key, value in key_value_pairs:
            table.append([directory_name, file_name, key, value])
    return table

def save_to_csv(table, output_file):
    """Save the table data to a CSV file."""
    headers = ["Directory", ".env File Name", "Environment Key", "Value"]
    with open(output_file, 'w', newline='') as csvfile:
        writer = csv.writer(csvfile)
        writer.writerow(headers)
        writer.writerows(table)
    print(f"Data saved to {output_file}")

def main(directory, output_format, output_file):
    """Main function to find .env files, read their contents, and display/save the results."""
    env_files = find_env_files(directory)
    if not env_files:
        print(f"No .env files found in '{directory}'.")
        return
    
    table = create_table(env_files, directory)
    headers = ["Directory", ".env File Name", "Environment Key", "Value"]

    if output_format == "table":
        print(tabulate(table, headers=headers, tablefmt="grid"))
    elif output_format == "csv":
        save_to_csv(table, output_file)
    else:
        print("Invalid output format. Choose 'table' or 'csv'.")

if __name__ == "__main__":
    # Set up command-line argument parsing
    parser = argparse.ArgumentParser(description="Find and display .env files in a directory.")
    parser.add_argument("directory", help="Directory to search for .env files")
    parser.add_argument("--format", choices=["table", "csv"], default="table", help="Output format (table or csv)")
    parser.add_argument("--output", default="output.csv", help="Output file name for CSV format")
    args = parser.parse_args()

    if not os.path.isdir(args.directory):
        print(f"Error: '{args.directory}' is not a valid directory.")
    else:
        main(args.directory, args.format, args.output)