import ast
import hashlib
from typing import Dict, Set


class WatermarkStripper(ast.NodeTransformer):
    """
    AST-based watermark / fingerprint stripper + light name obfuscator.
    - Renames user-defined functions, classes, and arguments
    - Removes leading docstrings from functions and classes
    - Preserves builtins, dunders, and imports
    """

    def __init__(self, prefix: str = "_w"):
        self.prefix = prefix
        self.name_map: Dict[str, str] = {}
        self.counter = 0
        self.protected: Set[str] = {
            "print", "len", "range", "list", "dict", "set", "tuple",
            "str", "int", "float", "bool", "None", "True", "False",
            "self", "cls", "super", "object", "type", "Exception",
            "__init__", "__name__", "__main__", "__file__", "__doc__",
            "__all__", "__dict__", "__class__", "__module__",
        }

    def _get_obfuscated_name(self, original: str) -> str:
        if original in self.protected or original.startswith("__"):
            return original
        if original not in self.name_map:
            digest = hashlib.md5(original.encode()).hexdigest()[:6]
            self.name_map[original] = f"{self.prefix}{self.counter}_{digest}"
            self.counter += 1
        return self.name_map[original]

    def visit_FunctionDef(self, node: ast.FunctionDef):
        node.name = self._get_obfuscated_name(node.name)
        if (node.body and isinstance(node.body[0], ast.Expr)
                and isinstance(node.body[0].value, ast.Constant)
                and isinstance(node.body[0].value.value, str)):
            node.body.pop(0)
        return self.generic_visit(node)

    def visit_ClassDef(self, node: ast.ClassDef):
        node.name = self._get_obfuscated_name(node.name)
        if (node.body and isinstance(node.body[0], ast.Expr)
                and isinstance(node.body[0].value, ast.Constant)
                and isinstance(node.body[0].value.value, str)):
            node.body.pop(0)
        return self.generic_visit(node)

    def visit_arg(self, node: ast.arg):
        node.arg = self._get_obfuscated_name(node.arg)
        return self.generic_visit(node)

    def visit_Name(self, node: ast.Name):
        if isinstance(node.ctx, (ast.Load, ast.Store, ast.Del)):
            if node.id in self.name_map:
                node.id = self.name_map[node.id]
        return node

    def visit_Attribute(self, node: ast.Attribute):
        self.generic_visit(node)
        return node


def sanitize_code(source_code: str) -> str:
    tree = ast.parse(source_code)
    stripper = WatermarkStripper()
    transformed = stripper.visit(tree)
    ast.fix_missing_locations(transformed)
    return ast.unparse(transformed)


if __name__ == "__main__":
    sample = '''
def calculate_sum(a, b):
    """This is a watermark docstring."""
    return a + b

class Calculator:
    """Class-level watermark."""
    def multiply(self, x, y):
        return x * y

result = calculate_sum(3, 4)
print(result)
'''

    print("=== ORIGINAL ===")
    print(sample)
    print("\n=== STRIPPED / OBFUSCATED ===")
    print(sanitize_code(sample))