"""Overlay and lazily processed library views."""
from __future__ import annotations

from collections import defaultdict
from dataclasses import dataclass
from typing import TYPE_CHECKING, Literal, Self, cast
import copy

from ..error import LibraryError
from ..pattern import Pattern, map_layers, map_targets
from .base import ILibrary, ILibraryView
from .capabilities import IBorrowing, IMaterializable
from .utils import (
    INameView,
    dangling_mode_t,
    _plan_source_names,
    _rename_patterns,
    _source_rename_map,
    _validate_dangling_mode,
    )
from .mapping import LibraryView

if TYPE_CHECKING:
    from collections.abc import Callable, Iterator, Mapping, Sequence

    import numpy
    from numpy.typing import NDArray

    from ..ports import Port
    from ..utils import layer_t


@dataclass
class _SourceLayer:
    """ One imported source layer tracked by an `OverlayLibrary`. """
    library: ILibraryView
    source_target_map: dict[str, str]
    child_graph: dict[str, set[str]]


@dataclass(frozen=True)
class _SourceEntry:
    """ Reference to a single visible source-backed cell in an overlay. """
    layer_index: int
    source_name: str


def _materialize_detached_pattern(view: ILibraryView, name: str) -> Pattern:
    if isinstance(view, IMaterializable):
        return view.materialize_detached(name)
    return view[name].deepcopy()


class _ProcessedLibraryView(ILibraryView, IMaterializable, IBorrowing):
    """Shared detached-materialization behavior for read-only processing views."""

    def __init__(
            self,
            source: ILibraryView,
            *,
            copy_through: bool,
            ) -> None:
        self._source = source
        self._copy_through = copy_through
        self._cache: dict[str, Pattern] = {}
        self._lookups_in_progress: list[str] = []

    def __getitem__(self, key: str) -> Pattern:
        return self.materialize(key, persist=True)

    def __iter__(self) -> Iterator[str]:
        return iter(self._source)

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

    def __contains__(self, key: object) -> bool:
        return key in self._source

    def _process_pattern(self, name: str, pattern: Pattern) -> Pattern:
        """Apply this view's processing to one detached source pattern."""
        raise NotImplementedError

    def _materialize_uncached_detached(self, name: str) -> Pattern:
        if name in self._lookups_in_progress:
            chain = ' -> '.join(self._lookups_in_progress + [name])
            raise LibraryError(
                f'Detected circular reference or recursive lookup of "{name}".\n'
                f'Lookup chain: {chain}\n'
                'This may be caused by an invalid (cyclical) reference, or buggy code.'
                )

        self._lookups_in_progress.append(name)
        try:
            pattern = _materialize_detached_pattern(self._source, name)
            pattern = self._process_pattern(name, pattern)
        finally:
            self._lookups_in_progress.pop()

        return pattern

    def materialize(self, name: str, *, persist: bool = True) -> Pattern:
        if name in self._cache:
            return self._cache[name]

        pattern = self._materialize_uncached_detached(name)

        if persist:
            self._cache[name] = pattern
        return pattern

    def materialize_detached(self, name: str) -> Pattern:
        if name in self._cache:
            return self._cache[name].deepcopy()
        return self._materialize_uncached_detached(name)

    def materialize_many_detached(
            self,
            names: Sequence[str],
            ) -> LibraryView:
        ordered_names = tuple(dict.fromkeys(names))
        result: dict[str, Pattern] = {}
        uncached = [name for name in ordered_names if name not in self._cache]

        for name in ordered_names:
            if name in self._cache:
                result[name] = self._cache[name].deepcopy()

        if uncached:
            if isinstance(self._source, IMaterializable):
                source_patterns = self._source.materialize_many_detached(uncached)
            else:
                source_patterns = LibraryView({name: self._source[name].deepcopy() for name in uncached})
            for name in uncached:
                if name in self._lookups_in_progress:
                    chain = ' -> '.join(self._lookups_in_progress + [name])
                    raise LibraryError(f'Detected circular reference or recursive lookup of "{name}".\nLookup chain: {chain}')
                self._lookups_in_progress.append(name)
                try:
                    result[name] = self._process_pattern(name, source_patterns[name])
                finally:
                    self._lookups_in_progress.pop()

        return LibraryView({name: result[name] for name in ordered_names})

    def source_order(self) -> tuple[str, ...]:
        return self._source.source_order()

    def borrowed_sources(self) -> tuple[ILibraryView, ...]:
        return (self._source,)

    def source_cell(self, name: str) -> tuple[ILibraryView, str] | None:
        if not self._copy_through or name not in self._source or name in self._cache:
            return None
        return self._source, name

    def child_graph(
            self,
            dangling: dangling_mode_t = 'error',
            ) -> dict[str, set[str]]:
        _validate_dangling_mode(dangling)
        return self._source.child_graph(dangling=dangling)

    def find_refs_local(
            self,
            name: str,
            parent_graph: dict[str, set[str]] | None = None,
            dangling: dangling_mode_t = 'error',
            ) -> dict[str, list[NDArray[numpy.float64]]]:
        _validate_dangling_mode(dangling)
        finder = getattr(self._source, 'find_refs_local', None)
        if callable(finder):
            return cast('dict[str, list[NDArray[numpy.float64]]]', finder(name, parent_graph=parent_graph, dangling=dangling))
        return super().find_refs_local(name, parent_graph=parent_graph, dangling=dangling)


