"""
Infrastructure module base class.

This module provides a concrete `Module` implementation that satisfies the
domain-level `IModule` protocol. It implements common conveniences used by
neural network layers, including:

- parameter registration and storage
- submodule registration and storage
- recursive parameter traversal (`parameters`, `named_parameters`)
- `__call__` forwarding to `forward` for ergonomic invocation

This class is part of the infrastructure layer and is intended to be subclassed
by concrete layers (e.g., Linear, Conv2D, activations-as-modules, containers).
"""

from __future__ import annotations

from collections.abc import Iterable
from typing import Dict, Iterator, Optional, Any

from ..domain._module import IModule
from ..domain._parameter import IParameter
from ..domain.device._device import Device


class Module(IModule):
    """
    Infrastructure base class for layers/modules.

    This class provides a lightweight foundation for implementing trainable
    layers. Subclasses typically:
    - create `Parameter` instances (or other `IParameter` implementations),
    - register them (explicitly via `register_parameter` or implicitly via attribute assignment),
    - implement `forward` to define computation.

    Attributes
    ----------
    _parameters : Dict[str, IParameter]
        Mapping from parameter name to parameter object for this module.
    _modules : Dict[str, Module]
        Mapping from child module name to child module object for this module.

    Notes
    -----
    - This implementation supports *recursive* traversal over registered submodules.
    - Parameters and submodules can be registered explicitly (register_*) or
      implicitly by assigning them as attributes, e.g.:
          self.weight = Parameter(...)
          self.block = Sequential(...)
    - `__call__` delegates to `forward`, matching common deep learning framework conventions.
    """

    def __init__(self) -> None:
        """
        Initialize an empty module with no registered parameters/submodules.
        """
        # Use super().__setattr__ to avoid triggering our __setattr__ logic.
        super().__setattr__("_parameters", {})  # type: ignore[assignment]
        super().__setattr__("_modules", {})  # type: ignore[assignment]
        self._parameters: Dict[str, IParameter]
        self._modules: Dict[str, "Module"]

    def __setattr__(self, name: str, value) -> None:
        """
        Intercept attribute assignment to auto-register Parameters and child Modules.

        This mirrors common deep learning framework behavior and ensures that:
        - optimizers can discover parameters via `parameters()`
        - containers can recurse into submodules
        """
        # Let internal bookkeeping attributes pass through without registration.
        if name in {"_parameters", "_modules"}:
            super().__setattr__(name, value)
            return

        # Lazy imports to avoid hard cycles at import time.
        try:
            from ._parameter import Parameter  # concrete infra Parameter
        except Exception:  # pragma: no cover
            Parameter = None  # type: ignore[assignment]

        # If user assigns None, treat it as "unregister" if present.
        if value is None:
            if hasattr(self, "_parameters") and name in self._parameters:
                self._parameters.pop(name, None)
            if hasattr(self, "_modules") and name in self._modules:
                self._modules.pop(name, None)
            super().__setattr__(name, value)
            return

        # Auto-register concrete Parameters.
        if Parameter is not None and isinstance(value, Parameter):
            self._parameters[name] = value

        # Auto-register child Modules (infrastructure modules).
        elif isinstance(value, Module):
            self._modules[name] = value

        super().__setattr__(name, value)

    def register_parameter(self, name: str, param: Optional[IParameter]) -> None:
        """
        Register a parameter with this module.

        Parameters
        ----------
        name : str
            Name under which the parameter will be stored (e.g., "weight", "bias").
        param : Optional[IParameter]
            Parameter instance to register. If None, registration is skipped.

        Notes
        -----
        - If `param` is None, nothing is registered.
        - If the name already exists, it is overwritten intentionally.
        - This also sets the attribute on the module so `self.<name>` works.
        """
        if param is None:
            return
        self._parameters[name] = param
        super().__setattr__(name, param)

    def register_module(self, name: str, module: Optional["Module"]) -> None:
        """
        Register a child module with this module.

        Parameters
        ----------
        name : str
            Name under which the module will be stored.
        module : Optional[Module]
            Child module to register. If None, registration is skipped.

        Notes
        -----
        - If `module` is None, nothing is registered.
        - This also sets the attribute on the module so `self.<name>` works.
        """
        if module is None:
            return
        self._modules[name] = module
        super().__setattr__(name, module)

    def parameters(self) -> Iterable[IParameter]:
        """
        Return an iterable over this module's parameters (recursive).

        Returns
        -------
        Iterable[IParameter]
            Iterable of parameters registered on this module and all submodules.
        """
        # Yield own parameters first
        for p in self._parameters.values():
            yield p
        # Then recurse
        for m in self._modules.values():
            yield from m.parameters()

    def named_parameters(self, prefix: str = "") -> Iterator[tuple[str, IParameter]]:
        """
        Return an iterator over (name, parameter) pairs (recursive).

        Parameters
        ----------
        prefix : str
            Prefix to prepend to parameter names (used for recursion).

        Returns
        -------
        Iterator[tuple[str, IParameter]]
            Iterator yielding (fully_qualified_name, parameter).
        """
        base = prefix + "." if prefix else ""

        for name, p in self._parameters.items():
            yield (f"{base}{name}", p)

        for child_name, child in self._modules.items():
            child_prefix = f"{base}{child_name}"
            yield from child.named_parameters(child_prefix)

    def forward(self, x):
        """
        Execute the forward computation of the module.

        Subclasses must implement this.
        """
        raise NotImplementedError

    def __call__(self, x):
        """
        Call the module as a function, delegating to `forward`.
        """
        return self.forward(x)

    # ------------------------------------------------------------------
    # Serialization hooks (opt-in contract)
    # ------------------------------------------------------------------
    def get_config(self) -> Dict[str, Any]:
        """
        Return a JSON-serializable configuration for this module.

        Subclasses that participate in JSON-based model save/load MUST
        override this method.

        Raises
        ------
        NotImplementedError
            If the module does not support JSON serialization.
        """
        raise NotImplementedError(
            f"{self.__class__.__name__} does not implement get_config(). "
            "This module cannot be serialized to JSON."
        )

    @classmethod
    def from_config(cls, cfg: Dict[str, Any]) -> "Module":
        """
        Reconstruct a module from a JSON configuration.

        Subclasses that participate in JSON-based model save/load MUST
        override this method.

        Raises
        ------
        NotImplementedError
            If the module does not support JSON deserialization.
        """
        raise NotImplementedError(
            f"{cls.__name__} does not implement from_config(). "
            "This module cannot be deserialized from JSON."
        )

    def to(self, device: Device) -> "Module":
        """
        Move this module (recursively) to `device` by moving all registered Parameters.

        Notes
        -----
        - Uses `_parameters` / `_modules` registries as the source of truth.
          This avoids touching properties / methods / non-parameter attributes.
        - Assumes each Parameter/Tensor implements `.to(Device) -> same-type-like`.
        - Rebinds attributes so `self.weight`, etc. now point to the moved objects.
        """
        if not isinstance(device, Device):
            raise TypeError(f"`device` must be a Device, got {type(device)!r}.")

        # Move parameters on this module
        for name, p in list(self._parameters.items()):
            # Some registries may include None; be defensive
            if p is None:
                continue

            moved = p.to(device)

            # Update registry and attribute binding WITHOUT triggering __setattr__ logic
            # (avoids double-registration / side effects).
            self._parameters[name] = moved  # type: ignore[assignment]
            super().__setattr__(name, moved)

        # Recurse into submodules
        for _, m in self._modules.items():
            m.to(device)

        # Sync module metadata if the module stores `.device`
        if hasattr(self, "device"):
            try:
                super().__setattr__("device", device)
            except Exception:
                # best-effort: some modules may expose read-only device properties
                pass

        self._to_extra_(device)

        return self

    def to_(self, device: Device) -> "Module":
        """
        Move this module and all of its parameters to `device` *in-place*.

        This method performs a recursive, in-place device migration of all
        parameters registered on this module and its submodules. Unlike
        `Module.to()`, which may rebind parameters to newly created objects,
        `to_()` attempts to preserve the identity of each parameter whenever
        possible.

        Behavior
        --------
        - For each registered parameter:
            - If the parameter implements `to_()`, it is migrated in-place
            (object identity is preserved).
            - Otherwise, the parameter is migrated out-of-place via `to(device)`
            and rebound on the module as a fallback.
        - All child modules are recursively migrated using the same rules.

        Parameters
        ----------
        device : Device
            Target device to which all parameters should be moved.

        Returns
        -------
        Module
            This module (`self`), after in-place migration.

        Notes
        -----
        - This method relies exclusively on the `_parameters` and `_modules`
        registries and does not inspect arbitrary attributes.
        - In-place migration is best-effort and depends on parameter support
        for `to_()`. Parameters that do not implement `to_()` will be replaced
        by newly created objects.
        - Autograd context is not preserved across device transfers; parameters
        should be treated as graph breaks after migration.
        - Optimizers that hold references to parameters remain valid only if
        all parameters support true in-place migration.
        """

        if not isinstance(device, Device):
            raise TypeError(f"`device` must be a Device, got {type(device)!r}.")

        for name, p in self._parameters.items():
            if p is None:
                continue
            if hasattr(p, "to_"):
                p.to_(device)  # in-place
            else:
                # fallback: out-of-place then rebind
                moved = p.to(device)
                self._parameters[name] = moved  # type: ignore[assignment]
                super().__setattr__(name, moved)

        for _, m in self._modules.items():
            m.to_(device)

        # Sync module metadata if the module stores `.device`
        if hasattr(self, "device"):
            try:
                super().__setattr__("device", device)
            except Exception:
                # best-effort: some modules may expose read-only device properties
                pass

        self._to_extra_(device)

        return self

    def _to_extra_(self, device: Device) -> None:
        """
        Optional hook for migrating non-parameter tensors during `to_()`.

        This method is called by `Module.to_()` after parameter migration and
        allows subclasses to move additional tensors that are **not registered
        as Parameters** (e.g., running statistics, buffers, cached tensors).

        The default implementation is a no-op. Subclasses should override this
        method only if they own such tensors and must ensure migration is done
        **in-place** to preserve object identity.

        Parameters
        ----------
        device : Device
            Target device to which extra tensors should be moved.

        Notes
        -----
        - This method must not allocate new tensors.
        - Tensors handled here are expected to have `requires_grad=False`.
        - Modules without extra non-parameter tensors do not need to override this.
        """
        # subclasses override if they own non-parameter tensors that must move
        return
