"""
C stub code generator.

Generates C wrapper code with proper memory management for:
- Primitive types (string, int, float, bool, unit)
- Option types (as nullable pointers)
- List types (as opaque OCaml values)
- Tuple types (as opaque OCaml values)
- 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 CStubGenerator:
    """Generate C wrapper code for OCaml functions."""

    # Map IR types to C types
    C_TYPE_MAP: Dict[str, str] = {
        "string": "const char*",
        "int": "int",
        "float": "double",
        "bool": "int",
        "unit": "void",
    }

    def generate_stubs(self, module: IRModule, module_name: str) -> str:
        """Generate C stub implementation file."""
        safe_name = sanitize_module_name(module_name)
        lines = [
            f"/* Generated by polyglot-ffi */",
            f"/* {safe_name}_stubs.c */",
            "",
            "#include <string.h>",
            "#include <caml/mlvalues.h>",
            "#include <caml/memory.h>",
            "#include <caml/alloc.h>",
            "#include <caml/callback.h>",
            "",
        ]

        for func in module.functions:
            lines.extend(self._generate_function_stub(func))
            lines.append("")

        return "\n".join(lines)

    def generate_header(self, module: IRModule, module_name: str) -> str:
        """Generate C header file."""
        safe_name = sanitize_module_name(module_name)
        guard = f"{safe_name.upper()}_STUBS_H"

        lines = [
            f"/* Generated by polyglot-ffi */",
            f"/* {safe_name}_stubs.h */",
            "",
            f"#ifndef {guard}",
            f"#define {guard}",
            "",
        ]

        # Function declarations
        for func in module.functions:
            c_return = self._get_c_type(func.return_type)
            params = ", ".join(f"{self._get_c_type(p.type)} {p.name}" for p in func.params)
            lines.append(f"{c_return} ml_{func.name}({params});")

        lines.append("")
        lines.append(f"#endif /* {guard} */")

        return "\n".join(lines)

    def _generate_function_stub(self, func: IRFunction) -> list:
        """Generate C stub for a single function."""
        lines = []

        # Function signature
        c_return = self._get_c_type(func.return_type)
        params = ", ".join(f"{self._get_c_type(p.type)} {p.name}" for p in func.params)

        lines.append(f"/* Wrapper for OCaml {func.name} function */")
        lines.append(f"{c_return} ml_{func.name}({params}) {{")
        lines.append("    CAMLparam0();")

        # Declare local variables
        num_locals = len(func.params) + 1
        local_vars = ", ".join(f"ml_{p.name}" for p in func.params) + ", ml_result"
        lines.append(f"    CAMLlocal{num_locals}({local_vars});")
        lines.append("")

        # Convert C parameters to OCaml values
        for param in func.params:
            lines.extend(self._convert_c_to_ocaml(param.name, param.type))

        # Call OCaml function
        if len(func.params) == 1:
            # Single parameter
            lines.append(
                f'    ml_result = caml_callback(*caml_named_value("{func.name}"), ml_{func.params[0].name});'
            )
        elif len(func.params) > 1:
            # Multiple parameters
            param_names = ", ".join(f"ml_{p.name}" for p in func.params)
            lines.append(
                f'    ml_result = caml_callback{len(func.params)}(*caml_named_value("{func.name}"), {param_names});'
            )
        else:
            # No parameters
            lines.append(
                f'    ml_result = caml_callback(*caml_named_value("{func.name}"), Val_unit);'
            )

        lines.append("")

        # Convert result back to C
        lines.extend(self._convert_ocaml_to_c(func.return_type))

        lines.append("}")

        return lines

    def _convert_c_to_ocaml(self, param_name: str, param_type: IRType) -> list:
        """
        Generate code to convert C value to OCaml value.

        Handles primitives and complex types (options, lists, tuples, custom types).
        Complex types are passed as opaque pointers and cast to OCaml values.
        """
        lines = []

        if param_type.is_primitive():
            if param_type.name == "string":
                lines.append(f"    ml_{param_name} = caml_copy_string({param_name});")
            elif param_type.name == "int":
                lines.append(f"    ml_{param_name} = Val_int({param_name});")
            elif param_type.name == "float":
                lines.append(f"    ml_{param_name} = caml_copy_double({param_name});")
            elif param_type.name == "bool":
                lines.append(f"    ml_{param_name} = Val_bool({param_name});")
            elif param_type.name == "unit":
                lines.append(f"    ml_{param_name} = Val_unit;")
        else:
            # Complex types (options, lists, tuples, custom types)
            # These are passed as opaque pointers and cast to value
            lines.append(f"    ml_{param_name} = (value){param_name};")

        return lines

    def _convert_ocaml_to_c(self, return_type: IRType) -> list:
        """
        Generate code to convert OCaml return value to C.

        Handles primitives and complex types. Complex types are returned
        as opaque pointers to maintain GC-safety.
        """
        lines = []
        c_type = self._get_c_type(return_type)

        if return_type.is_primitive():
            if return_type.name == "string":
                lines.append(f"    {c_type} result = strdup(String_val(ml_result));")
                lines.append(f"    CAMLreturnT({c_type}, result);")
            elif return_type.name == "int":
                lines.append(f"    {c_type} result = Int_val(ml_result);")
                lines.append(f"    CAMLreturnT({c_type}, result);")
            elif return_type.name == "float":
                lines.append(f"    {c_type} result = Double_val(ml_result);")
                lines.append(f"    CAMLreturnT({c_type}, result);")
            elif return_type.name == "bool":
                lines.append(f"    {c_type} result = Bool_val(ml_result);")
                lines.append(f"    CAMLreturnT({c_type}, result);")
            elif return_type.name == "unit":
                lines.append("    CAMLreturn0;")
        else:
            # Complex types (options, lists, tuples, custom types)
            # Return as opaque pointer (cast from value)
            # Note: This keeps the value alive and GC-safe
            lines.append(f"    {c_type} result = (void*)ml_result;")
            lines.append(f"    CAMLreturnT({c_type}, result);")

        return lines

    def _get_c_type(self, ir_type: IRType) -> str:
        """
        Convert IR type to C type string.

        Primitives map to C native types.
        Complex types (options, lists, tuples, custom) map to void* (opaque).
        """
        if ir_type.is_primitive():
            return self.C_TYPE_MAP.get(ir_type.name, "char*")

        # Complex types are opaque pointers in C
        if ir_type.kind in (
            TypeKind.OPTION,
            TypeKind.LIST,
            TypeKind.TUPLE,
            TypeKind.CUSTOM,
            TypeKind.RECORD,
            TypeKind.VARIANT,
        ):
            return "void*"

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