ritz/framework.py

"""
Ritz Framework - Python Prototype
This runs in actual Python for debugging/prototyping!
"""
from contextlib import contextmanager
from typing import Generator, Any, List, Dict, Optional, Callable
import json

# Global context for building component trees
_context_stack = []

class Signal:
    """Reactive signal - simplified for Python prototype"""
    def __init__(self, initial_value):
        self.value = initial_value
        self.subscribers = []
    
    def __str__(self):
        return str(self.value)
    
    def __iadd__(self, other):
        self.value += other
        self._notify()
        return self
    
    def _notify(self):
        for subscriber in self.subscribers:
            subscriber()
    
    def subscribe(self, callback):
        self.subscribers.append(callback)

class VNode:
    """Virtual node for component tree"""
    def __init__(self, tag: str, props: Dict = None, children: List = None):
        self.tag = tag
        self.props = props or {}
        self.children = children or []
    
    def to_html(self) -> str:
        """Render to HTML string"""
        if self.children:
            children_html = ''.join(
                child.to_html() if isinstance(child, VNode) else str(child)
                for child in self.children
            )
            attrs = ' '.join(f'{k}="{v}"' for k, v in self.props.items() if k != 'onclick')
            attrs = ' ' + attrs if attrs else ''
            return f'<{self.tag}{attrs}>{children_html}</{self.tag}>'
        else:
            attrs = ' '.join(f'{k}="{v}"' for k, v in self.props.items() if k != 'onclick')
            attrs = ' ' + attrs if attrs else ''
            return f'<{self.tag}{attrs} />'
    
    def to_ir(self) -> Dict:
        """Convert to intermediate representation"""
        # Filter out non-serializable props like onclick functions
        serializable_props = {
            k: v for k, v in self.props.items() 
            if not callable(v)
        }
        
        return {
            'type': 'element',
            'tag': self.tag,
            'props': serializable_props,
            'children': [
                child.to_ir() if isinstance(child, VNode) 
                else {'type': 'text', 'value': str(child)}
                for child in self.children
            ]
        }

class Element:
    """Context manager for building elements"""
    def __init__(self, tag: str, *children, **props):
        self.tag = tag
        self.props = props
        self.children = list(children) if children else []
        self.vnode = VNode(tag, props, self.children)
    
    def __enter__(self):
        _context_stack.append(self)
        # Store the initial length so we know what was added in context
        self._initial_children_len = len(self.children)
        return self
    
    def __exit__(self, *args):
        _context_stack.pop()
        # Don't rebuild vnode here - let the component decorator handle it

def component(func: Callable) -> Callable:
    """Decorator for components that use yield"""
    def wrapper(*args, **kwargs):
        gen = func(*args, **kwargs)
        if hasattr(gen, '__next__'):  # It's a generator
            root_element = None
            initial_stack_len = len(_context_stack)
            
            # Track the first context that gets created
            first_context = None
            
            # Execute the generator to build the tree
            for item in gen:
                # Check if a new context was just created
                if not first_context and len(_context_stack) > initial_stack_len:
                    first_context = _context_stack[initial_stack_len]
                
                current_context = _context_stack[-1] if _context_stack else None
                
                if current_context:
                    # We're inside a with statement, add to its children
                    if isinstance(item, Element):
                        current_context.children.append(item.vnode)
                    elif isinstance(item, VNode):
                        current_context.children.append(item) 
                    elif isinstance(item, str):
                        current_context.children.append(item)
                    else:
                        # Try to convert to string
                        current_context.children.append(str(item))
                else:
                    # No context, this should be the root
                    if isinstance(item, Element):
                        root_element = item.vnode
                    elif isinstance(item, VNode):
                        root_element = item
            
            # Return the first context if it was created, otherwise root_element
            if first_context:
                first_context.vnode = VNode(
                    first_context.tag,
                    first_context.props,
                    first_context.children
                )
                return first_context.vnode
            
            return root_element
        else:
            return func(*args, **kwargs)
    
    return wrapper

# Factory functions for HTML elements
def _make_element(tag: str) -> Callable:
    def element_factory(*children, **props):
        if len(children) == 1 and isinstance(children[0], str):
            # Shorthand: P('text') creates element with text
            return Element(tag, children[0], **props)
        return Element(tag, *children, **props)
    return element_factory

# Generate standard HTML elements
for tag in ['main', 'div', 'p', 'button', 'label', 'h1', 'h2', 'h3', 
            'article', 'section', 'header', 'footer', 'span', 'strong',
            'em', 'li', 'ul', 'ol', 'form', 'input', 'nav', 'time']:
    globals()[tag.capitalize()] = _make_element(tag)

def signal(initial_value) -> Signal:
    """Create a reactive signal"""
    return Signal(initial_value)

# Helper function for debugging
def render(component_func, *args, **kwargs):
    """Render a component to HTML and IR"""
    result = component_func(*args, **kwargs)
    
    if isinstance(result, VNode):
        print("=== HTML Output ===")
        print(result.to_html())
        print("\n=== IR Output ===")
        print(json.dumps(result.to_ir(), indent=2))
    else:
        print("Component did not return a VNode")
    
    return result