from __future__ import annotations

from dataclasses import dataclass, field
from typing import TYPE_CHECKING

from blends.stack.criteria import (
    is_definition_endpoint,
    is_exported_scope_endpoint,
    is_reference_endpoint,
)
from blends.stack.partial_path import (
    PartialPathDatabase,
    PartialPathIndexRequest,
    PartialPathLimits,
    index_partial_paths_for_file,
)
from blends.stack.selection import (
    DefinitionCandidate,
    prune_shadowed_candidates,
    sort_candidates_deterministically,
)
from blends.stack.view import StackGraphView

if TYPE_CHECKING:
    from blends.models import NId, SyntaxGraph
    from blends.stack.partial_path import PartialPath, PartialPathFileRecord

_FILE_HANDLE: int = 1
_MAX_INDEX_WORK: int = 50_000


@dataclass(frozen=True, slots=True)
class DefinitionResolver:
    view: StackGraphView
    _cache: dict[NId, NId | None] = field(default_factory=dict)

    @classmethod
    def from_syntax_graph(cls, graph: SyntaxGraph, file_path: str) -> DefinitionResolver:
        view = StackGraphView.from_syntax_graph(graph, path=file_path)
        database = PartialPathDatabase()
        index_partial_paths_for_file(
            graph,
            request=PartialPathIndexRequest(
                file_path=file_path,
                file_handle=_FILE_HANDLE,
                limits=PartialPathLimits(max_work_per_phase=_MAX_INDEX_WORK),
            ),
            database=database,
        )
        record = database.get_file(_FILE_HANDLE)
        cache: dict[NId, NId | None] = {}
        if record is not None:
            cache = _build_cache(view, record)
        return cls(view=view, _cache=cache)

    def resolve(self, ref_n_id: NId) -> NId | None:
        return self._cache.get(ref_n_id)


def _is_complete_direct_path(path: PartialPath) -> bool:
    return (
        path.symbol_stack_precondition.can_match_empty()
        and path.scope_stack_precondition.can_match_empty()
        and path.symbol_stack_postcondition.can_match_empty()
        and (
            path.scope_stack_postcondition.can_match_empty()
            or path.scope_stack_postcondition.has_variable()
        )
    )


def _is_indexable_scope_def_path(path: PartialPath) -> bool:
    pre_sym = path.symbol_stack_precondition
    if len(pre_sym.symbols) != 1:
        return False
    scopes = pre_sym.symbols[0].scopes
    if scopes is not None and not scopes.has_variable():
        return False
    if not path.symbol_stack_postcondition.can_match_empty():
        return False
    post_sco = path.scope_stack_postcondition
    return post_sco.can_match_empty() or post_sco.has_variable()


def _scan_partial_paths(
    view: StackGraphView,
    paths: tuple[PartialPath, ...],
) -> tuple[dict[int, list[DefinitionCandidate]], dict[int, list[int]]]:
    direct_candidates: dict[int, list[DefinitionCandidate]] = {}
    scope_def_index: dict[int, list[int]] = {}
    for path in paths:
        s = path.start_node_index
        e = path.end_node_index
        if is_exported_scope_endpoint(view, s) and is_definition_endpoint(view, e):
            if _is_indexable_scope_def_path(path):
                sym_id = path.symbol_stack_precondition.symbols[0].symbol_id
                scope_def_index.setdefault(sym_id, []).append(e)
            continue
        is_direct = is_reference_endpoint(view, s) and is_definition_endpoint(view, e)
        if is_direct and _is_complete_direct_path(path):
            direct_candidates.setdefault(s, []).append(
                DefinitionCandidate(definition_node_index=e, edges=path.edges)
            )
    return direct_candidates, scope_def_index


def _resolve_direct_candidates(
    view: StackGraphView,
    direct_candidates: dict[int, list[DefinitionCandidate]],
) -> tuple[dict[NId, NId | None], set[int]]:
    cache: dict[NId, NId | None] = {}
    resolved_ref_indices: set[int] = set()
    for ref_idx, candidates in direct_candidates.items():
        pruned = prune_shadowed_candidates(candidates)
        ordered = sort_candidates_deterministically(pruned)
        if ordered:
            cache[view.index_to_nid[ref_idx]] = view.index_to_nid[ordered[0].definition_node_index]
            resolved_ref_indices.add(ref_idx)
    return cache, resolved_ref_indices


def _resolve_cross_scope(
    view: StackGraphView,
    scope_def_index: dict[int, list[int]],
    resolved_ref_indices: set[int],
    cache: dict[NId, NId | None],
) -> None:
    for i, is_ref in enumerate(view.node_is_reference):
        if not is_ref or i in resolved_ref_indices:
            continue
        if not is_reference_endpoint(view, i):
            continue
        ref_sym_id = view.symbol_id_at(i)
        if ref_sym_id is None:
            continue
        def_indices = scope_def_index.get(ref_sym_id)
        if not def_indices:
            continue
        cache[view.index_to_nid[i]] = view.index_to_nid[def_indices[0]]


def _build_cache(
    view: StackGraphView,
    record: PartialPathFileRecord,
) -> dict[NId, NId | None]:
    direct_candidates, scope_def_index = _scan_partial_paths(view, record.paths)
    cache, resolved_ref_indices = _resolve_direct_candidates(view, direct_candidates)
    if scope_def_index:
        _resolve_cross_scope(view, scope_def_index, resolved_ref_indices, cache)
    return cache
