#!/opt/homebrew/bin/python3

import re
import glob
import os
import graphviz

##################################################################################

######### USAGE #########
# 1. Place this file at the topmost folder that contains all of the SwiftUI 
#  files you wish to be included. The script will recuresively find files in
#  any subfolder.

# 2. Change this to the path to your main SwiftUI view file, relative to the
#  script. If the script is in the same folder as the file, just put the filename.
MAIN_FILE = 'AppView.swift'

# 3. Get need dot and graphviz. Easiest way is to run:
#  $ pip install graphviz

# 4. Ensure `dot` is in PATH. If you use homebrew, this should do it:
os.environ["PATH"] += os.pathsep +  "/opt/homebrew/bin"

# Ignore any `PreferenceKey` and `EnvironmentKey` types
DROP_KEYS = True

# Ignore anything that isn't a `View`
DROP_NON_VIEWS = True

# Style subgraphs differently
STYLE_SUBGRAPH = False

##################################################################################

class Node:
	def __init__(self, name, children, inline_children):
		self.name = name
		self.children = children
		self.inline_children = inline_children

found = []

def analyze(file):
	# print(name)
	found.append(file)
	base = os.path.basename(file)
	viewname = os.path.splitext(base)[0]
	
	children = []
	inline_structs = []
	
	with open(file) as f:
		lines = f.readlines()
		
		for rawline in lines:
			line = rawline.strip()
			if line.strip("private").strip("public").strip().startswith("struct"):
				regex_result = re.search("struct ([A-Za-z0-9_]*)", line)
				if regex_result:
					name = regex_result.group(1)
					if name != viewname and name not in inline_structs and "_" not in name:
						isView = "View" in line
						if not DROP_NON_VIEWS or isView:
							iskey = "PreferenceKey" in line or "EnvironmentKey" in line
							if not DROP_KEYS or not iskey:
								inline_structs.append(name)
    
		in_view = False
		
		for rawline in lines:
			line = rawline.strip()
			is_comment = line.startswith("/*") or line.startswith("*") or line.startswith("//") or line.startswith("import")
			is_property_decl = " var " in line or " let " in line
			if not is_comment and not is_property_decl:
				regex_result = re.search("(?!\.)([A-Z][a-zA-z0-9]*)(?!\.)\s*(\(|{)", line)
				if regex_result:
					name = regex_result.group(1)
				
					if name not in inline_structs and "_" not in name:
						for f in glob.glob('**/' + name + '.swift', recursive=True):
							if f != file:
								child = analyze(f)
								children.append(child)
	
	return Node(viewname, children, inline_structs)
						
			
def print_node(node, prefix):
	print(prefix + node.name)
	for inline in node.inline_children:
		print(prefix + " (" + inline + ")")
	for child in node.children:
		print_node(child, prefix + "  ")
	
# keep track of already created edges since we don't care about multiple edges between the same nodes
edgesIDs = [] 

def graph(dot, node):
	dot.node(node.name)
	for child in node.children:
		id = node.name + child.name
		if not id in edgesIDs:
			edge = dot.edge(node.name, child.name)
			edgesIDs.append(id)
			graph(dot, child)
	if len(node.inline_children) > 0:
		with dot.subgraph(name='cluster_' + node.name) as c:
			if STYLE_SUBGRAPH:
				c.node_attr['shape'] = 'egg'
				c.node_attr['fontcolor'] = 'darkgray'
				c.node_attr['style'] = ''
			c.graph_attr['style'] = 'dotted'
			
			for inline in node.inline_children:
				id = node.name + inline
				if id not in edgesIDs:
					edgesIDs.append(id)
					c.node(inline)
					c.edge(node.name, inline)

def run():
	main_node = analyze(MAIN_FILE)

	print_node(main_node, "")

	dot = graphviz.Digraph(comment='View Tree From: ' + main_node.name, node_attr={'shape': 'box', 'style': 'filled'})
	graph(dot, main_node)
	# print(dot.source)

	if STYLE_SUBGRAPH:
		source = dot
	else:
		source = dot.unflatten(stagger=3)
	source.render('./graph', format='png', view=True)
	
run()