#!/usr/bin/env python3
import json
import argparse
import sys
import os
import re # Import re

def parse_rules(rules_str):
    """ Parses rules like 'Action:C,D Another:Attr1' into a dictionary """
    rules = {}
    if not rules_str: return rules
    try:
        parts = rules_str.strip().split()
        for part in parts:
            if ':' not in part: raise ValueError(f"Rule part '{part}' missing ':' separator")
            base, attrs_str = part.split(':', 1)
            attrs = [attr.strip() for attr in attrs_str.split(',') if attr.strip()]
            if not base or not attrs: raise ValueError(f"Invalid rule part '{part}': base or attrs empty")
            rules[base.strip()] = attrs
    except Exception as e:
        print(f"Error: Could not parse rules string '{rules_str}'. Format: 'Base1:Attr1,Attr2 Base2:AttrX'. Error: {e}", file=sys.stderr)
        sys.exit(1)
    return rules


def apply_fixes(obfuscated_file_path, pairs_file_path, rules_dict):
    """
    Reads the obfuscated file, applies attribute preservation rules,
    and replaces 'name = "..."' strings with 'name = "strategy"'.
    """
    try:
        with open(obfuscated_file_path, 'r', encoding='utf-8') as f:
            code = f.read()
    except FileNotFoundError:
        print(f"Error: Fixer script cannot find obfuscated file: {obfuscated_file_path}", file=sys.stderr)
        sys.exit(1)
    except Exception as e:
        print(f"Error: Fixer script failed reading obfuscated file {obfuscated_file_path}: {e}", file=sys.stderr)
        sys.exit(1)

    # Load the name mappings (pairs) if the file exists
    pairs = {}
    if os.path.exists(pairs_file_path):
        try:
            with open(pairs_file_path, 'r', encoding='utf-8') as f:
                pairs = json.load(f)
        except json.JSONDecodeError as e:
             print(f"Error: Fixer script failed decoding JSON from pairs file {pairs_file_path}: {e}", file=sys.stderr)
             sys.exit(1) # Exit if pairs file is corrupt
        except Exception as e:
             print(f"Error: Fixer script failed reading pairs file {pairs_file_path}: {e}", file=sys.stderr)
             sys.exit(1)
    else:
        print(f"Warning: Fixer script did not find pairs file '{pairs_file_path}'. Dynamic name variable matching might default to 'name'.", file=sys.stderr)


    original_code = code # Keep a copy to see if changes were made
    fixes_applied_count = 0
    masking_applied_count = 0
    regex_matches_found = 0

    # --- 1. Apply the attribute preservation logic (existing code) ---
    print(f"  Fixer: Applying attribute preservation rules to '{os.path.basename(obfuscated_file_path)}'...")
    for base_name, attrs_to_preserve in rules_dict.items():
        obfuscated_base_name = pairs.get(base_name, base_name)
        for original_attr in attrs_to_preserve:
            obfuscated_attr_name = pairs.get(original_attr)
            if obfuscated_attr_name:
                string_to_find = f"{obfuscated_base_name}.{obfuscated_attr_name}"
                replacement_string = f"{obfuscated_base_name}.{original_attr}"
                if string_to_find in code:
                    initial_count = code.count(string_to_find)
                    code = code.replace(string_to_find, replacement_string)
                    fixes_applied_count += initial_count
    if fixes_applied_count > 0:
        print(f"  Fixer: Applied {fixes_applied_count} attribute restorations based on rules.")


    # --- 2. Replace 'varName = "OriginalName"' string with 'varName = "strategy"' ---
    print(f"  Fixer: Checking for name assignments to replace with 'strategy' in '{os.path.basename(obfuscated_file_path)}'...")

    # Find the potentially obfuscated name for the variable 'name' itself
    obfuscated_name_var = pairs.get("name", "name") # Default to "name" if not found in pairs

    # Dynamically create the regex pattern using the (potentially obfuscated) variable name
    name_pattern_str = rf'(^\s*(?:self\.)?{re.escape(obfuscated_name_var)}\s*=\s*["\'])([^"\']+)(["\'])'
    name_pattern = re.compile(name_pattern_str, re.MULTILINE)

    replacement_name = "strategy" # The fixed string to insert

    def replace_with_strategy_string(match):
        nonlocal masking_applied_count, regex_matches_found
        regex_matches_found += 1
        prefix = match.group(1) # e.g., ' llIllIIIIllllllIIl = "'
        original_name_in_string = match.group(2) # e.g., 'ZD-SET-2'
        suffix = match.group(3) # e.g., '"'

        # Replace if the current string is not already 'strategy'
        if original_name_in_string != replacement_name:
            masking_applied_count += 1
            return prefix + replacement_name + suffix
        else:
            # Already correct, return the original match
            return match.group(0)

    # Apply the substitution using the dynamically compiled pattern
    code = name_pattern.sub(replace_with_strategy_string, code)

    # Report results of name masking
    if masking_applied_count > 0:
         print(f"  Fixer: Replaced {masking_applied_count} name assignment strings with '\"strategy\"'.")
    elif regex_matches_found > 0:
         # Matched the pattern but didn't need to change anything
         print(f"  Fixer: Found {regex_matches_found} lines like '{obfuscated_name_var} = ...', but the value was already 'strategy'.")
    else:
         # Didn't even find the pattern
         print(f"  Fixer: No lines matching '{obfuscated_name_var} = \"...\"' were found by the regex.")


    # --- 3. Overwrite the original obfuscated file *only* if changes were made ---
    if code != original_code:
        try:
            with open(obfuscated_file_path, 'w', encoding='utf-8') as f:
                f.write(code)
            print(f"  Fixer: Applied changes (attributes and/or name strings) to '{os.path.basename(obfuscated_file_path)}'")
        except Exception as e:
            print(f"Error: Fixer script failed writing changes to {obfuscated_file_path}: {e}", file=sys.stderr)
            sys.exit(1)
    else:
         print(f"  Fixer: No changes made by the fixer for '{os.path.basename(obfuscated_file_path)}'")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Post-process Carbon obfuscated files. Restores specific attributes based on --rules and replaces name='...' strings with name='strategy'.", # Updated description
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    parser.add_argument("-f", "--obfuscated-file", required=True, help="Path to the obfuscated .py file to fix (will be modified in-place).")
    parser.add_argument("-p", "--pairs-file", required=True, help="Path to the .pairs.json file containing name mappings from carbon.py.")
    parser.add_argument("-r", "--rules", default="", help="Optional rules string defining attribute exceptions, e.g., 'Action:C,D AnotherObject:Attr1'")
    args = parser.parse_args()
    rules_dictionary = parse_rules(args.rules)
    if not rules_dictionary and args.rules:
         print("Error: Rule parsing resulted in empty rules dictionary. Please check format.", file=sys.stderr)
         sys.exit(1)
    apply_fixes(args.obfuscated_file, args.pairs_file, rules_dictionary)