import os
import re

def sanitize_content(content):
    # 1. Replace paper title variations
    title_regex = re.compile(
        r"XpertGPT:\s*Mixture\s*of\s*Experts\s*with\s*Parallelized\s*Multi-Scale\s*Information\s*Transmission(\s*for\s*Data-Constrained\s*Pretraining)?",
        re.IGNORECASE
    )
    content = title_regex.sub("AnonymousModel: Mixture of Experts with Parallelized Multi-Scale Information Transmission", content)
    
    # 2. Replace specific author names
    authors = [
        (re.compile(r"Soham\s*Jain", re.IGNORECASE), "Anonymous Author"),
        (re.compile(r"Harsh\s*Singh", re.IGNORECASE), "Anonymous Author"),
        (re.compile(r"Divija\s*Dewan", re.IGNORECASE), "Anonymous Author"),
        (re.compile(r"Atul\s*Dev", re.IGNORECASE), "Anonymous Author"),
        (re.compile(r"\bSoham\b", re.IGNORECASE), "Anonymous"),
        (re.compile(r"\bHarsh\b", re.IGNORECASE), "Anonymous"),
        (re.compile(r"\bDivija\b", re.IGNORECASE), "Anonymous"),
        (re.compile(r"\bAtul\b", re.IGNORECASE), "Anonymous"),
        (re.compile(r"\bJain\b", re.IGNORECASE), "Anonymous"),
        (re.compile(r"\bDewan\b", re.IGNORECASE), "Anonymous"),
    ]
    
    # We should be careful about replacing "Singh" and "Dev" as they might appear in code/comments
    # but we can replace them if they are in context of citation/names.
    # To be safe, let's replace "Singh" and "Dev" only when capitalized as part of names.
    authors.extend([
        (re.compile(r"\bSingh\b"), "Anonymous"),
        (re.compile(r"\bDev\b"), "Anonymous"),
    ])
    
    for pattern, replacement in authors:
        content = pattern.sub(replacement, content)
        
    # 3. Anonymize Hugging Face username, repo URLs, and BibTeX citation keys
    content = re.sub(r"SRJ5035", "anonymous_user", content)
    content = re.sub(r"swi_glu_sw_64_16_8_4_xpert_gpt", "anonymous_model", content)
    content = re.sub(r"jain2026xpertgpt", "anonymous2026model", content, flags=re.IGNORECASE)
    
    return content

def main():
    target_dir = os.path.dirname(os.path.abspath(__file__))
    print(f"[Sanitizer] Scanning folder: {target_dir}")
    
    files_to_sanitize = []
    for root, dirs, files in os.walk(target_dir):
        # Skip git and cache directories
        if '.git' in dirs:
            dirs.remove('.git')
        if '__pycache__' in dirs:
            dirs.remove('__pycache__')
            
        for f in files:
            if f in ['sanitize_codebase.py', 'upload_project_hf.py']:
                continue
            # Sanitize text-based files
            if f.endswith(('.py', '.md', '.txt', '.cff', '.json', '.sh')):
                files_to_sanitize.append(os.path.join(root, f))
                
    for filepath in files_to_sanitize:
        print(f"[Sanitizer] Sanitizing: {os.path.relpath(filepath, target_dir)}")
        try:
            with open(filepath, 'r', encoding='utf-8', errors='ignore') as f:
                content = f.read()
                
            sanitized = sanitize_content(content)
            
            with open(filepath, 'w', encoding='utf-8') as f:
                f.write(sanitized)
        except Exception as e:
            print(f"[Sanitizer] Error processing {filepath}: {e}")
            
    print("[Sanitizer] Sanitization complete!")

if __name__ == "__main__":
    main()