Decorators

class Decorators:
    """Various decorators intended to be attached to a KCLayout."""

    def __init__(self, kcl: KCLayout) -> None:
        """Just set the standard `@cell` decorator."""
        self._cell = kcl.cell

    @overload
    def module_cell[**KCellParams, KC: ProtoTKCell[Any]](
        self,
        _func: Callable[KCellParams, KC],
        /,
    ) -> Callable[KCellParams, KC]: ...

    @overload
    def module_cell[**KCellParams, KC: ProtoTKCell[Any]](
        self, /, **kwargs: Unpack[ModuleCellKWargs]
    ) -> Callable[[Callable[KCellParams, KC]], Callable[KCellParams, KC]]: ...

    def module_cell[**KCellParams, KC: ProtoTKCell[Any]](
        self,
        _func: Callable[KCellParams, KC] | None = None,
        /,
        **kwargs: Unpack[ModuleCellKWargs],
    ) -> (
        Callable[KCellParams, KC]
        | Callable[[Callable[KCellParams, KC]], Callable[KCellParams, KC]]
    ):
        """Constructs the `@module_cell` decorator on KCLayout.decorators."""

        def mc(
            **kwargs: Unpack[ModuleCellKWargs],
        ) -> Callable[[Callable[KCellParams, KC]], Callable[KCellParams, KC]]:
            return _module_cell(self._cell, **kwargs)  # ty:ignore[invalid-argument-type]

        return mc(**kwargs) if _func is None else mc(**kwargs)(_func)

class KCellDecorator[**KCellParams, K: ProtoKCell[Any, Any]](Protocol):
    """Signature of the `@cell` decorator."""

    def __call__(
        self, **kwargs: Unpack[KCellDecoratorKWargs]
    ) -> Callable[[Callable[KCellParams, K]], Callable[KCellParams, K]]:
        """__call__ implementation."""
        ...

class KCellDecoratorKWargs(TypedDict, total=False):
    """KWargs for `@cell`."""

    set_settings: bool
    set_name: bool
    check_ports: bool
    check_pins: bool
    check_instances: CheckInstances | None
    snap_ports: bool
    add_port_layers: bool
    cache: Cache[int, Any] | dict[int, Any] | None
    basename: str | None
    drop_params: list[str]
    register_factory: bool
    overwrite_existing: bool | None
    layout_cache: bool | None
    info: dict[str, MetaData] | None
    post_process: Iterable[Callable[[TKCell], None]]
    debug_names: bool | None

class ModuleCellKWargs(TypedDict, total=False):
    """KWargs for `@module_cell`."""

    set_settings: bool
    set_name: bool
    check_ports: bool
    check_pins: bool
    check_instances: CheckInstances | None
    snap_ports: bool
    add_port_layers: bool
    cache: Cache[int, Any] | dict[int, Any] | None
    drop_params: list[str]
    register_factory: bool
    overwrite_existing: bool | None
    layout_cache: bool | None
    info: dict[str, MetaData] | None
    post_process: Iterable[Callable[[TKCell], None]]
    debug_names: bool | None

class ModuleDecorator(Protocol):
    """Signature of the `@module_cell` decorator."""

    def __call__[**KCellParams, K: ProtoKCell[Any, Any]](
        self, /, **kwargs: Unpack[ModuleCellKWargs]
    ) -> Callable[[Callable[KCellParams, K]], Callable[KCellParams, K]]:
        """__call__ implementation."""
        ...

