import re
import sys

def fix_documentation_improved(file_path):
    with open(file_path, 'r') as f:
        content = f.read()
    
    lines = content.split('\n')
    result = []
    i = 0
    
    while i < len(lines):
        line = lines[i]
        
        # Check if this is an undocumented public struct/enum/trait/fn
        if re.match(r'^pub (struct|enum|trait|fn|type|const|static) \w+', line):
            if i == 0 or not lines[i-1].strip().startswith('///'):
                # Extract the name and kind for documentation
                match = re.search(r'pub (struct|enum|trait|fn|type|const|static) (\w+)', line)
                if match:
                    kind = match.group(1)
                    name = match.group(2)
                    # Add basic documentation
                    result.append(f'/// {name}')
                    if kind in ['struct', 'enum']:
                        result.append(f'/// ')
                        result.append(f'/// TODO: Add detailed documentation for this {kind}')
                    elif kind == 'fn':
                        result.append(f'/// ')
                        result.append(f'/// TODO: Add detailed documentation for this function')
                    elif kind == 'trait':
                        result.append(f'/// ')
                        result.append(f'/// TODO: Add detailed documentation for this trait')
        
        # Check if this is an undocumented public field (more robust pattern)
        elif re.match(r'^\s+pub \w+:', line):
            if i == 0 or not (lines[i-1].strip().startswith('///') or lines[i-1].strip().startswith('#[')):
                # Extract field name
                field_match = re.search(r'pub (\w+):', line)
                if field_match:
                    field_name = field_match.group(1)
                    indent = re.match(r'^(\s+)', line).group(1)
                    # Create more descriptive documentation
                    readable_name = field_name.replace('_', ' ').title()
                    result.append(f'{indent}/// {readable_name}')
        
        # Check for undocumented enum variants
        elif re.match(r'^\s+[A-Z]\w*[,{]?$', line) or re.match(r'^\s+[A-Z]\w*\([^)]*\)[,{]?$', line):
            if i == 0 or not lines[i-1].strip().startswith('///'):
                # Extract variant name
                variant_match = re.search(r'^\s+([A-Z]\w*)', line)
                if variant_match:
                    variant_name = variant_match.group(1)
                    indent = re.match(r'^(\s+)', line).group(1)
                    result.append(f'{indent}/// {variant_name} variant')
        
        result.append(line)
        i += 1
    
    return '\n'.join(result)

if __name__ == "__main__":
    if len(sys.argv) != 2:
        print("Usage: python fix_docs_improved.py <file_path>")
        sys.exit(1)
    
    file_path = sys.argv[1]
    fixed_content = fix_documentation_improved(file_path)
    
    with open(file_path, 'w') as f:
        f.write(fixed_content)
    
    print(f"Improved documentation in {file_path}")