"""HUGR package and pointed package interfaces."""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Generic, TypeVar, cast

from typing_extensions import deprecated

import hugr._serialization.extension as ext_s
import hugr.model as model
from hugr import ext
from hugr._hugr.linking import link_modules
from hugr.envelope import (
    EnvelopeConfig,
    _make_envelope,
    _make_envelope_str,
    read_envelope,
    read_envelope_str,
)
from hugr.hugr.base import Hugr
from hugr.ops import FuncDecl, FuncDefn, Op

if TYPE_CHECKING:
    from hugr.ext import Extension, ExtensionRegistry
    from hugr.hugr.node_port import Node

__all__ = [
    "Package",
    "PackagePointer",
    "ModulePointer",
    "ExtensionPointer",
    "NodePointer",
    "FuncDeclPointer",
    "FuncDefnPointer",
]


@dataclass(frozen=True)
class Package:
    """A package of HUGR modules and extensions.


    The HUGRs may refer to the included extensions or those not included.
    """

    #: HUGR modules in the package.
    modules: list[Hugr]
    #: Extensions included in the package.
    extensions: list[Extension] = field(default_factory=list)

    def _to_serial(self) -> ext_s.Package:
        return ext_s.Package(
            modules=[m._to_serial() for m in self.modules],
            extensions=[e._to_serial() for e in self.extensions],
        )

    @staticmethod
    def from_bytes(
        envelope: bytes, extensions: ExtensionRegistry | None = None
    ) -> Package:
        """Deserialize a byte string to a Package object.

        Some envelope formats can be read from a string. See :meth:`from_str`.

        Args:
            envelope: The byte string representing a Package.
            extensions: If not None, an extension registry to resolve the custom
                operations and types.

        Returns:
            The deserialized Package object.
        """
        return read_envelope(envelope, extensions)

    @staticmethod
    def from_str(envelope: str, extensions: ExtensionRegistry | None = None) -> Package:
        """Deserialize a string to a Package object.

        Not all envelope formats can be read from a string.
        See :meth:`from_bytes` for a more general method.

        Args:
            envelope: The string representing a Package.
            extensions: If not None, an extension registry to resolve the custom
                operations and types.

        Returns:
            The deserialized Package object.
        """
        return read_envelope_str(envelope, extensions)

    @staticmethod
    def from_model(
        package: model.Package,
        extensions: list[Extension] | None = None,
        *,
        resolve_from: ExtensionRegistry | None = None,
    ):
        """Deserialize a hugr model to a Package object.

        Args:
            package: The hugr model representing a Package.
            extensions: If not None, the extensions to include in the package.
            resolve_from: If not None, an extension registry to resolve the custom
                operations and types. Defaults to the standard extensions, plus any
                extensions included in the package.
        """
        from hugr.ext import ExtensionRegistry

        if extensions is None:
            extensions = []
        pkg = Package([Hugr.from_model(hugr) for hugr in package.modules], extensions)

        # Resolve extensions in the modules and inside the bundled extension definitions
        pkg.resolve_extensions(resolve_from or ExtensionRegistry())

        return pkg

    def to_bytes(self, config: EnvelopeConfig | None = None) -> bytes:
        """Serialize the package to a HUGR envelope byte string.

        Some envelope formats can be encoded into a string. See :meth:`to_str`.
        """
        config = config or EnvelopeConfig.BINARY
        return _make_envelope(self, config)

    def to_str(self, config: EnvelopeConfig | None = None) -> str:
        """Serialize the package to a HUGR envelope string.

        Not all envelope formats can be encoded into a string.
        See :meth:`to_bytes` for a more general method.
        """
        config = config or EnvelopeConfig.TEXT
        return _make_envelope_str(self, config)

    @deprecated("Use HUGR envelopes instead. See the `to_bytes` and `to_str` methods.")
    def to_json(self) -> str:
        """Serialize the package to a printable HUGR envelope string."""
        return self._to_serial().model_dump_json()

    @classmethod
    @deprecated(
        "Use HUGR envelopes instead. See the `from_bytes` and `from_str` methods."
    )
    def from_json(cls, json_str: str) -> Package:
        """Deserialize a JSON string to a Package object.

        Args:
            json_str: The JSON string representing a Package.

        Returns:
            The deserialized Package object.
        """
        return ext_s.Package.model_validate_json(json_str).deserialize()

    def to_model(self) -> model.Package:
        """Export the package as its hugr model representation.

        At the moment this does not yet contain the extensions.
        """
        return model.Package([module.to_model() for module in self.modules])

    def used_extensions(
        self, resolve_from: ext.ExtensionRegistry | None = None
    ) -> ext.ExtensionResolutionResult:
        """Get the extensions used by this Package, optionally resolving unresolved
        types and operations.

        This method modifies the HUGR modules and packaged extensions in-place when
        resolve_from is provided, replacing Custom operations with ExtOp operations and
        opaque types with ExtType when their extensions are found in the registry.

        Args:
            resolve_from: Optional extension registry to resolve against.
                If None, opaque types and Custom ops will not be resolved.

        Returns:
            The result of resolving the extensions, containing the used
            extensions and a list of referenced but unresolved extensions.

        Example:
            >>> from hugr.build import Dfg
            >>> Dfg(tys.Qubit).hugr.used_extensions().ids()
            {'prelude'}
        """
        from hugr.ext import ExtensionResolutionResult
        from hugr.std import _std_extensions

        packaged_extensions = ext.ExtensionRegistry.from_extensions(self.extensions)

        if resolve_from is not None:
            resolve_from.extend(_std_extensions())
            resolve_from.extend(packaged_extensions)

        # Packaged extensions are always marked as "used".
        result = ExtensionResolutionResult(
            used_extensions=packaged_extensions,
        )

        for module in self.modules:
            result.extend(module.used_extensions(resolve_from))

        result._extend_with_transitive_ops(resolve_from)

        return result

    def resolve_extensions(self, registry: ext.ExtensionRegistry) -> Package:
        """Resolve extension references in the types and operations of the package.

        This method modifies the HUGR modules and packaged extensions in-place,
        replacing Custom operations with ExtOp operations and
        opaque types with ExtType when their extensions are found in the registry.

        This is an alias for :meth:`used_extensions` that discards the computed
        extensions.

        Args:
            registry: The extension registry to resolve against.
        """
        self.used_extensions(resolve_from=registry)
        return self

    def link(self, *other: Package):
        """Link this package with other packages, returning a new package containing the
        extensions of all packages, as well as a single module created from linking the
        modules from all packages.

        Args:
            *other: Other packages to link with.

        Returns:
            A new package containing the modules and extensions of all packages.
        """
        from hugr.ext import ExtensionRegistry

        modules = self.modules[:]
        registry = ExtensionRegistry.from_extensions(self.extensions)
        for pkg in other:
            modules.extend(pkg.modules)
            for new_ext in pkg.extensions:
                registry.register(new_ext)

        if len(modules) == 0:
            return Package([], list(registry.all_extensions))

        result_module_bytes = modules[0].to_bytes(include_extensions=registry)
        for module in modules[1:]:
            result_module_bytes = link_modules(
                result_module_bytes, module.to_bytes(include_extensions=registry)
            )

        return Package(
            [Hugr.from_bytes(result_module_bytes)],
            list(registry.all_extensions),
        )