@final
class WrappedKCellFunc[**KCellParams, KC: ProtoTKCell[Any]]:
    _f: Callable[KCellParams, KC]
    _f_orig: Callable[KCellParams, ProtoTKCell[Any]]
    _f_schematic: Callable[KCellParams, TSchematic[Any]] | None = None
    cache: Cache[Hashable, Any] | dict[Hashable, Any]
    name: str
    kcl: KCLayout
    output_type: type[KC]
    lvs_equivalent_ports: list[list[str]] | None = None
    ports_definition: PortsDefinition | None = None
    tags: set[str]
    signature: inspect.Signature
    _sig_params: SignatureParams

    def __init__(
        self,
        *,
        kcl: KCLayout,
        f: Callable[KCellParams, ProtoTKCell[Any]],
        sig: inspect.Signature,
        output_type: type[KC],
        cache: Cache[Hashable, KC] | dict[Hashable, KC],
        set_settings: bool,
        set_name: bool,
        check_ports: bool,
        check_pins: bool,
        check_instances: CheckInstances,
        check_unnamed_cells: CheckUnnamedCells,
        snap_ports: bool,
        add_port_layers: bool,
        basename: str | None,
        drop_params: Sequence[str],
        overwrite_existing: bool | None,
        layout_cache: bool | None,
        info: dict[str, MetaData] | None,
        post_process: Iterable[Callable[[ProtoTKCell[Any]], None]],
        debug_names: bool,
        lvs_equivalent_ports: list[list[str]] | None = None,
        ports: PortsDefinition | None = None,
        tags: Sequence[str] | None = None,
        schematic_function: Callable[KCellParams, TSchematic[Any]] | None = None,
        type_serializers: Sequence[tuple[type | UnionType, Callable[[Any], Any]]] = (
            (
                SymmetricalCrossSection
                | CrossSection
                | DCrossSection
                | AsymmetricalCrossSection
                | AsymmetricCrossSection
                | DAsymmetricCrossSection,
                attrgetter("name"),
            ),
        ),
        type_hints_serializer: dict[
            type | UnionType | TypeAliasType, Callable[[Any], Any]
        ]
        | None = None,
    ) -> None:
        self.kcl = kcl
        self.output_type = output_type
        self.name = basename or get_function_name(f)
        self.ports_definition = ports.copy() if ports is not None else None
        self.tags = set(tags) if tags else set()
        self._f_schematic = schematic_function

        # Pre-compute static signature metadata once at decoration time
        sig_params = SignatureParams(sig)
        self.signature = sig
        self._sig_params = sig_params
        # Resolve annotations to concrete types for the type-hint serializers.
        # Guard against ``NameError`` etc. raised by ``from __future__ import
        # annotations`` functions whose hints reference ``TYPE_CHECKING``-only
        # imports; an unresolvable hint simply opts out of hint-based
        # serialization rather than breaking decoration.
        try:
            hints = get_type_hints(f)
        except Exception:
            hints = {}
        if type_hints_serializer is None:
            # ``cast`` because ty cannot recognize a PEP-695 ``type`` alias used
            # as a value (the ``CrossSectionSpec`` dict key) as assignable to
            # ``TypeAliasType``, even though it is one at runtime.
            type_hints_serializer = cast(
                "dict[type | UnionType | TypeAliasType, Callable[[Any], Any]]",
                {CrossSectionSpec: kcl_cross_section_serializer(kcl=kcl)},
            )

        @functools.wraps(f)
        def wrapper_autocell(
            *args: KCellParams.args, **kwargs: KCellParams.kwargs
        ) -> KC:
            params = _parse_params(
                sig_params.defaults, sig_params.names, kcl, args, kwargs
            )

            with kcl.thread_lock:
                cell_ = wrapped_cell(**params)
                if cell_.destroyed():
                    # If any cell has been destroyed, we should clean up the cache.
                    # Delete all the KCell entrances in the cache which have
                    # `destroyed() == True`
                    deleted_cell_hashes: list[Hashable] = [
                        _hash_item
                        for _hash_item, _cell_item in cache.items()
                        if _cell_item.destroyed()
                    ]
                    for _dch in deleted_cell_hashes:
                        del cache[_dch]
                    cell_ = wrapped_cell(**params)

                if info is not None:
                    cell_.info.update(info)

                return cell_

        @cached(
            cache=cache,
            lock=RLock(),
            key=get_keys_function(
                hints=hints,
                drop_args=drop_params,
                serialize_types=type_serializers,
                serialize_hints=type_hints_serializer,
            ),
        )
        def wrapped_cell(**params: Any) -> KC:

            _params_to_original(params)

            old_future_name: str | None = None
            if set_name:
                if basename is not None:
                    name = get_cell_name(basename, **params)
                else:
                    name = get_cell_name(self.name, **params)
                old_future_name = kcl._future_cell_name
                kcl._future_cell_name = name
                if layout_cache:
                    if overwrite_existing:
                        for c in list(kcl.cells(kcl._future_cell_name)):
                            kcl[c.cell_index()].delete()
                    else:
                        layout_cell = kcl.layout_cell(kcl._future_cell_name)
                        if layout_cell is not None:
                            logger.debug(
                                "Loading {} from layout cache",
                                kcl._future_cell_name,
                            )
                            return kcl.get_cell(layout_cell.cell_index(), output_type)
                logger.debug(f"Constructing {kcl._future_cell_name}")
                name_: str | None = name
            else:
                name_ = None
            cell = f(**params)  # ty:ignore[missing-argument]
            if cell is None:
                raise TypeError(
                    f"The cell function {self.name!r} in {str(self.file)!r}"
                    " returned None. Did you forget to return the cell or component"
                    " at the end of the function?"
                )
            if not isinstance(cell, ProtoTKCell):
                raise TypeError(
                    f"The cell function {self.name!r} in {str(self.file)!r}"
                    f" returned {type(cell)=}. The `@cell` decorator only supports"
                    " KCell/DKCell or any SubClass such as Component."
                )

            logger.debug("Constructed {}", name_ or cell.name)

            if cell.locked:
                # If the cell is locked, it likely comes
                # from a cache and should be copied first
                cell = cell.dup(new_name=kcl._future_cell_name)
            if overwrite_existing:
                _overwrite_existing(name_, cell, kcl)
            if set_name and name_:
                if debug_names and cell.kcl.layout_cell(name_) is not None:
                    logger.opt(depth=4).error(
                        "KCell with name {name} exists already. Duplicate "
                        "occurrence in module '{module}' at "
                        "line {lno}",
                        name=name_,
                        module=f.__module__,
                        function_name=get_function_name(f),
                        lno=inspect.getsourcelines(f)[1],
                    )
                    raise CellNameError(f"KCell with name {name_} exists already.")

                cell.name = name_
                kcl._future_cell_name = old_future_name
            if set_settings:
                _set_settings(cell, f, drop_params, params, sig_params.units, basename)
            if check_ports:
                _check_ports(cell)
            if check_pins:
                _check_pins(cell)
            match check_unnamed_cells:
                case CheckUnnamedCells.RAISE | CheckUnnamedCells.WARNING:
                    unnamed_cells: list[str] = []
                    for ci in cell.kdb_cell.each_child_cell():
                        c = cell.kcl[ci]
                        if re.fullmatch(_fixed_unnamed_pattern, c.name):
                            factory_name = c.basename or c.function_name
                            factory_string = (
                                f"factory_name={factory_name!r}"
                                if factory_name
                                else "Cell without cell function"
                            )
                            unnamed_cells.append(f"{c.name} ({factory_string})")
                    if unnamed_cells:
                        msg = (
                            f"Cell {cell.name!r} has"
                            " unnamed cells instantiated:\n" + "\n".join(unnamed_cells)
                        )
                        if check_unnamed_cells == CheckUnnamedCells.RAISE:
                            raise ValueError(msg)
                        logger.warning(msg)

            _check_instances(cell, kcl, check_instances)
            cell.insert_vinsts(recursive=False)
            if snap_ports:
                _snap_ports(cell, kcl)
            if add_port_layers:
                _add_port_layers(cell, kcl)
            _post_process(cell, post_process)
            cell.base.lock()
            _check_cell(cell, kcl)
            if self.ports_definition is not None:
                port_lengths = 0
                for direction in Direction:
                    port_lengths += len(self.ports_definition.get(direction, []))
                mapping = {0: "right", 1: "top", 2: "left", 3: "bottom"}
                if len(cell.ports) != port_lengths:
                    received_ports = PortsDefinition()
                    for port in cell.ports:
                        mapped: Direction = Direction(mapping[port.trans.angle])
                        if mapped not in received_ports:
                            received_ports[mapped] = []  # ty:ignore[invalid-key]
                        received_ports[mapped].append(port.name)  # ty:ignore[invalid-key]
                    raise ValueError(
                        "The `@cell` decorator defines ports, but they do not match"
                        " the extracted ports. Declared ports: "
                        f"{self.ports_definition}"
                        ", Received ports: "
                        f"{received_ports}"
                    )

                if check_ports:
                    found_errors = False
                    for port in cell.ports:
                        if (
                            port.name
                            not in self.ports_definition[mapping[port.trans.angle]]  # ty:ignore[invalid-key]
                        ):
                            found_errors = True
                    if found_errors:
                        received_ports = PortsDefinition()
                        for port in cell.ports:
                            mapped = Direction(mapping[port.trans.angle])
                            if mapped not in received_ports:
                                received_ports[mapped] = []  # ty:ignore[invalid-key]
                            received_ports[mapped].append(port.name)  # ty:ignore[invalid-key]
                        raise ValueError(
                            "The `@cell` decorator defines ports, but they do not"
                            " match the extracted ports. Declared ports: "
                            f"{self.ports_definition}"
                            ", Received ports: "
                            f"{received_ports}"
                        )
                else:
                    port_names: list[str | None] = []
                    for direction in Direction:
                        if direction in self.ports_definition:
                            port_names.extend(self.ports_definition[direction])  # ty:ignore[invalid-key]

                    for port in cell.ports:
                        if port.name not in port_names:
                            found_errors = True
                    if found_errors:
                        raise ValueError(
                            "The `@cell` decorator defines ports, but they do not"
                            " match the extracted ports. Declared ports: "
                            f"{port_names}"
                            ", Received ports: "
                            f"{[p.name for p in cell.ports]}"
                        )

            return output_type(base=cell.base)

        self._f = wrapper_autocell
        self._f_orig = f
        self.cache = cache
        self.lvs_equivalent_ports = lvs_equivalent_ports
        functools.update_wrapper(self, f)

    def __call__(self, *args: KCellParams.args, **kwargs: KCellParams.kwargs) -> KC:
        return self._f(*args, **kwargs)

    def __len__(self) -> int:
        del_cells = [hk for hk, kc in self.cache.items() if kc._destroyed()]
        for hk in del_cells:
            del self.cache[hk]

        return len(self.cache)

    @functools.cached_property
    def file(self) -> Path:
        return _get_path(self._f_orig)

    @property
    def qualified_name(self) -> str:
        return f"{self._f_orig.__module__}.{self._f_orig.__qualname__}"  # ty:ignore[unresolved-attribute]

    def prune(self) -> None:
        cells = [c for c in self.cache.values() if not c._destroyed()]
        caller_cis = {
            ci for cell in cells for ci in cell.caller_cells() if not cell._destroyed()
        }
        caller_cis |= {c.cell_index() for c in cells}
        self.kcl.delete_cells(list(caller_cis))

        self.kcl.cleanup()
        self.cache.clear()

    def schematic_driven(self) -> bool:
        return self._f_schematic is not None

    def get_schematic(
        self, *args: KCellParams.args, **kwargs: KCellParams.kwargs
    ) -> TSchematic[Any]:
        if self._f_schematic is None:
            raise ValueError(
                f"(D)KCellFunction {self.name} is not schematic driven and therefore"
                " cannot return a schematic."
            )
        return self._f_schematic(*args, **kwargs)

    def metadata_providers(self) -> tuple[_FactoryMetadataProviderRecord, ...]:
        """Return all metadata provider records registered for this factory."""
        return self.kcl.metadata_providers_for(self)

    def has_metadata(self) -> bool:
        """Whether any metadata providers are registered for this factory."""
        return bool(self.metadata_providers())

    def get_metadata(
        self, *args: KCellParams.args, **kwargs: KCellParams.kwargs
    ) -> FactoryMetadata:
        """Resolve all metadata providers into a single ``FactoryMetadata``.

        Accepts the same parameters as the cell factory itself.
        """
        params = _parse_params(
            self._sig_params.defaults,
            self._sig_params.names,
            self.kcl,
            args,
            kwargs,
        )
        return resolve_metadata(self.metadata_providers(), params)

