# pioton/preprocessor.py
import re
import sys
from pathlib import Path

class PimportTransformer:
    """Transforms pimport statements to standard Python import calls."""
    
    def __init__(self):
        # Patterns for different pimport forms
        self.patterns = [
            # pimport numpy as np
            (re.compile(r'pimport\s+([a-zA-Z_][a-zA-Z0-9_]*)(?:\s+as\s+([a-zA-Z_][a-zA-Z0-9_]*))?'),
             lambda m: self._transform_simple(m)),
            
            # from .module pimport X as x, Y as y, Z
            (re.compile(r'from\s+(\..*?)\s+pimport\s+([^#]*)'),
             lambda m: self._transform_from(m)),
        ]
    
    def _transform_simple(self, match):
        """pimport numpy as np -> _pimport('numpy', 'np')"""
        module = match.group(1)
        alias = match.group(2) or module
        return f"{alias} = _pimport('{module}')"
    
    def _transform_from(self, match):
        """from .module pimport X, Y as y -> _pimport_from('.module', ['X', 'Y as y'])"""
        module_path = match.group(1)
        names = [name.strip() for name in match.group(2).split(',')]
        names_str = ', '.join(f"'{name}'" for name in names)
        return f"_pimport_from('{module_path}', [{names_str}])"
    
    def transform_code(self, code: str) -> str:
        """Transform all pimport statements in code."""
        lines = code.split('\n')
        transformed = []
        
        for line in lines:
            original = line
            for pattern, handler in self.patterns:
                match = pattern.match(line.strip())
                if match:
                    line = handler(match)
                    # Add comment showing original
                    line += f"  # Original: {original.strip()}"
                    break
            transformed.append(line)
        
        return '\n'.join(transformed)