import os
import re
from pathlib import Path

def generate_architectural_graph(directory, output_file):
    # Setup node definitions and edges
    node_types = {} # entity_name -> type (container/context/task)
    edges = []
    
    # Regex to extract Gen blocks
    block_re = re.compile(r"```Gen\s*(.*?)\s*```", re.DOTALL)
    
    for md_file in Path(directory).rglob("*.md"):
        content = md_file.read_text()
        blocks = block_re.findall(content)
        
        for block in blocks:
            lines = [l.strip() for l in block.split('\n') if l.strip()]
            props = {}
            current_entity = None
            current_type = None

            # Parse lines for entity definition and properties
            for line in lines:
                if ':' in line:
                    key, val = line.split(':', 1)
                    key, val = key.strip(), val.strip()
                    props[key] = val
                    
                    if key in ['comprehension', 'modulator', 'container', 'context', 'task']:
                        current_entity = val
                        current_type = key
                        node_types[current_entity] = current_type

            if not current_entity: continue

            # Morphism Logic: Contained By (Parent <- Child)
            if 'contained_by' in props:
                parent = props['contained_by']
                edges.append((parent, current_entity, "contains"))

            # Morphism Logic: Expands (Parent -> Child)
            if 'expands' in props:
                contains_val = props['expands']
                match = re.match(r"(\d+x)?\s*(.*)", contains_val)
                label = match.group(1) if match.group(1) else "expands"
                child = match.group(2)
                edges.append((current_entity, child, label))

    # Build Mermaid syntax
    mermaid_lines = ["graph TD"]
    
    # Define Nodes with Types (Styling)
    # Containers = Rectangles [], Contexts = Hexagons {{}}, Tasks = Ovals (())
    for name, n_type in node_types.items():
        if n_type == 'comprehension':
            mermaid_lines.append(f'    {name}["comprehension: {name}"]')
        elif n_type == 'modulator':
            mermaid_lines.append(f'    {name}["modulator: {name}"]')
        elif n_type == 'container':
            mermaid_lines.append(f'    {name}["container: {name}"]')
        elif n_type == 'context':
            mermaid_lines.append(f'    {name}{{{{"context: {name}"}}}}')
        elif n_type == 'task':
            mermaid_lines.append(f'    {name}(("task: {name}"))')


    # Add Edges
    for source, target, label in edges:
        # Determine relationship type string for the arrow
        s_type = node_types.get(source, "unknown")
        t_type = node_types.get(target, "unknown")
        # Example: container tre -- contains --> task bbgg
        mermaid_lines.append(f'    {source} -- "{label}" --> {target}')

    # Save to file
    with open(output_file, "w") as f:
        f.write("# Architectural Hierarchy Graph\n\n")
        f.write("```mermaid\n")
        f.write("\n".join(mermaid_lines))
        f.write("\n```")

if __name__ == "__main__":
    generate_architectural_graph("..", "architecture_graph.mmd")
    print("Success: architecture_graph.md has been generated.")