class PortLoadView(_ProcessedLibraryView):
    """
    Read-only view which loads or applies ports on first materialization.

    The wrapped source remains untouched; this view owns a separate processed
    cache so direct-copy workflows can continue to use the raw source view.
    The view borrows its source: callers must keep the source open for the
    lifetime of the view and close the source themselves.

    Graph queries and source ordering are delegated to the wrapped source,
    while `source_cell()` exposes unchanged layout provenance and `__getitem__`
    and `materialize_many()` return port-imported patterns.
    """

    def __init__(
            self,
            source: ILibraryView,
            *,
            layers: Sequence[layer_t] = (),
            max_depth: int = 0,
            skip_subcells: bool = True,
            ports: Mapping[str, Mapping[str, Port]] | None = None,
            replace: bool = False,
            ) -> None:
        super().__init__(source, copy_through=True)
        self._layers = tuple(layers)
        self._max_depth = max_depth
        self._skip_subcells = skip_subcells
        self._ports = {
            name: copy.deepcopy(dict(cell_ports))
            for name, cell_ports in (ports or {}).items()
        }
        self._replace = replace

    def _process_pattern(self, name: str, pattern: Pattern) -> Pattern:
        from ..utils.ports2data import data_to_ports        # noqa: PLC0415

        if self._layers:
            pattern = data_to_ports(
                layers=self._layers,
                library=self,
                pattern=pattern,
                name=name,
                max_depth=self._max_depth,
                skip_subcells=self._skip_subcells,
                )
        if name in self._ports:
            ports = copy.deepcopy(self._ports[name])
            if self._replace:
                pattern.ports = ports
            else:
                pattern.ports.update(ports)
        return pattern


class LayerMappedView(_ProcessedLibraryView):
    """
    Read-only view which remaps shape and label layers on materialization.

    The wrapped source remains untouched. By default, source-aware writers
    must materialize and serialize every mapped cell. With `copy_through=True`,
    unmaterialized cells may instead be copied unchanged from their source;
    persistent access maps and caches a cell, disabling copy-through for it.
    """

    def __init__(
            self,
            source: ILibraryView,
            map_layer: Callable[[layer_t], layer_t],
            *,
            copy_through: bool = False,
            ) -> None:
        super().__init__(source, copy_through=copy_through)
        self._map_layer = map_layer

    def _process_pattern(self, name: str, pattern: Pattern) -> Pattern:
        _ = name
        pattern.shapes = map_layers(pattern.shapes, self._map_layer)
        pattern.labels = map_layers(pattern.labels, self._map_layer)
        return pattern