@final
class WrappedVKCellFunc[**VKCellParams, VK: VKCell]:
    _f: Callable[VKCellParams, VK]
    _f_orig: Callable[VKCellParams, VKCell]
    cache: Cache[Hashable, VK] | dict[Hashable, Any]
    name: str
    kcl: KCLayout
    output_type: type[VK]
    lvs_equivalent_ports: list[list[str]] | None = None
    ports_definition: PortsDefinition | None = None
    tags: set[str]
    signature: inspect.Signature
    _sig_params: SignatureParams

    def __init__(
        self,
        *,
        kcl: KCLayout,
        f: Callable[VKCellParams, VKCell],
        sig: inspect.Signature,
        output_type: type[VK],
        cache: Cache[Hashable, VK] | dict[Hashable, VK],
        set_settings: bool,
        set_name: bool,
        check_ports: bool,
        check_pins: bool,
        check_unnamed_cells: CheckUnnamedCells,
        add_port_layers: bool,
        basename: str | None,
        drop_params: Sequence[str],
        info: dict[str, MetaData] | None,
        post_process: Iterable[Callable[[VKCell], None]],
        lvs_equivalent_ports: list[list[str]] | None = None,
        ports: PortsDefinition | None = None,
        tags: Sequence[str] | None = None,
        type_serializers: Sequence[tuple[type | UnionType, Callable[[Any], Any]]] = (
            (
                SymmetricalCrossSection
                | CrossSection
                | DCrossSection
                | AsymmetricalCrossSection
                | AsymmetricCrossSection
                | DAsymmetricCrossSection,
                attrgetter("name"),
            ),
        ),
        type_hints_serializer: dict[
            type | UnionType | TypeAliasType, Callable[[Any], Any]
        ]
        | None = None,
    ) -> None:
        self.kcl = kcl
        self.output_type = output_type
        self.name = basename or get_function_name(f)
        self.ports_definitions = ports.copy() if ports is not None else None
        self.tags = set(tags) if tags else set()

        # Pre-compute static signature metadata once at decoration time
        sig_params = SignatureParams(sig)
        self.signature = sig
        self._sig_params = sig_params
        # Resolve annotations to concrete types for the type-hint serializers (see
        # the equivalent block in ``WrappedKCellFunc``); an unresolvable hint simply
        # opts out of hint-based serialization rather than breaking decoration.
        try:
            hints = get_type_hints(f)
        except Exception:
            hints = {}
        if type_hints_serializer is None:
            type_hints_serializer = cast(
                "dict[type | UnionType | TypeAliasType, Callable[[Any], Any]]",
                {CrossSectionSpec: kcl_cross_section_serializer(kcl=kcl)},
            )

        @functools.wraps(f)
        def wrapper_autocell(
            *args: VKCellParams.args, **kwargs: VKCellParams.kwargs
        ) -> VK:
            params = _parse_params(
                sig_params.defaults, sig_params.names, kcl, args, kwargs
            )

            with kcl.thread_lock:
                cell_ = wrapped_cell(**params)

                if info is not None:
                    cell_.info.update(info)

                return cell_

        @cached(
            cache=cache,
            lock=RLock(),
            key=get_keys_function(
                hints=hints,
                drop_args=drop_params,
                serialize_types=type_serializers,
                serialize_hints=type_hints_serializer,
            ),
        )
        def wrapped_cell(**params: Any) -> VK:

            _params_to_original(params)

            old_future_name: str | None = None
            if set_name:
                if basename is not None:
                    name = get_cell_name(basename, **params)
                else:
                    name = get_cell_name(self.name, **params)
                old_future_name = kcl._future_cell_name
                kcl._future_cell_name = name
                logger.debug(f"Constructing {kcl._future_cell_name}")
                name_: str | None = name
            else:
                name_ = None
            cell = f(**params)  # ty:ignore[missing-argument]
            if cell is None:
                raise TypeError(
                    f"The cell function {self.name!r} in {str(self.file)!r}"
                    " returned None. Did you forget to return the cell or component"
                    " at the end of the function?"
                )
            if not isinstance(cell, VKCell):
                raise TypeError(
                    f"The cell function {self.name!r} in {str(self.file)!r}"
                    f" returned {type(cell)=}. The `@vcell` decorator only supports"
                    " VKCell or any SubClass such as ComponentAllAngle."
                )

            logger.debug("Constructed {}", name_ or cell.name)

            if cell.locked:
                # If the cell is locked, it comes from a cache (most likely)
                # and should be copied first
                cell = cell.dup(new_name=kcl._future_cell_name)
            if set_name and name_:
                cell.name = name_
                kcl._future_cell_name = old_future_name
            if set_settings:
                _set_settings(cell, f, drop_params, params, sig_params.units, basename)
            if check_ports:
                _check_ports(cell)
            if check_pins:
                _check_pins(cell)
            match check_unnamed_cells:
                case CheckUnnamedCells.RAISE | CheckUnnamedCells.WARNING:
                    unnamed_cells: set[str] = set()
                    for inst in cell.insts:
                        c = inst.cell
                        if re.fullmatch(_fixed_unnamed_pattern, c.name or ""):
                            factory_name = c.basename or c.function_name
                            factory_string = (
                                f"factory_name={factory_name!r}"
                                if factory_name
                                else "Cell without cell function"
                            )
                            unnamed_cells.add(f"{c.name} ({factory_string})")
                    if unnamed_cells:
                        msg = (
                            f"Cell {cell.name!r} has"
                            " unnamed cells instantiated:\n" + "\n".join(unnamed_cells)
                        )
                        if check_unnamed_cells == CheckUnnamedCells.RAISE:
                            raise ValueError(msg)
                        logger.warning(msg)
            if add_port_layers:
                _add_port_layers_vkcell(cell, kcl)
            _post_process(cell, post_process)
            cell.base.lock()
            _check_cell(cell, kcl)
            if self.ports_definition is not None:
                port_lengths = 0
                for direction in Direction:
                    port_lengths += len(self.ports_definition.get(direction, []))
                mapping = {0: "right", 1: "top", 2: "left", 3: "bottom"}
                if len(cell.ports) != port_lengths:
                    received_ports = PortsDefinition()
                    for port in cell.ports:
                        mapped: Direction = Direction(mapping[port.trans.angle])
                        if mapped not in received_ports:
                            received_ports[mapped] = []  # ty:ignore[invalid-key]
                        received_ports[mapped].append(port.name)  # ty:ignore[invalid-key]
                    raise ValueError(
                        "The `@cell` decorator defines ports, but they do not match"
                        " the extracted ports. Declared ports: "
                        f"{self.ports_definition}"
                        ", Received ports: "
                        f"{received_ports}"
                    )

                port_names: list[str | None] = []
                for direction in Direction:
                    if direction in self.ports_definition:
                        port_names.extend(self.ports_definition[direction])  # ty:ignore[invalid-key]

                for port in cell.ports:
                    if port.name not in port_names:
                        found_errors = True
                if found_errors:
                    raise ValueError(
                        "The `@cell` decorator defines ports, but they do not"
                        " match the extracted ports. Declared ports: "
                        f"{port_names}"
                        ", Received ports: "
                        f"{[p.name for p in cell.ports]}"
                    )

            return output_type(base=cell.base)

        self._f = wrapper_autocell
        self._f_orig = f
        self.cache = cache
        self.lvs_equivalent_ports = lvs_equivalent_ports
        functools.update_wrapper(self, f)

    def __call__(self, *args: Any, **kwargs: Any) -> VK:
        return self._f(*args, **kwargs)

    def __len__(self) -> int:
        return len(self.cache)

    @functools.cached_property
    def file(self) -> Path:
        return _get_path(self._f_orig)

    @property
    def qualified_name(self) -> str:
        return f"{self._f_orig.__module__}.{self._f_orig.__qualname__}"  # ty:ignore[unresolved-attribute]

    def metadata_providers(self) -> tuple[_FactoryMetadataProviderRecord, ...]:
        """Return all metadata provider records registered for this factory."""
        return self.kcl.metadata_providers_for(self)

    def has_metadata(self) -> bool:
        """Whether any metadata providers are registered for this factory."""
        return bool(self.metadata_providers())

    def get_metadata(
        self, *args: VKCellParams.args, **kwargs: VKCellParams.kwargs
    ) -> FactoryMetadata:
        """Resolve all metadata providers into a single ``FactoryMetadata``.

        Accepts the same parameters as the cell factory itself.
        """
        params = _parse_params(
            self._sig_params.defaults,
            self._sig_params.names,
            self.kcl,
            args,
            kwargs,
        )
        return resolve_metadata(self.metadata_providers(), params)