#!/usr/bin/env python3
import os
import sys
import subprocess
import shutil
import argparse
import time # Keep time import

# Add these near the top
FIXER_SCRIPT_NAME = "fix_attributes.py" 

def main(carbon_script_path, source_dir, target_dir, fixer_script_path=None, fixer_rules=None): # Add fixer args
    """
    Runs the carbon.py obfuscator, optionally followed by a fixer script.
    """
    # --- Input validation (keep as before) ---
    if not os.path.isfile(carbon_script_path):
        print(f"Error: Carbon script not found at '{carbon_script_path}'")
        sys.exit(1)
    if not os.path.isdir(source_dir):
        print(f"Error: Source directory not found at '{source_dir}'")
        sys.exit(1)
    # Add validation for fixer script if provided
    if fixer_script_path and not os.path.isfile(fixer_script_path):
        print(f"Error: Fixer script specified but not found at '{fixer_script_path}'")
        sys.exit(1)
    if fixer_script_path and not fixer_rules:
         print("Warning: Fixer script specified, but no rules provided via --fixer-rules.")
    if fixer_rules and not fixer_script_path:
         print("Warning: Fixer rules provided, but no fixer script specified via --fixer-script. Rules will be ignored.")


    # --- Path setup (keep as before) ---
    carbon_script_path = os.path.abspath(carbon_script_path)
    source_dir = os.path.abspath(source_dir)
    target_dir = os.path.abspath(target_dir)
    if fixer_script_path:
        fixer_script_path = os.path.abspath(fixer_script_path)


    # --- Target directory handling (keep as before) ---
    try:
        os.makedirs(target_dir, exist_ok=True)
        print(f"Ensured target directory exists: '{target_dir}'")
    except OSError as e:
        print(f"Error creating target directory '{target_dir}': {e}")
        sys.exit(1)

    # --- File iteration setup (keep as before) ---
    print(f"Scanning source directory: '{source_dir}'...")
    start_time = time.time()
    processed_count = 0
    error_count = 0
    skipped_count = 0
    fixed_count = 0 # Track files processed by fixer
    error_files = []

    try:
        filenames = sorted(os.listdir(source_dir))
    except OSError as e:
        print(f"Error listing files in source directory '{source_dir}': {e}")
        sys.exit(1)

    total_files = len(filenames)
    print(f"Found {total_files} items in source directory.")

    # --- Main Processing Loop ---
    for filename in filenames:
        source_file_path = os.path.join(source_dir, filename)
        pairs_file_path = source_file_path + ".pairs.json" # Expected pairs file path

        # --- Skip logic (keep as before, maybe add skipping .pairs.json) ---
        if not filename.endswith(".py") or \
           filename.endswith("-obf.py") or \
           filename.endswith(".pairs.json") or \
           not os.path.isfile(source_file_path) or \
           source_file_path == carbon_script_path or \
           (fixer_script_path and source_file_path == fixer_script_path):
            skipped_count += 1
            # Clean up stray pairs files if the source wasn't processed
            if os.path.exists(pairs_file_path):
                try: os.remove(pairs_file_path)
                except OSError: pass
            continue

        print(f"\nProcessing: '{filename}'...")

        base_name = filename[:-3] # Remove .py
        temp_obfuscated_path = os.path.join(source_dir, f"{base_name}-obf.py")
        final_target_path = os.path.join(target_dir, f"{base_name}-obf.py")

        # --- Pre-run Checks (keep as before) ---
        if os.path.exists(temp_obfuscated_path):
             print(f"  Warning: Removing existing temporary file '{temp_obfuscated_path}'")
             try: os.remove(temp_obfuscated_path)
             except OSError as e:
                 print(f"  Error removing existing temp file: {e}"); error_count += 1; error_files.append(filename); continue
        # Also remove potentially stale pairs file before running carbon
        if os.path.exists(pairs_file_path):
             try: os.remove(pairs_file_path)
             except OSError as e:
                 print(f"  Warning: Could not remove stale pairs file '{pairs_file_path}': {e}")


        obfuscation_success = False # Flag to track if carbon step worked
        try:
            # --- Run Carbon Obfuscator (keep as before) ---
            print(f"  Running obfuscator via: {os.path.basename(carbon_script_path)}")
            process_carbon = subprocess.run(
                [sys.executable, carbon_script_path],
                input=source_file_path + '\n', text=True, capture_output=True,
                check=False, encoding='utf-8', timeout=120
            )

            # --- Check Carbon Output ---
            if os.path.exists(temp_obfuscated_path):
                obfuscation_success = True # Mark success for this stage
                print(f"  Obfuscation successful (temporary file created).")
            else:
                 # Carbon failed to create output
                 print(f"  Error: Obfuscated file not found at '{temp_obfuscated_path}' after running carbon.py.")
                 print(f"  Carbon script exit code: {process_carbon.returncode}")
                 if process_carbon.stdout: print(f"  --- Carbon stdout ---:\n{process_carbon.stdout.strip()}\n  --- End Carbon stdout ---")
                 if process_carbon.stderr: print(f"  --- Carbon stderr ---:\n{process_carbon.stderr.strip()}\n  --- End Carbon stderr ---")
                 error_count += 1
                 error_files.append(filename)
                 # Attempt cleanup of pairs file if carbon failed
                 if os.path.exists(pairs_file_path):
                     try: os.remove(pairs_file_path)
                     except OSError: pass
                 continue # Skip to next file

            # --- Run Fixer Script (if specified and carbon succeeded) ---
            fixer_success = True # Assume success if fixer not run or runs ok
            if obfuscation_success and fixer_script_path and fixer_rules:
                if os.path.exists(pairs_file_path):
                    print(f"  Running fixer script: {os.path.basename(fixer_script_path)}")
                    try:
                        process_fixer = subprocess.run(
                            [
                                sys.executable, fixer_script_path,
                                "--obfuscated-file", temp_obfuscated_path,
                                "--pairs-file", pairs_file_path,
                                "--rules", fixer_rules
                            ],
                            capture_output=True, text=True, check=False, # Check returncode manually
                            encoding='utf-8', timeout=60 # Shorter timeout for fixer
                        )
                        if process_fixer.returncode == 0:
                            print(f"  Fixer script finished successfully.")
                            if process_fixer.stdout.strip(): # Print fixer output if any
                                print(f"  --- Fixer stdout ---:\n{process_fixer.stdout.strip()}\n  --- End Fixer stdout ---")
                            fixed_count +=1
                        else:
                            # Fixer script reported an error
                            print(f"  Error: Fixer script failed with exit code {process_fixer.returncode}.")
                            if process_fixer.stdout: print(f"  --- Fixer stdout ---:\n{process_fixer.stdout.strip()}\n  --- End Fixer stdout ---")
                            if process_fixer.stderr: print(f"  --- Fixer stderr ---:\n{process_fixer.stderr.strip()}\n  --- End Fixer stderr ---")
                            fixer_success = False
                            error_count += 1
                            error_files.append(filename)

                    except subprocess.TimeoutExpired:
                        print(f"  Error: Timeout expired while running fixer script on '{filename}'.")
                        fixer_success = False; error_count += 1
                        error_files.append(filename)
                    except Exception as e:
                        print(f"  Error: Failed running fixer script: {e}")
                        fixer_success = False; error_count += 1
                        error_files.append(filename)
                else:
                    # Pairs file missing, fixer cannot run reliably
                    print(f"  Warning: Cannot run fixer script because pairs file '{pairs_file_path}' is missing.")
                    # Decide if this should be counted as an error or just a warning
                    # error_count += 1
                    # fixer_success = False # Treat as failure if pairs file is mandatory for fixer

            # --- Move Final File (only if carbon and fixer succeeded) ---
            if obfuscation_success and fixer_success:
                try:
                    print(f"  Moving '{temp_obfuscated_path}' to '{final_target_path}'...")
                    shutil.move(temp_obfuscated_path, final_target_path)
                    processed_count += 1
                except OSError as e:
                    print(f"  Error moving final obfuscated file: {e}")
                    error_count += 1
                    error_files.append(filename)
                    # Attempt cleanup of temp file if move failed
                    if os.path.exists(temp_obfuscated_path):
                        try: os.remove(temp_obfuscated_path)
                        except OSError: pass

        # --- Catch exceptions during the processing of a single file ---
        except subprocess.TimeoutExpired:
             print(f"  Error: Timeout expired while running carbon.py on '{filename}'.")
             error_count += 1
             error_files.append(filename)
        except FileNotFoundError:
            print(f"  Critical Error: Failed to execute Python interpreter ('{sys.executable}') or Carbon script ('{carbon_script_path}').")
            sys.exit(1)
        except Exception as e:
            print(f"  An unexpected error occurred during processing '{filename}': {e}")
            error_count += 1
            error_files.append(filename)
        finally:
             # --- Cleanup: Always remove the pairs file ---
             if os.path.exists(pairs_file_path):
                 try:
                     os.remove(pairs_file_path)
                 except OSError as e:
                     print(f"  Warning: Could not remove pairs file '{pairs_file_path}': {e}")
             # Cleanup temp obfuscated file if it still exists (e.g., move failed)
             if os.path.exists(temp_obfuscated_path) and not os.path.exists(final_target_path):
                  print(f"  Cleaning up leftover temp file: {temp_obfuscated_path}")
                  try: os.remove(temp_obfuscated_path)
                  except OSError: pass


    # --- Final Report (adjust counts) ---
    end_time = time.time()
    duration = end_time - start_time
    print("\n" + "=" * 40)
    print("Batch Obfuscation Summary:")
    print(f"  Source Directory: '{source_dir}'")
    print(f"  Target Directory: '{target_dir}'")
    print("-" * 40)
    print(f"  Processed successfully: {processed_count} file(s)")
    if fixer_script_path:
        print(f"  Files processed by fixer: {fixed_count} file(s)") # Add fixer count
    print(f"  Skipped files:        {skipped_count} (non-.py, *-obf.py, etc.)")
    print(f"  Encountered errors:   {error_count} file(s)")

    if error_files:
        print("  --------------------")
        print("  Files with errors:")
        # Use set to ensure uniqueness if a file somehow errored multiple times
        for error_file in sorted(list(set(error_files))):
            print(f"    - {error_file}")
        print("  --------------------")

    print(f"  Total items scanned:  {total_files}")
    print(f"  Duration:             {duration:.2f} seconds")
    print("=" * 40)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Batch obfuscate Python files using carbon.py, optionally applying post-processing fixes.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    # --- Add arguments for fixer script ---
    parser.add_argument("-c", "--carbon-script", required=True, help="Path to the carbon.py obfuscator script.")
    parser.add_argument("-s", "--source-dir", required=True, help="Directory containing the Python scripts to obfuscate.")
    parser.add_argument("-t", "--target-dir", required=True, help="Directory where the obfuscated scripts will be saved.")
    parser.add_argument("--fixer-script", default=None, help=f"Optional path to the fixer script (e.g., '{FIXER_SCRIPT_NAME}').")
    parser.add_argument("--fixer-rules", default=None, help="Optional rules string for the fixer script, e.g., 'Action:C,D Base:Attr1'")

    args = parser.parse_args()

    # --- Keep validation/warning about carbon script in source dir ---
    if os.path.abspath(args.source_dir) == os.path.dirname(os.path.abspath(args.carbon_script)):
        print("Warning: The carbon script appears to be inside the source directory.")

    # Pass new args to main
    main(args.carbon_script, args.source_dir, args.target_dir, args.fixer_script, args.fixer_rules)