"""
Classic source-backed lazy GDSII reader built on the pure-python klamath path.

This module provides the non-Arrow half of Masque's lazy GDS architecture:

- `GdsLibrarySource` scans a GDS stream once to discover library metadata,
  struct order, and child edges without materializing every cell.
- cells are materialized on demand through the classic `gdsii` decoder
  whenever a caller indexes the lazy view
- the source can be wrapped in `PortsLibraryView` or merged through
  `OverlayLibrary`

The public surface intentionally parallels `gdsii.lazy_arrow` closely so that
callers can swap between the classic and Arrow-backed implementations with
minimal changes.
"""
from __future__ import annotations

from dataclasses import dataclass
from typing import IO, TYPE_CHECKING, Any, cast
from collections import defaultdict
import gzip
import io
import logging
import mmap
import pathlib

import klamath
from klamath import records

from . import klamath as gdsii
from .lazy_write import write as write, writefile as writefile
from ..utils import is_gzipped
from ...error import LibraryError
from ...library import (
    ILibraryView,
    IMaterializable,
    PortsLibraryView,
    dangling_mode_t,
)
from ...library.utils import _validate_dangling_mode

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

    import numpy
    from numpy.typing import NDArray

    from ...pattern import Pattern
    from ...ports import Port


logger = logging.getLogger(__name__)


@dataclass
class _SourceHandle:
    """ Owns the underlying stream and any companion file handle for a source. """
    path: pathlib.Path | None
    stream: IO[bytes]
    handle: IO[bytes] | None = None

    def close(self) -> None:
        self.stream.close()
        if self.handle is not None and self.handle is not self.stream:
            self.handle.close()
            self.handle = None


@dataclass(frozen=True)
class _CellScan:
    """ Scan-time metadata for one cell in the source stream. """
    offset: int
    children: set[str]


def _open_source_stream(
        filename: str | pathlib.Path,
        *,
        use_mmap: bool,
        ) -> _SourceHandle:
    path = pathlib.Path(filename).expanduser().resolve()
    if is_gzipped(path):
        if use_mmap:
            logger.info('Asked to mmap a gzipped file, reading into memory instead...')
            with gzip.open(path, mode='rb') as stream:
                data = stream.read()
            return _SourceHandle(path=path, stream=io.BytesIO(data))
        stream = cast('IO[bytes]', gzip.open(path, mode='rb'))      # noqa: SIM115
        return _SourceHandle(path=path, stream=stream)

    if use_mmap:
        handle = path.open(mode='rb', buffering=0)
        mapped = cast('IO[bytes]', mmap.mmap(handle.fileno(), 0, access=mmap.ACCESS_READ))
        return _SourceHandle(path=path, stream=mapped, handle=handle)

    stream = path.open(mode='rb')
    return _SourceHandle(path=path, stream=stream)


def _scan_library(
        stream: IO[bytes],
        ) -> tuple[dict[str, Any], list[str], dict[str, _CellScan]]:
    library_info = gdsii._read_header(stream)
    order: list[str] = []
    cells: dict[str, _CellScan] = {}

    found_struct = records.BGNSTR.skip_past(stream)
    while found_struct:
        name = records.STRNAME.skip_and_read(stream).decode('ASCII')
        offset = stream.tell()
        elements = klamath.library.read_elements(stream)
        children = {
            element.struct_name.decode('ASCII')
            for element in elements
            if isinstance(element, klamath.elements.Reference)
        }
        order.append(name)
        cells[name] = _CellScan(offset=offset, children=children)
        found_struct = records.BGNSTR.skip_past(stream)

    return library_info, order, cells


class GdsLibrarySource(ILibraryView, IMaterializable):
    """
    Read-only library backed by a seekable GDS stream.

    Cells are scanned once up front to discover order and child edges, then
    materialized one at a time through the classic `gdsii.read_elements` path.

    The source owns the stream lifetime, preserves on-disk ordering through
    `source_order()`, and answers graph queries from scan metadata whenever
    possible so callers can inspect hierarchy without forcing a full load.
    """

    def __init__(
            self,
            *,
            source: _SourceHandle,
            library_info: dict[str, Any],
            cell_order: Sequence[str],
            cells: dict[str, _CellScan],
            ) -> None:
        self.path = source.path
        self.library_info = library_info
        self._source = source
        self._cell_order = tuple(cell_order)
        self._cells = cells
        self._cache: dict[str, Pattern] = {}
        self._lookups_in_progress: list[str] = []

    @classmethod
    def from_file(
            cls,
            filename: str | pathlib.Path,
            *,
            use_mmap: bool = True,
            ) -> GdsLibrarySource:
        source = _open_source_stream(filename, use_mmap=use_mmap)
        source.stream.seek(0)
        library_info, cell_order, cells = _scan_library(source.stream)
        return cls(source=source, library_info=library_info, cell_order=cell_order, cells=cells)

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

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

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

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

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

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

        if name not in self._cells:
            raise KeyError(name)

        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.\n'
                'If you are lazy-loading a file, try a non-lazy load and check for reference cycles.'
                )

        self._lookups_in_progress.append(name)
        try:
            self._source.stream.seek(self._cells[name].offset)
            pat = gdsii.read_elements(self._source.stream, raw_mode=True)
        finally:
            self._lookups_in_progress.pop()

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

    def _raw_children(self, name: str) -> set[str]:
        return set(self._cells[name].children)

    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._cell_order:
            if name in self._cache:
                graph[name] = {child for child, refs in self._cache[name].refs.items() if child is not None and refs}
            else:
                graph[name] = self._raw_children(name)

        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],
            ) -> ILibraryView:
        return super().subtree(tops)

    def with_ports_from_data(
            self,
            *,
            layers: Sequence[tuple[int, int] | int],
            max_depth: int = 0,
            skip_subcells: bool = True,
            ports: Mapping[str, Mapping[str, Port]] | None = None,
            replace: bool = False,
            ) -> PortsLibraryView:
        return PortsLibraryView(
            self,
            layers=layers,
            max_depth=max_depth,
            skip_subcells=skip_subcells,
            ports=ports,
            replace=replace,
        )

    def with_port_overrides(
            self,
            ports: Mapping[str, Mapping[str, Port]],
            *,
            replace: bool = False,
            ) -> PortsLibraryView:
        return PortsLibraryView(
            self,
            ports=ports,
            replace=replace,
        )

    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()):
            if parent in self._cache:
                for ref in self._cache[parent].refs.get(name, []):
                    instances[parent].append(ref.as_transforms())
                continue
            pat = self.materialize(parent, persist=False)
            for ref in pat.refs.get(name, []):
                instances[parent].append(ref.as_transforms())
        return instances

    def close(self) -> None:
        self._source.close()

    def __enter__(self) -> GdsLibrarySource:
        return self

    def __exit__(self, *_args: object) -> None:
        self.close()


def read(
        stream: IO[bytes],
        ) -> tuple[GdsLibrarySource, dict[str, Any]]:
    source = _SourceHandle(path=None, stream=stream)
    stream.seek(0)
    library_info, cell_order, cells = _scan_library(stream)
    lib = GdsLibrarySource(source=source, library_info=library_info, cell_order=cell_order, cells=cells)
    return lib, library_info


def readfile(
        filename: str | pathlib.Path,
        *,
        use_mmap: bool = True,
        ) -> tuple[GdsLibrarySource, dict[str, Any]]:
    lib = GdsLibrarySource.from_file(filename, use_mmap=use_mmap)
    return lib, lib.library_info