import tempfile
import os
import sys
import subprocess
import numpy as np
from setuptools import setup
from typing import Any, Optional, Dict
import json


class CythonCompilationError(Exception):
    """Cython compilation error"""
    pass


class CythonRuntimeError(Exception):
    """Compiled code execution error"""
    pass


class CythonImportError(Exception):
    """Cython import error"""
    pass




def cython(pyx_content: str, auto_convert: bool = True, function_name: Optional[str] = None, 
           args: tuple = (), kwargs: Optional[Dict[str, Any]] = None):
    """
    Compiles and executes Cython code.
    
    Args:
        pyx_content: Cython code to compile
        auto_convert: Automatically convert result via eval
        function_name: Name of function to call (default 'main')
        args: Positional arguments for the function
        kwargs: Keyword arguments for the function
        
    Returns:
        Code execution result
        
    Raises:
        CythonImportError: If Cython is not installed
        CythonCompilationError: If compilation fails
        CythonRuntimeError: If code execution fails
    """
    try:
        import Cython
    except ImportError:
        raise CythonImportError(
            "Cython is not installed. Install it: pip install Cython"
        )
    
    if kwargs is None:
        kwargs = {}
    
    function_name = function_name or 'main'
    
    with tempfile.TemporaryDirectory() as temp_dir:
        pyx_path = os.path.join(temp_dir, "module.pyx")
        main_path = os.path.join(temp_dir, "main.py")
        setup_path = os.path.join(temp_dir, "setup.py")
        
        # Write .pyx file
        try:
            with open(pyx_path, 'w', encoding='utf-8') as f:
                f.write(pyx_content)
        except Exception as e:
            raise CythonCompilationError(
                f"Failed to write .pyx file: {e}"
            )
        
        # Create main.py with argument handling
        enhanced_main = f'''
import sys
import os
import json
import traceback

sys.path.insert(0, os.getcwd())

try:
    import module
except Exception as e:
    print(json.dumps({{
        "error": "import",
        "message": str(e),
        "traceback": traceback.format_exc()
    }}))
    sys.exit(1)

def capture_output():
    try:
        # Load arguments
        args = json.loads("""{json.dumps(args)}""")
        kwargs = json.loads("""{json.dumps(kwargs)}""")
        
        result = None
        
        if hasattr(module, '{function_name}'):
            # Call specified function with arguments
            func = getattr(module, '{function_name}')
            result = func(*args, **kwargs)
        elif hasattr(module, 'main'):
            # Fallback to main without arguments (for backward compatibility)
            result = module.main()
        else:
            # Execute module code and capture output
            import io
            from contextlib import redirect_stdout
            f = io.StringIO()
            with redirect_stdout(f):
                import importlib
                importlib.reload(module)
            result = f.getvalue()
        
        return result
    except Exception as e:
        print(json.dumps({{
            "error": "runtime",
            "message": str(e),
            "traceback": traceback.format_exc()
        }}))
        sys.exit(1)

if __name__ == "__main__":
    output = capture_output()
    if output is not None:
        print(json.dumps({{"success": True, "result": str(output)}}))
    else:
        print(json.dumps({{"success": True, "result": None}}))
'''
        
        try:
            with open(main_path, 'w', encoding='utf-8') as f:
                f.write(enhanced_main)
        except Exception as e:
            raise CythonCompilationError(
                f"Failed to create main.py: {e}"
            )
        
        # Create setup.py
        setup_content = '''
from setuptools import setup
from Cython.Build import cythonize
import sys

try:
    setup(ext_modules=cythonize("module.pyx", compiler_directives={'language_level': '3'}))
except Exception as e:
    print(f"Setup error: {e}", file=sys.stderr)
    sys.exit(1)
'''
        try:
            with open(setup_path, 'w', encoding='utf-8') as f:
                f.write(setup_content)
        except Exception as e:
            raise CythonCompilationError(
                f"Failed to create setup.py: {e}"
            )
        
        # Compile
        try:
            build_result = subprocess.run(
                [sys.executable, "setup.py", "build_ext", "--inplace"],
                cwd=temp_dir,
                capture_output=True,
                text=True,
                timeout=30
            )
            
            if build_result.returncode != 0:
                error_msg = build_result.stderr or build_result.stdout
                raise CythonCompilationError(
                    f"Cython compilation error:\n{error_msg}"
                )
        except subprocess.TimeoutExpired:
            raise CythonCompilationError(
                "Compilation timeout (30 seconds)"
            )
        except CythonCompilationError:
            raise
        except Exception as e:
            raise CythonCompilationError(
                f"Unexpected compilation error: {e}"
            )
        
        # Execute
        try:
            result = subprocess.run(
                [sys.executable, "main.py"],
                cwd=temp_dir,
                capture_output=True,
                text=True,
                timeout=30
            )
            
            if result.returncode != 0:
                try:
                    error_data = json.loads(result.stdout)
                    if error_data.get("error") == "import":
                        raise CythonRuntimeError(
                            f"Module import error:\n{error_data['message']}\n\n"
                            f"Traceback:\n{error_data['traceback']}"
                        )
                    elif error_data.get("error") == "runtime":
                        raise CythonRuntimeError(
                            f"Execution error:\n{error_data['message']}\n\n"
                            f"Traceback:\n{error_data['traceback']}"
                        )
                except json.JSONDecodeError:
                    raise CythonRuntimeError(
                        f"Execution error:\n{result.stdout}\n{result.stderr}"
                    )
            
            # Parse result
            try:
                output_data = json.loads(result.stdout)
                if output_data.get("success"):
                    res = output_data.get("result")
                    if res is None:
                        return None
                    
                    # Try to convert result
                    if auto_convert and res:
                        try:
                            return eval(res)
                        except (SyntaxError, NameError):
                            return res
                    return res
            except json.JSONDecodeError:
                # Fallback to old output format
                res = result.stdout.strip()
                if auto_convert and res:
                    try:
                        return eval(res)
                    except (SyntaxError, NameError):
                        return res
                return res
                
        except subprocess.TimeoutExpired:
            raise CythonRuntimeError(
                "Execution timeout (30 seconds)"
            )
        except CythonRuntimeError:
            raise
        except Exception as e:
            raise CythonRuntimeError(
                f"Unexpected execution error: {e}"
            )

def runfile(path: str, function_name: Optional[str] = None, 
            args: tuple = (), kwargs: Optional[Dict[str, Any]] = None):
    """
    Compiles and executes Cython code from file.
    
    Args:
        path: Path to .pyx file
        function_name: Name of function to call
        args: Positional arguments
        kwargs: Keyword arguments
        
    Returns:
        Execution result
        
    Raises:
        FileNotFoundError: If file not found
        CythonCompilationError: If compilation fails
        CythonRuntimeError: If execution fails
    """
    try:
        with open(path, encoding='utf-8') as file:
            content = file.read()
    except FileNotFoundError:
        raise FileNotFoundError(f"File not found: {path}")
    except Exception as e:
        raise CythonCompilationError(f"Failed to read file {path}: {e}")
    
    return cython(content, function_name=function_name, args=args, kwargs=kwargs)

def setupfile(path: str = 'all'):
    path = '*' if path == all else path
    from Cython.Build import cythonize

    setup(
        ext_modules=cythonize(f"{path}.pyx"),
        include_dirs=[np.get_include()]
    )