import re
import sys

files = ['ai_client.py', 'aggregate.py', 'mcp_client.py', 'shell_runner.py']
for file_path in files:
 print(f"Checking {file_path}...")
 with open(file_path, 'r', encoding='utf-8') as f:
  content = f.read()
  # Find all function definitions
  # This regex is simplified and might miss some edge cases (like multi-line defs)
  # But it's better than nothing.
  defs = re.finditer(r'def\s+([a-zA-Z_][a-zA-Z0-9_]*)\s*\((.*?)\)\s*(->\s*.*?)?:', content, re.DOTALL)
  for m in defs:
   name = m.group(1)
   args = m.group(2).strip()
   ret = m.group(3)
   if not ret:
    print(f"  Missing return type: {name}({args})")
    # Check arguments
   if args:
    arg_list = [a.strip() for a in args.split(',')]
    for arg in arg_list:
     if not arg or arg == 'self' or arg == 'cls':
      continue
     if ':' not in arg and '=' not in arg:
      print(f"  Missing arg type: {name} -> {arg}")
     elif ':' not in arg and '=' in arg:
     # arg=val (missing type)
      print(f"  Missing arg type: {name} -> {arg}")