from typing import Dict, List, Optional, Any
import xml.etree.ElementTree as ET
from io import StringIO
import os
from pathlib import Path
from ..decorators import MoldoFunction

class MoldoParser:
    def __init__(self):
        self.imported_functions: Dict[str, callable] = {}
        self.imported_modules: Dict[str, Any] = {}
        self.module_paths: Dict[str, str] = {}

    def parse(self, code: str) -> str:
        """
        Parse Moldo code and generate equivalent Python code.
        
        Args:
            code: The Moldo code to parse
            
        Returns:
            Generated Python code
        """
        # Wrap the code in a root element for parsing
        wrapped_code = f"<moldo>{code}</moldo>"
        
        try:
            root = ET.fromstring(wrapped_code)
        except ET.ParseError as e:
            raise ValueError(f"Invalid Moldo code: {e}")
        
        python_code = []
        
        # First pass: collect imports
        for block in root:
            if block.tag == "mblock" and block.attrib.get("type") == "import":
                content = block.text.strip() if block.text else ""
                self.import_python_module(content)
        
        # Add module imports at the start
        header = ["# Generated by Moldo"]
        for module_name, path in self.module_paths.items():
            header.append(f"import sys")
            header.append(f"sys.path.append(os.path.dirname({path!r}))")
        header.append("")
        for module_name in self.imported_modules:
            header.append(f"import {module_name}")
        
        # Second pass: process other blocks
        for block in root:
            if block.tag == "mblock":
                block_type = block.attrib.get("type", "")
                content = block.text.strip() if block.text else ""
                
                if block_type == "print":
                    if not content.startswith('"') and not content.startswith("'") and not content.startswith("f"):
                        content = repr(content)
                    python_code.append(f"print({content})")
                elif block_type == "input":
                    python_code.append(f"print({content!r}, end='')")
                elif block_type == "variable":
                    python_code.append(content)
                elif block_type == "call":
                    # Map function calls to their module functions
                    for name, func in self.imported_functions.items():
                        if content.startswith(f"{name}("):
                            content = f"{func.__module__}.{func.__name__}{content[len(name):]}"
                            break
                    python_code.append(f"print({content})")
                elif block_type != "import":  # Skip import blocks in second pass
                    python_code.append(f"# Unknown block type: {block_type}")
                    python_code.append(f"# Content: {content}")
        
        return "\n".join(["import os"] + header + python_code)

    def import_python_module(self, module_path: str) -> None:
        """
        Import a Python module and register its Moldo-decorated functions.
        
        Args:
            module_path: Path to the Python module to import
        """
        import importlib.util
        import sys
        
        abs_path = os.path.abspath(module_path)
        module_name = Path(module_path).stem
        
        spec = importlib.util.spec_from_file_location(module_name, abs_path)
        if spec is None or spec.loader is None:
            raise ImportError(f"Could not import module from {abs_path}")
            
        module = importlib.util.module_from_spec(spec)
        sys.modules[module_name] = module
        spec.loader.exec_module(module)
        
        # Store the module and its path
        self.imported_modules[module_name] = module
        self.module_paths[module_name] = abs_path
        
        # Find all MoldoFunction instances in the module
        for name, obj in module.__dict__.items():
            if isinstance(obj, MoldoFunction):
                self.imported_functions[obj.reference_name] = obj

    def get_available_functions(self) -> List[str]:
        """
        Get a list of all available Moldo functions.
        
        Returns:
            List of function reference names
        """
        return list(self.imported_functions.keys())

    def execute_function(self, name: str, *args, **kwargs) -> Any:
        """
        Execute a Moldo function by its reference name.
        
        Args:
            name: The reference name of the function
            *args: Positional arguments to pass to the function
            **kwargs: Keyword arguments to pass to the function
            
        Returns:
            The result of the function execution
        """
        if name not in self.imported_functions:
            raise ValueError(f"Function {name} not found")
        return self.imported_functions[name](*args, **kwargs)
