from __future__ import annotations

from typing import TYPE_CHECKING

from blends.stack.partial_path.partial_path import PartialPath
from blends.stack.partial_path.partial_stacks import (
    PartialScopedSymbol,
    PartialScopeStack,
    PartialSymbolStack,
)
from blends.stack.partial_path.variables import ScopeStackVariable, SymbolStackVariable

if TYPE_CHECKING:
    from blends.stack.view import StackGraphView, SymbolInterner

_BRIDGE_SYMBOL_VAR_ID = 1
_BRIDGE_SCOPE_ATTACHED_VAR_ID = 2
_BRIDGE_SCOPE_STACK_VAR_ID = 1


def _make_bridge_partial_path(
    start_node_index: int,
    pre_bound_symbol_id: int,
    post_bound_symbol_id: int,
    module_symbol_ids: tuple[int, ...],
) -> PartialPath:
    symbol_var = SymbolStackVariable(var_id=_BRIDGE_SYMBOL_VAR_ID)
    scope_attached_var = ScopeStackVariable(var_id=_BRIDGE_SCOPE_ATTACHED_VAR_ID)
    scope_stack_var = ScopeStackVariable(var_id=_BRIDGE_SCOPE_STACK_VAR_ID)

    pre_symbol_stack = PartialSymbolStack(
        symbols=(
            PartialScopedSymbol(
                symbol_id=pre_bound_symbol_id,
                scopes=PartialScopeStack.from_variable(scope_attached_var),
            ),
        ),
        variable=symbol_var,
    )
    post_module_symbols = tuple(
        PartialScopedSymbol(symbol_id=sid, scopes=None) for sid in module_symbol_ids
    )
    post_symbol_stack = PartialSymbolStack(
        symbols=(
            *post_module_symbols,
            PartialScopedSymbol(
                symbol_id=post_bound_symbol_id,
                scopes=PartialScopeStack.from_variable(scope_attached_var),
            ),
        ),
        variable=symbol_var,
    )
    pre_scope_stack = PartialScopeStack.from_variable(scope_stack_var)
    post_scope_stack = PartialScopeStack.from_variable(scope_attached_var)

    return PartialPath(
        start_node_index=start_node_index,
        end_node_index=0,
        symbol_stack_precondition=pre_symbol_stack,
        symbol_stack_postcondition=post_symbol_stack,
        scope_stack_precondition=pre_scope_stack,
        scope_stack_postcondition=post_scope_stack,
        edges=(),
    )


def _module_path_to_symbol_ids(module_path: str, symbols: SymbolInterner) -> tuple[int, ...] | None:
    parts = [p for p in module_path.split(".") if p]
    if not parts:
        return None
    segments_with_dots: list[str] = []
    for i, part in enumerate(parts):
        segments_with_dots.append(part)
        if i < len(parts) - 1:
            segments_with_dots.append(".")
    ids: list[int] = []
    for seg in segments_with_dots:
        sid = symbols.symbol_to_id.get(seg)
        if sid is None:
            return None
        ids.append(sid)
    return tuple(ids)


def _resolve_relative_stem(
    current_module_path: str,
    relative_level: int,
    trailing_stem: str | None,
    *,
    is_package: bool = False,
) -> str | None:
    parts = [p for p in current_module_path.split(".") if p]
    effective_level = relative_level - 1 if is_package else relative_level
    if effective_level > len(parts):
        return None
    package_parts = parts[:-effective_level] if effective_level > 0 else parts
    if trailing_stem:
        return ".".join([*package_parts, trailing_stem]) if package_parts else trailing_stem
    return ".".join(package_parts) if package_parts else None


def _effective_module_stem(view: StackGraphView, node_index: int) -> str | None:
    stem = view.node_import_module_stem[node_index]
    level = view.node_import_relative_level[node_index]
    if level > 0:
        is_package = view.file_path.endswith("/__init__.py")
        return _resolve_relative_stem(view.module_path, level, stem or None, is_package=is_package)
    return stem or None


def collect_star_import_stems(view: StackGraphView) -> dict[str, int]:
    result: dict[str, int] = {}
    for node_index, is_bridge in enumerate(view.node_is_import_binding):
        if not is_bridge:
            continue
        symbol_id = view.node_symbol_id[node_index]
        if symbol_id is None:
            continue
        if view.symbols.id_to_symbol[symbol_id] != "*":
            continue
        stem = _effective_module_stem(view, node_index)
        if stem:
            result[stem] = node_index
    return result


def _find_parent_scope(view: StackGraphView, node_index: int) -> int | None:
    for idx, edges in enumerate(view.outgoing):
        for sink, _ in edges:
            if sink == node_index:
                return idx
    return None


def expand_star_bridges(
    consumer_view: StackGraphView,
    star_node_index: int,
    target_view: StackGraphView,
) -> tuple[PartialPath, ...]:
    module_stem = _effective_module_stem(consumer_view, star_node_index)
    if not module_stem:
        return ()
    module_symbol_ids = _module_path_to_symbol_ids(module_stem, consumer_view.symbols)
    if module_symbol_ids is None:
        return ()
    scope_index = _find_parent_scope(consumer_view, star_node_index)
    if scope_index is None:
        return ()
    bridges: list[PartialPath] = []
    for def_node_index in target_view.exported_definition_nodes:
        symbol_id = target_view.node_symbol_id[def_node_index]
        if symbol_id is None:
            continue
        bridges.append(
            _make_bridge_partial_path(scope_index, symbol_id, symbol_id, module_symbol_ids)
        )
    return tuple(bridges)


def compute_bridge_paths_from_view(view: StackGraphView) -> tuple[PartialPath, ...]:
    bridges: list[PartialPath] = []
    for node_index, is_bridge in enumerate(view.node_is_import_binding):
        if not is_bridge:
            continue
        module_stem = _effective_module_stem(view, node_index)
        if not module_stem:
            continue
        module_symbol_ids = _module_path_to_symbol_ids(module_stem, view.symbols)
        if module_symbol_ids is None:
            continue
        pre_bound_symbol_id = view.node_symbol_id[node_index]
        if pre_bound_symbol_id is None:
            continue
        if view.symbols.id_to_symbol[pre_bound_symbol_id] == "*":
            continue
        import_source_symbol = view.node_import_source_symbol[node_index]
        source_symbol_str = (
            import_source_symbol
            if import_source_symbol is not None
            else view.symbols.id_to_symbol[pre_bound_symbol_id]
        )
        post_bound_symbol_id = view.symbols.symbol_to_id.get(source_symbol_str)
        if post_bound_symbol_id is None:
            post_bound_symbol_id = view.symbols.intern(source_symbol_str)
        bridges.append(
            _make_bridge_partial_path(
                node_index, pre_bound_symbol_id, post_bound_symbol_id, module_symbol_ids
            )
        )
    return tuple(bridges)
