import ast
import os
import json
import sys
import pathlib

class SDMMapper:
 def __init__(self):
  self.files = {} # path -> {"functions": {}, "classes": {}}
  self.functions_global = {} # name -> {"file": str, "class": str, "callers": set()}
  self.current_file = ""
  self.current_class = None
  self.current_function = None
  self.project_root = pathlib.Path.cwd().resolve()

 def get_rel_path(self, path):
  p = pathlib.Path(path).resolve()
  try:
   return str(p.relative_to(self.project_root)).replace("\\", "/")
  except (ValueError, RuntimeError):
   return str(p).replace("\\", "/")

 def collect_symbols(self, dirs):
  for d in dirs:
   if not os.path.exists(d): continue
   for root, _, files in os.walk(d):
    for f in files:
     if f.endswith(".py"):
      path = os.path.join(root, f)
      rel_path = self.get_rel_path(path)
      try:
       with open(path, "r", encoding="utf-8-sig") as file:
        tree = ast.parse(file.read(), filename=path)
        if rel_path not in self.files:
         self.files[rel_path] = {"functions": {}, "classes": {}}
        
        for node in ast.walk(tree):
         if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
          full_name = node.name
          # In first pass, we just note definitions. 
          # Class-member identification happens in visit_ClassDef.
          pass
         elif isinstance(node, ast.ClassDef):
          if node.name not in self.files[rel_path]["classes"]:
           self.files[rel_path]["classes"][node.name] = {"methods": {}, "variables": {}}
      except Exception as e:
       print(f"Error collecting symbols from {path}: {e}", file=sys.stderr)

 def analyze_files(self, dirs):
  for d in dirs:
   if not os.path.exists(d): continue
   for root, _, files in os.walk(d):
    for f in files:
     if f.endswith(".py"):
      self.analyze_file(os.path.join(root, f))

 def analyze_file(self, path):
  self.current_file = self.get_rel_path(path)
  if self.current_file not in self.files:
   self.files[self.current_file] = {"functions": {}, "classes": {}}
  
  try:
   with open(path, "r", encoding="utf-8-sig") as file:
    tree = ast.parse(file.read(), filename=path)
    visitor = SDMVisitor(self)
    visitor.visit(tree)
  except Exception as e:
   print(f"Error analyzing {path}: {e}", file=sys.stderr)

class SDMVisitor(ast.NodeVisitor):
 def __init__(self, mapper):
  self.mapper = mapper
  self.current_class = None
  self.current_function = None

 def visit_ClassDef(self, node):
  old_class = self.current_class
  self.current_class = node.name
  if self.current_class not in self.mapper.files[self.mapper.current_file]["classes"]:
   self.mapper.files[self.mapper.current_file]["classes"][self.current_class] = {"methods": {}, "variables": {}}
  self.generic_visit(node)
  self.current_class = old_class

 def visit_FunctionDef(self, node):
  old_func = self.current_function
  self.current_function = node.name
  
  full_name = f"{self.current_class}.{node.name}" if self.current_class else node.name
  if full_name not in self.mapper.functions_global:
   self.mapper.functions_global[full_name] = {
    "file": self.mapper.current_file,
    "class": self.current_class,
    "callers": set()
   }
  
  self.generic_visit(node)
  self.current_function = old_func

 def visit_AsyncFunctionDef(self, node):
  self.visit_FunctionDef(node)

 def visit_Call(self, node):
  name = None
  if isinstance(node.func, ast.Name):
   name = node.func.id
  elif isinstance(node.func, ast.Attribute):
   name = node.func.attr
  
  if name:
   # Try to find if it's a known function/method
   potential_matches = [n for n in self.mapper.functions_global if n == name or n.endswith("." + name)]
   for match in potential_matches:
    match_file = self.mapper.functions_global[match]["file"]
    # EXTERNAL FILTER: Only add caller if it's from a different file
    if match_file != self.mapper.current_file:
     caller_name = f"{self.current_class}.{self.current_function}" if self.current_class else (self.current_function or "module")
     # Include file name for external clarity
     self.mapper.functions_global[match]["callers"].add(f"{self.mapper.current_file}:{caller_name}")
  
  self.generic_visit(node)

 def visit_Attribute(self, node):
  if isinstance(node.value, ast.Name) and node.value.id == "self" and self.current_class:
   attr_name = node.attr
   class_data = self.mapper.files[self.mapper.current_file]["classes"][self.current_class]
   if attr_name not in class_data["variables"]:
    class_data["variables"][attr_name] = {"mutations": [], "usages": set()}
   
   if isinstance(node.ctx, ast.Store):
    class_data["variables"][attr_name]["mutations"].append({
     "file": self.mapper.current_file,
     "line": node.lineno,
     "method": self.current_function
    })
   elif isinstance(node.ctx, ast.Load):
    class_data["variables"][attr_name]["usages"].add(self.mapper.current_file)
  self.generic_visit(node)

def main():
 target = "."
 if len(sys.argv) > 1:
  target = sys.argv[1]
 
 mapper = SDMMapper()
 dirs = ["src", "simulation", "tests"]
 
 if os.path.isfile(target):
  mapper.collect_symbols(dirs)
  mapper.analyze_file(target)
 else:
  search_dirs = [target] if target in dirs else dirs
  mapper.collect_symbols(search_dirs)
  mapper.analyze_files(search_dirs)

 # Build the final grouped report
 report = {}
 
 # 1. Add functions/methods
 for full_name, data in mapper.functions_global.items():
  f_path = data["file"]
  if f_path not in report: report[f_path] = {"functions": {}, "classes": {}}
  
  # External callers only
  callers = sorted(list(data["callers"]))
  if not callers:
   continue

  tag = f"[C: {', '.join(callers)}]"
  if data["class"]:
   c_name = data["class"]
   if c_name not in report[f_path]["classes"]:
    report[f_path]["classes"][c_name] = {"methods": {}, "variables": {}}
   m_name = full_name.split(".")[-1]
   report[f_path]["classes"][c_name]["methods"][m_name] = tag
  else:
   report[f_path]["functions"][full_name] = tag

 # 2. Add class variables
 for f_path, f_data in mapper.files.items():
  if f_path not in report: continue
  for c_name, c_data in f_data["classes"].items():
   if c_name not in report[f_path]["classes"]: 
    report[f_path]["classes"][c_name] = {"methods": {}, "variables": {}}
   
   class_vars_summary = []
   for v_name, v_data in c_data["variables"].items():
    # EXTERNAL FILTER: Only include mutations/usages from different files
    ext_muts = [f"{m['file']}:{m['line']}, {m['method']}" for m in v_data["mutations"] if m['file'] != f_path]
    ext_usages = [u for u in v_data["usages"] if u != f_path]
    
    if not ext_muts and not ext_usages:
     continue

    m_tag = f"[M: {'; '.join(ext_muts or ['None'])}]"
    u_tag = f"[U: {', '.join(sorted(list(ext_usages or ['None'])))}]"
    tag = f"{m_tag} {u_tag}"
    report[f_path]["classes"][c_name]["variables"][v_name] = tag
    class_vars_summary.append(f"{v_name}: {tag}")
   
   if class_vars_summary:
    report[f_path]["classes"][c_name]["class_tag"] = "\n".join(class_vars_summary)

 print(json.dumps(report, indent=1))

if __name__ == "__main__":
 main()