from __future__ import annotations

from dataclasses import dataclass, field
from typing import TYPE_CHECKING

from blends.stack.criteria import (
    is_definition_endpoint,
    is_import_binding_endpoint,
    is_reference_endpoint,
)
from blends.stack.forward_stitcher import (
    ForwardStitcherConfig,
    ForwardStitcherLoadRequests,
    ForwardStitcherStats,
    ForwardStitcherStatsSnapshot,
    _snapshot_stats,
)
from blends.stack.partial_path.errors import PartialPathResolutionError
from blends.stack.partial_path.stack_conditions import create_seed_path
from blends.stack.stacks import StackState
from blends.stack.stitcher_utils import (
    _path_priority,
    _should_cancel,
    _stack_state_from_partial,
)

if TYPE_CHECKING:
    from blends.stack.partial_path import (
        FileHandle,
        PartialPath,
        PartialScopeStack,
        PartialSymbolStack,
    )
    from blends.stack.partial_path_db import PartialPathDb
    from blends.stack.view import StackGraphView


MultiFileStateKey = tuple["FileHandle", int, int, "PartialSymbolStack", "PartialScopeStack"]


@dataclass(frozen=True, slots=True)
class MultiFileForwardStitcherPhaseResult:
    complete_paths: tuple[tuple[FileHandle, PartialPath], ...]
    load_requests: ForwardStitcherLoadRequests
    stats: ForwardStitcherStatsSnapshot


@dataclass(slots=True)
class MultiFileForwardStitcher:
    view: StackGraphView
    file_handle: FileHandle
    db: PartialPathDb
    views: dict[FileHandle, StackGraphView] = field(default_factory=dict)
    config: ForwardStitcherConfig = field(default_factory=ForwardStitcherConfig)

    _current_queue: list[tuple[FileHandle, PartialPath]] = field(default_factory=list)
    _previous_phase_paths: tuple[tuple[FileHandle, PartialPath], ...] = ()
    _stats: ForwardStitcherStats = field(default_factory=ForwardStitcherStats)
    _is_complete: bool = False
    _seen_states: set[MultiFileStateKey] = field(default_factory=set)

    @classmethod
    def from_references(  # noqa: PLR0913
        cls,
        *,
        view: StackGraphView,
        file_handle: FileHandle,
        db: PartialPathDb,
        reference_nodes: tuple[int, ...],
        config: ForwardStitcherConfig | None = None,
        views: dict[FileHandle, StackGraphView] | None = None,
    ) -> MultiFileForwardStitcher:
        seed_paths: list[tuple[FileHandle, PartialPath]] = []
        for node_index in sorted(reference_nodes):
            if not is_reference_endpoint(view, node_index):
                continue
            seed_paths.append((file_handle, create_seed_path(view, node_index)))

        stitcher = cls(
            view=view,
            file_handle=file_handle,
            db=db,
            views=views if views is not None else {file_handle: view},
            config=config or ForwardStitcherConfig(),
        )
        stitcher._current_queue = seed_paths
        stitcher._is_complete = len(seed_paths) == 0

        if stitcher.config.enable_cycle_guard:
            for current_fh, seed_path in seed_paths:
                state_key: MultiFileStateKey = (
                    current_fh,
                    seed_path.start_node_index,
                    seed_path.end_node_index,
                    seed_path.symbol_stack_postcondition,
                    seed_path.scope_stack_postcondition,
                )
                stitcher._seen_states.add(state_key)

        return stitcher

    def process_next_phase(self) -> MultiFileForwardStitcherPhaseResult:
        if self._is_complete or self._phase_limit_reached():
            self._is_complete = True
            return self._build_phase_result(ForwardStitcherLoadRequests(), ())

        self._stats.phases_executed += 1
        self._previous_phase_paths = tuple(self._current_queue)

        complete_paths: list[tuple[FileHandle, PartialPath]] = []
        self._current_queue = self._expand_phase(self._current_queue, complete_paths)
        if len(self._current_queue) == 0:
            self._is_complete = True

        load_requests = _collect_load_requests(self._current_queue)
        if self._phase_limit_reached():
            self._is_complete = True

        return self._build_phase_result(load_requests, tuple(complete_paths))

    def previous_phase_paths(self) -> tuple[tuple[FileHandle, PartialPath], ...]:
        return self._previous_phase_paths

    def is_complete(self) -> bool:
        return self._is_complete

    def current_load_requests(self) -> ForwardStitcherLoadRequests:
        return _collect_load_requests(self._current_queue)

    def _phase_limit_reached(self) -> bool:
        return (
            self.config.max_phases is not None
            and self._stats.phases_executed >= self.config.max_phases
        )

    def _build_phase_result(
        self,
        load_requests: ForwardStitcherLoadRequests,
        complete_paths: tuple[tuple[FileHandle, PartialPath], ...],
    ) -> MultiFileForwardStitcherPhaseResult:
        return MultiFileForwardStitcherPhaseResult(
            complete_paths=complete_paths,
            load_requests=load_requests,
            stats=_snapshot_stats(self._stats),
        )

    def _expand_phase(
        self,
        current_queue: list[tuple[FileHandle, PartialPath]],
        complete_paths: list[tuple[FileHandle, PartialPath]],
    ) -> list[tuple[FileHandle, PartialPath]]:
        next_queue: list[tuple[FileHandle, PartialPath]] = []
        work_performed_phase = 0

        for index, (current_fh, current_path) in enumerate(current_queue):
            if _should_cancel(self.config):
                self._stats.cancelled = True
                self._is_complete = True
                return []
            if (
                self.config.max_work_per_phase is not None
                and work_performed_phase >= self.config.max_work_per_phase
            ):
                next_queue.extend(current_queue[index:])
                break
            if (
                self.config.max_total_work is not None
                and self._stats.candidates_considered >= self.config.max_total_work
            ):
                self._is_complete = True
                break
            candidates_processed = self._extend_with_candidates(
                current_fh, current_path, next_queue, complete_paths
            )
            work_performed_phase += candidates_processed

        if self.config.max_queue_size is not None and len(next_queue) > self.config.max_queue_size:
            pruned_count = len(next_queue) - self.config.max_queue_size
            self._stats.queue_pruned_count += pruned_count
            next_queue.sort(key=lambda item: _path_priority(item[1]))
            kept_pairs = next_queue[: self.config.max_queue_size]
            pruned_pairs = next_queue[self.config.max_queue_size :]
            if self.config.enable_cycle_guard:
                for pruned_fh, pruned_path in pruned_pairs:
                    state_key: MultiFileStateKey = (
                        pruned_fh,
                        pruned_path.start_node_index,
                        pruned_path.end_node_index,
                        pruned_path.symbol_stack_postcondition,
                        pruned_path.scope_stack_postcondition,
                    )
                    self._seen_states.discard(state_key)
            next_queue = kept_pairs

        return next_queue

    def _extend_with_candidates(
        self,
        current_fh: FileHandle,
        current_path: PartialPath,
        next_queue: list[tuple[FileHandle, PartialPath]],
        complete_paths: list[tuple[FileHandle, PartialPath]],
    ) -> int:
        candidates = self.db.get_candidates(
            file_handle=current_fh,
            end_node_index=current_path.end_node_index,
            symbol_stack=current_path.symbol_stack_postcondition,
            scope_stack=current_path.scope_stack_postcondition,
        )
        self._stats.candidates_considered += len(candidates)
        current_view = self.views.get(current_fh, self.view)
        for candidate_id in candidates:
            candidate_path = self.db.get_path(candidate_id)
            if candidate_path is None:
                continue
            candidate_fh: FileHandle = candidate_id[0]
            try:
                extended = current_path.concatenate(current_view, candidate_path)
            except PartialPathResolutionError:
                continue
            if self.config.enable_cycle_guard:
                state_key: MultiFileStateKey = (
                    candidate_fh,
                    extended.start_node_index,
                    extended.end_node_index,
                    extended.symbol_stack_postcondition,
                    extended.scope_stack_postcondition,
                )
                if state_key in self._seen_states:
                    self._stats.cycle_guard_hits += 1
                    continue
                self._seen_states.add(state_key)
            self._stats.concatenations_succeeded += 1
            candidate_view = self.views.get(candidate_fh, self.view)
            self._route_extended(candidate_fh, extended, candidate_view, next_queue, complete_paths)
        return len(candidates)

    def _route_extended(
        self,
        candidate_fh: FileHandle,
        extended: PartialPath,
        candidate_view: StackGraphView,
        next_queue: list[tuple[FileHandle, PartialPath]],
        complete_paths: list[tuple[FileHandle, PartialPath]],
    ) -> None:
        if not _is_complete_path(candidate_view, extended):
            next_queue.append((candidate_fh, extended))
            return
        complete_paths.append((candidate_fh, extended))
        self._stats.complete_paths_found += 1
        if is_import_binding_endpoint(candidate_view, extended.end_node_index):
            next_queue.append((candidate_fh, extended))


