"""
Python wrapper generator.

Generates Python wrapper code with support for:
- Primitive types (string, int, float, bool, unit)
- Option types (Optional[T])
- List types (List[T])
- Tuple types (Tuple[T1, T2, ...])
- Custom types (records, variants)
"""

from typing import Dict

from polyglot_ffi.ir.types import IRModule, IRType, IRFunction, TypeKind
from polyglot_ffi.utils.naming import sanitize_module_name


class PythonGenerator:
    """Generate Python wrapper code."""

    # Map IR types to Python type hints
    PY_TYPE_MAP: Dict[str, str] = {
        "string": "str",
        "int": "int",
        "float": "float",
        "bool": "bool",
        "unit": "None",
    }

    # Map IR types to ctypes
    CTYPES_MAP: Dict[str, str] = {
        "string": "ctypes.c_char_p",
        "int": "ctypes.c_int",
        "float": "ctypes.c_double",
        "bool": "ctypes.c_bool",
        "unit": "None",
    }

    @staticmethod
    def _sanitize_identifier(name: str) -> str:
        """Sanitize a name to be a valid Python identifier."""
        # Replace hyphens and other non-identifier chars with underscores
        sanitized = name.replace("-", "_").replace(".", "_")
        # Ensure it starts with a letter or underscore
        if sanitized and not (sanitized[0].isalpha() or sanitized[0] == "_"):
            sanitized = "_" + sanitized
        return sanitized

    def generate(self, module: IRModule, module_name: str) -> str:
        """Generate Python wrapper module."""
        # Sanitize module name for Python identifiers and filenames
        safe_name = sanitize_module_name(module_name)
        error_class = safe_name.capitalize() + "Error"

        lines = [
            "# Generated by polyglot-ffi",
            f"# {safe_name}_py.py",
            "",
            "import ctypes",
            "from pathlib import Path",
            "from typing import Optional, List, Tuple, Any",
            "",
            "# Load the shared library",
            f'_lib_path = Path(__file__).parent / "lib{safe_name}.so"',
            "_lib = ctypes.CDLL(str(_lib_path))",
            "",
            f"class {error_class}(Exception):",
            f'    """Raised when {module_name} operations fail"""',
            "    pass",
            "",
        ]

        # Configure ctypes for each function
        for func in module.functions:
            lines.append(f"# Configure {func.name}")

            # Set argtypes
            if func.params:
                argtypes = [self._get_ctypes(p.type) for p in func.params]
                lines.append(f"_lib.ml_{func.name}.argtypes = [{', '.join(argtypes)}]")

            # Set restype
            restype = self._get_ctypes(func.return_type)
            lines.append(f"_lib.ml_{func.name}.restype = {restype}")
            lines.append("")

        # Generate wrapper functions
        for func in module.functions:
            lines.extend(self._generate_function_wrapper(func, module_name, error_class))
            lines.append("")

        return "\n".join(lines)

    def _generate_function_wrapper(
        self, func: IRFunction, module_name: str, error_class: str
    ) -> list:
        """Generate Python wrapper for a single function."""
        lines = []

        # Function signature
        params = [f"{p.name}: {self._get_py_type(p.type)}" for p in func.params]
        params_str = ", ".join(params) if params else ""
        return_type = self._get_py_type(func.return_type)

        lines.append(f"def {func.name}({params_str}) -> {return_type}:")

        # Docstring
        if func.doc:
            lines.append(f'    """{func.doc}"""')
        else:
            lines.append(f'    """Call OCaml {func.name} function"""')

        lines.append("    try:")

        # Convert arguments
        call_args = []
        for param in func.params:
            if param.type.name == "string":
                call_args.append(f"{param.name}.encode('utf-8')")
            else:
                call_args.append(param.name)

        # Make the call
        call_str = ", ".join(call_args) if call_args else ""
        lines.append(f"        result = _lib.ml_{func.name}({call_str})")

        # Handle return value
        if func.return_type.name == "string":
            lines.append("        if result is None:")
            lines.append(f'            raise {error_class}("{func.name} returned NULL")')
            lines.append("        return result.decode('utf-8')")
        elif func.return_type.name == "unit":
            lines.append("        return None")
        else:
            lines.append("        return result")

        # Error handling
        lines.append("    except Exception as e:")
        lines.append(f'        raise {error_class}(f"{func.name} failed: {{e}}")')

        return lines

    def _get_py_type(self, ir_type: IRType) -> str:
        """
        Convert IR type to Python type hint.

        Handles:
        - Primitives: str, int, float, bool, None
        - Options: Optional[T]
        - Lists: List[T]
        - Tuples: Tuple[T1, T2, ...]
        - Custom types: Class names
        """
        if ir_type.is_primitive():
            return self.PY_TYPE_MAP.get(ir_type.name, "str")

        # Handle option types
        if ir_type.kind == TypeKind.OPTION:
            if ir_type.params:
                inner_type = self._get_py_type(ir_type.params[0])
                return f"Optional[{inner_type}]"
            return "Optional[Any]"

        # Handle list types
        if ir_type.kind == TypeKind.LIST:
            if ir_type.params:
                inner_type = self._get_py_type(ir_type.params[0])
                return f"List[{inner_type}]"
            return "List[Any]"

        # Handle tuple types
        if ir_type.kind == TypeKind.TUPLE:
            if ir_type.params:
                tuple_types = [self._get_py_type(p) for p in ir_type.params]
                return f"Tuple[{', '.join(tuple_types)}]"
            return "Tuple[Any, ...]"

        # Handle custom types (records, variants)
        if ir_type.kind in (TypeKind.CUSTOM, TypeKind.RECORD, TypeKind.VARIANT):
            # Use the type name as a class name (with capitalization)
            return ir_type.name.capitalize()

        raise ValueError(f"Unsupported type for Python generation: {ir_type}")

    def _get_ctypes(self, ir_type: IRType) -> str:
        """
        Convert IR type to ctypes type.

        For complex types (options, lists, tuples, custom types),
        we use c_void_p as they're opaque pointers from C's perspective.
        """
        if ir_type.is_primitive():
            return self.CTYPES_MAP.get(ir_type.name, "ctypes.c_char_p")

        # Complex types are passed as opaque pointers (v0.2.0+)
        if ir_type.kind in (
            TypeKind.OPTION,
            TypeKind.LIST,
            TypeKind.TUPLE,
            TypeKind.CUSTOM,
            TypeKind.RECORD,
            TypeKind.VARIANT,
        ):
            return "ctypes.c_void_p"

        raise ValueError(f"Unsupported type for ctypes: {ir_type}")