@dataclass(frozen=True)
class PackagePointer:
    """Classes that point to packages and their inner contents."""

    #: Package pointed to.
    package: Package


@dataclass(frozen=True)
class ModulePointer(PackagePointer):
    """Pointer to a module in a package.

    Args:
        package: Package pointed to.
        module_index: Index of the module in the package.
    """

    #: Index of the module in the package.
    module_index: int

    @property
    def module(self) -> Hugr:
        """Hugr definition of the module."""
        return self.package.modules[self.module_index]

    def to_executable_package(self) -> ExecutablePackage:
        """Create an executable package from a module containing a main function.

        Raises:
            ValueError: If the module does not contain a main function.
        """
        module = self.module
        try:
            main_node = next(
                n
                for n in module.children()
                if isinstance((f_def := module[n].op), FuncDefn)
                and f_def.f_name == "main"
            )
        except StopIteration as e:
            msg = "Module does not contain a main function"
            raise ValueError(msg) from e
        return ExecutablePackage(self.package, self.module_index, main_node)


@dataclass(frozen=True)
class ExtensionPointer(PackagePointer):
    """Pointer to an extension in a package.

    Args:
        package: Package pointed to.
        extension_index: Index of the extension in the package.
    """

    #: Index of the extension in the package.
    extension_index: int

    @property
    def extension(self) -> Extension:
        """Extension definition."""
        return self.package.extensions[self.extension_index]


OpType = TypeVar("OpType", bound=Op)


@dataclass(frozen=True)
class NodePointer(Generic[OpType], ModulePointer):
    """Pointer to a node in a module.

    Args:
        package: Package pointed to.
        module_index: Index of the module in the package.
        node: Node pointed to
    """

    #: Node pointed to.
    node: Node

    @property
    def node_op(self) -> OpType:
        """Get the operation of the node."""
        return cast(OpType, self.module[self.node].op)


@dataclass(frozen=True)
class FuncDeclPointer(NodePointer[FuncDecl]):
    """Pointer to a function declaration in a module.

    Args:
        package: Package pointed to.
        module_index: Index of the module in the package.
        node: Node containing the function declaration.
    """

    @property
    def func_decl(self) -> FuncDecl:
        """Function declaration."""
        return self.node_op


@dataclass(frozen=True)
class FuncDefnPointer(NodePointer[FuncDefn]):
    """Pointer to a function definition in a module.

    Args:
        package: Package pointed to.
        module_index: Index of the module in the package.
        node: Node containing the function definition
    """

    @property
    def func_defn(self) -> FuncDefn:
        """Function definition."""
        return self.node_op


@dataclass(frozen=True)
class ExecutablePackage(FuncDefnPointer):
    """PackagePointer with a defined entrypoint node.

    Args:
        package: Package pointed to.
        module_index: Index of the module in the package.
        node: Node containing the entry point function definition.
    """

    @property
    def entry_point_node(self) -> Node:
        """Get the entry point node of the package."""
        return self.node