def _collect_load_requests(
    paths: list[tuple[FileHandle, PartialPath]],
) -> ForwardStitcherLoadRequests:
    node_requests: list[tuple[FileHandle, int]] = []
    root_requests: list[tuple[FileHandle, PartialSymbolStack]] = []
    seen_nodes: set[tuple[FileHandle, int]] = set()
    seen_roots: set[tuple[FileHandle, PartialSymbolStack]] = set()

    for current_fh, path in paths:
        if path.end_node_index == 0:
            root_key = (current_fh, path.symbol_stack_postcondition)
            if root_key in seen_roots:
                continue
            seen_roots.add(root_key)
            root_requests.append(root_key)
        else:
            node_key = (current_fh, path.end_node_index)
            if node_key in seen_nodes:
                continue
            seen_nodes.add(node_key)
            node_requests.append(node_key)

    return ForwardStitcherLoadRequests(
        node_requests=tuple(node_requests),
        root_requests=tuple(root_requests),
    )


def _is_complete_path(
    def_view: StackGraphView,
    path: PartialPath,
) -> bool:
    if not is_definition_endpoint(def_view, path.end_node_index):
        return False
    start_state = _stack_state_from_partial(
        path.symbol_stack_precondition, path.scope_stack_precondition
    )
    end_state = _stack_state_from_partial(
        path.symbol_stack_postcondition, path.scope_stack_postcondition
    )
    if end_state.scope_stack is not None and path.scope_stack_postcondition.has_variable():
        end_state = StackState(
            symbol_stack=end_state.symbol_stack,
            scope_stack=None,
        )
    if start_state.symbol_stack is not None:
        return False
    if start_state.scope_stack is not None:
        return False
    if end_state.symbol_stack is not None:
        return False
    return end_state.scope_stack is None