class OverlayLibrary(ILibrary, IMaterializable, IBorrowing):
    """
    Mutable overlay over one or more source libraries.

    Source-backed cells remain lazy until accessed through `__getitem__`, which
    persistently materializes a detached, overlay-owned `Pattern`.

    Source libraries must remain open and must not be mutated after they are
    added. The overlay borrows each source and snapshots its names, hierarchy,
    and initial visible-name mapping while retaining the source itself for lazy
    pattern materialization.
    """

    def __init__(self) -> None:
        self._layers: list[_SourceLayer] = []
        self._entries: dict[str, Pattern | _SourceEntry] = {}
        self._order: list[str] = []
        self._target_remap: dict[str, str] = {}

    def __iter__(self) -> Iterator[str]:
        return (name for name in self._order if name in self._entries)

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

    def __contains__(self, key: object) -> bool:
        return key in self._entries

    def __getitem__(self, key: str) -> Pattern:
        return self.materialize(key, persist=True)

    def __setitem__(
            self,
            key: str,
            value: Pattern | Callable[[], Pattern],
            ) -> None:
        if key in self._entries:
            raise LibraryError(f'"{key}" already exists in the library. Overwriting is not allowed!')
        pattern = value() if callable(value) else value
        self._entries[key] = pattern
        if key not in self._order:
            self._order.append(key)

    def __delitem__(self, key: str) -> None:
        if key not in self._entries:
            raise KeyError(key)
        del self._entries[key]

    def _merge(self, key_self: str, other: Mapping[str, Pattern], key_other: str) -> None:
        self[key_self] = copy.deepcopy(other[key_other])

    def add_source(
            self,
            source: Mapping[str, Pattern] | ILibraryView,
            *,
            rename_theirs: Callable[[INameView, str], str] | None = _rename_patterns,
            rename_when: Literal['conflict', 'always'] = 'conflict',
            ) -> dict[str, str]:
        """
        Add a source-backed library layer.

        The source must remain open, and its names, hierarchy, and pattern
        contents must remain unchanged for the lifetime of this overlay.

        Args:
            rename_theirs: Function used to choose visible names for imported
                source cells. Its `INameView` argument contains existing and
                previously reserved names, but does not support pattern lookup.
                By default, conflicting single-use names are made unique;
                pass `None` to reject every conflict.
            rename_when: If `'conflict'`, only conflicting names are renamed.
                If `'always'`, every imported source name is passed through
                `rename_theirs`.
        """
        view = source if isinstance(source, ILibraryView) else LibraryView(source)
        source_order = list(view.source_order())
        child_graph = view.child_graph(dangling='include')

        source_to_visible = _plan_source_names(
            self,
            source_order,
            rename_theirs = rename_theirs,
            rename_when = rename_when,
            )
        layer = _SourceLayer(
            library=view,
            source_target_map=dict(source_to_visible),
            child_graph=child_graph,
        )
        layer_index = len(self._layers)
        self._layers.append(layer)

        for source_name, visible_name in source_to_visible.items():
            self._entries[visible_name] = _SourceEntry(layer_index=layer_index, source_name=source_name)
            if visible_name not in self._order:
                self._order.append(visible_name)

        return _source_rename_map(source_to_visible)

    def rename(
            self,
            old_name: str,
            new_name: str,
            move_references: bool = False,
            ) -> OverlayLibrary:
        if old_name not in self._entries:
            raise LibraryError(f'"{old_name}" does not exist in the library.')
        if old_name == new_name:
            return self
        if new_name in self._entries:
            raise LibraryError(f'"{new_name}" already exists in the library.')

        entry = self._entries.pop(old_name)
        self._entries[new_name] = entry

        idx = self._order.index(old_name)
        self._order[idx] = new_name

        if move_references:
            self.move_references(old_name, new_name)
        return self

    def _resolve_target(self, target: str) -> str:
        seen: set[str] = set()
        current = target
        while current in self._target_remap:
            if current in seen:
                raise LibraryError(f'Cycle encountered while resolving target remap for {target!r}')
            seen.add(current)
            current = self._target_remap[current]
        return current

    def _set_target_remap(self, old_target: str, new_target: str) -> None:
        resolved_new = self._resolve_target(new_target)
        if resolved_new == old_target:
            raise LibraryError(f'Ref target remap would create a cycle: {old_target!r} -> {new_target!r}')
        self._target_remap[old_target] = resolved_new
        for key in list(self._target_remap):
            self._target_remap[key] = self._resolve_target(self._target_remap[key])

    def move_references(self, old_target: str, new_target: str) -> OverlayLibrary:
        if old_target == new_target:
            return self
        self._set_target_remap(old_target, new_target)
        for entry in list(self._entries.values()):
            if isinstance(entry, Pattern) and old_target in entry.refs:
                entry.refs[new_target].extend(entry.refs[old_target])
                del entry.refs[old_target]
        return self

    def _effective_target(self, layer: _SourceLayer, target: str) -> str:
        visible = layer.source_target_map.get(target, target)
        return self._resolve_target(visible)

    def _remap_source_pattern(self, layer: _SourceLayer, source_pat: Pattern) -> Pattern:
        def remap(target: str | None) -> str | None:
            return None if target is None else self._effective_target(layer, target)

        if source_pat.refs:
            source_pat.refs = map_targets(source_pat.refs, remap)
        return source_pat

    def materialize(self, name: str, *, persist: bool = True) -> Pattern:
        if name not in self._entries:
            raise KeyError(name)
        entry = self._entries[name]
        if isinstance(entry, Pattern):
            return entry

        layer = self._layers[entry.layer_index]
        source_pat = _materialize_detached_pattern(layer.library, entry.source_name)
        pat = self._remap_source_pattern(layer, source_pat)
        if persist:
            self._entries[name] = pat
        return pat

    def materialize_detached(self, name: str) -> Pattern:
        if name not in self._entries:
            raise KeyError(name)
        entry = self._entries[name]
        if isinstance(entry, Pattern):
            return entry.deepcopy()
        layer = self._layers[entry.layer_index]
        source_pat = _materialize_detached_pattern(layer.library, entry.source_name)
        return self._remap_source_pattern(layer, source_pat)

    def materialize_many_detached(
            self,
            names: Sequence[str],
            ) -> LibraryView:
        ordered_names = tuple(dict.fromkeys(names))
        missing = next((name for name in ordered_names if name not in self._entries), None)
        if missing is not None:
            raise KeyError(missing)

        result: dict[str, Pattern] = {}
        grouped: dict[int, list[tuple[str, str]]] = defaultdict(list)
        for name in ordered_names:
            entry = self._entries[name]
            if isinstance(entry, Pattern):
                result[name] = entry.deepcopy()
            else:
                grouped[entry.layer_index].append((name, entry.source_name))

        for layer_index, cells in grouped.items():
            layer = self._layers[layer_index]
            source_names = [source_name for _name, source_name in cells]
            if isinstance(layer.library, IMaterializable):
                source_patterns = layer.library.materialize_many_detached(source_names)
            else:
                source_patterns = LibraryView({
                    source_name: layer.library[source_name].deepcopy()
                    for source_name in source_names
                    })
            for name, source_name in cells:
                result[name] = self._remap_source_pattern(layer, source_patterns[source_name])

        return LibraryView({name: result[name] for name in ordered_names})

    def child_graph(
            self,
            dangling: dangling_mode_t = 'error',
            ) -> dict[str, set[str]]:
        _validate_dangling_mode(dangling)
        graph: dict[str, set[str]] = {}
        for name in self._order:
            if name not in self._entries:
                continue
            entry = self._entries[name]
            if isinstance(entry, Pattern):
                graph[name] = {child for child, refs in entry.refs.items() if child is not None and refs}
                continue
            layer = self._layers[entry.layer_index]
            children = {self._effective_target(layer, child) for child in layer.child_graph.get(entry.source_name, set())}
            graph[name] = children

        existing = set(graph)
        dangling_refs = set().union(*(children - existing for children in graph.values()))
        if dangling == 'error':
            if dangling_refs:
                raise self._dangling_refs_error(cast('set[str]', dangling_refs), 'building child graph')
            return graph
        if dangling == 'ignore':
            return {name: {child for child in children if child in existing} for name, children in graph.items()}

        for child in dangling_refs:
            graph.setdefault(cast('str', child), set())
        return graph

    def subtree(
            self,
            tops: str | Sequence[str],
            ) -> Self:
        if isinstance(tops, str):
            tops = (tops,)

        graph = self.child_graph(dangling='include')
        keep = self._referenced_patterns_from_graph(graph, tops=tops)
        keep &= set(self)
        keep |= set(tops)

        new = type(self)()
        new._layers = [
            _SourceLayer(
                library=layer.library,
                source_target_map=dict(layer.source_target_map),
                child_graph={name: set(children) for name, children in layer.child_graph.items()},
                )
            for layer in self._layers
            ]
        new._order = [name for name in self._order if name in keep and name in self._entries]
        new._entries = {name: self._entries[name] for name in new._order}
        new._target_remap = dict(self._target_remap)
        return new

    def find_refs_local(
            self,
            name: str,
            parent_graph: dict[str, set[str]] | None = None,
            dangling: dangling_mode_t = 'error',
            ) -> dict[str, list[NDArray[numpy.float64]]]:
        _validate_dangling_mode(dangling)
        instances: dict[str, list[NDArray[numpy.float64]]] = defaultdict(list)
        if parent_graph is None:
            graph_mode = 'ignore' if dangling == 'ignore' else 'include'
            parent_graph = self.parent_graph(dangling=graph_mode)

        if name not in self:
            if name not in parent_graph:
                return instances
            if dangling == 'error':
                raise self._dangling_refs_error({name}, f'finding local refs for {name!r}')
            if dangling == 'ignore':
                return instances

        for parent in parent_graph.get(name, set()):
            pat = self.materialize(parent, persist=False)
            for ref in pat.refs.get(name, []):
                instances[parent].append(ref.as_transforms())
        return instances

    def source_order(self) -> tuple[str, ...]:
        return tuple(name for name in self._order if name in self._entries)

    def borrowed_sources(self) -> tuple[ILibraryView, ...]:
        return tuple(layer.library for layer in self._layers)

    def source_cell(self, name: str) -> tuple[ILibraryView, str] | None:
        entry = self._entries.get(name)
        if not isinstance(entry, _SourceEntry):
            return None
        layer = self._layers[entry.layer_index]
        children = layer.child_graph.get(entry.source_name, set())
        if any(self._effective_target(layer, child) != child for child in children):
            return None
        return layer.library, entry.source_name