from __future__ import annotations

from dataclasses import dataclass, field
from typing import TYPE_CHECKING

from blends.stack.criteria import (
    is_complete_definition_path,
    is_reference_endpoint,
)
from blends.stack.partial_path.errors import PartialPathResolutionError
from blends.stack.partial_path.stack_conditions import create_seed_path
from blends.stack.stitcher_utils import (
    _path_priority,
    _should_cancel,
    _stack_state_from_partial,
)

if TYPE_CHECKING:
    from collections.abc import Callable

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

from blends.stack.stacks import StackState

StateKey = tuple[int, int, "PartialSymbolStack", "PartialScopeStack"]


@dataclass(frozen=True, slots=True)
class ForwardStitcherConfig:
    max_work_per_phase: int | None = None
    max_phases: int | None = None
    max_total_work: int | None = None
    max_queue_size: int | None = None
    enable_dedupe: bool = True
    enable_cycle_guard: bool = True
    is_cancelled: Callable[[], bool] | None = None
    deadline: float | None = None


@dataclass(slots=True)
class ForwardStitcherStats:
    phases_executed: int = 0
    candidates_considered: int = 0
    concatenations_succeeded: int = 0
    complete_paths_found: int = 0
    cancelled: bool = False
    cycle_guard_hits: int = 0
    queue_pruned_count: int = 0


@dataclass(frozen=True, slots=True)
class ForwardStitcherStatsSnapshot:
    phases_executed: int
    candidates_considered: int
    concatenations_succeeded: int
    complete_paths_found: int
    cancelled: bool
    cycle_guard_hits: int
    queue_pruned_count: int


@dataclass(frozen=True, slots=True)
class ForwardStitcherLoadRequests:
    node_requests: tuple[tuple[FileHandle, int], ...] = ()
    root_requests: tuple[tuple[FileHandle, PartialSymbolStack], ...] = ()


@dataclass(frozen=True, slots=True)
class ForwardStitcherPhaseResult:
    complete_paths: tuple[PartialPath, ...]
    stats: ForwardStitcherStatsSnapshot


@dataclass(slots=True)
class ForwardStitcher:
    view: StackGraphView
    file_handle: FileHandle
    db: PartialPathDb
    config: ForwardStitcherConfig = field(default_factory=ForwardStitcherConfig)

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

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

        stitcher = cls(
            view=view,
            file_handle=file_handle,
            db=db,
            config=config or ForwardStitcherConfig(),
        )
        stitcher._current_queue = seed_paths
        stitcher._is_complete = len(seed_paths) == 0

        if stitcher.config.enable_cycle_guard:
            for seed_path in seed_paths:
                state_key: StateKey = (
                    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) -> ForwardStitcherPhaseResult:
        if self._is_complete or self._phase_limit_reached():
            self._is_complete = True
            return self._build_phase_result(())

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

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

        if self._phase_limit_reached():
            self._is_complete = True

        return self._build_phase_result(tuple(complete_paths))

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

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

    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,
        complete_paths: tuple[PartialPath, ...],
    ) -> ForwardStitcherPhaseResult:
        return ForwardStitcherPhaseResult(
            complete_paths=complete_paths,
            stats=_snapshot_stats(self._stats),
        )

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

        for index, 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_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=_path_priority)
            kept = next_queue[: self.config.max_queue_size]
            pruned = next_queue[self.config.max_queue_size :]
            if self.config.enable_cycle_guard:
                for pruned_path in pruned:
                    state_key: StateKey = (
                        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

        return next_queue

    def _extend_with_candidates(
        self,
        current_path: PartialPath,
        next_queue: list[PartialPath],
        complete_paths: list[PartialPath],
    ) -> int:
        candidates = self.db.get_candidates(
            file_handle=self.file_handle,
            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)

        for candidate_id in candidates:
            candidate_path = self.db.get_path(candidate_id)
            if candidate_path is None:
                continue

            try:
                extended = current_path.concatenate(self.view, candidate_path)
            except PartialPathResolutionError:
                continue

            if self.config.enable_cycle_guard:
                state_key: StateKey = (
                    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

            if _is_complete_path(self.view, extended):
                complete_paths.append(extended)
                self._stats.complete_paths_found += 1
            else:
                next_queue.append(extended)

        return len(candidates)


def _snapshot_stats(stats: ForwardStitcherStats) -> ForwardStitcherStatsSnapshot:
    return ForwardStitcherStatsSnapshot(
        phases_executed=stats.phases_executed,
        candidates_considered=stats.candidates_considered,
        concatenations_succeeded=stats.concatenations_succeeded,
        complete_paths_found=stats.complete_paths_found,
        cancelled=stats.cancelled,
        cycle_guard_hits=stats.cycle_guard_hits,
        queue_pruned_count=stats.queue_pruned_count,
    )


def _is_complete_path(
    view: StackGraphView,
    path: PartialPath,
) -> bool:
    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,
        )
    return is_complete_definition_path(
        view,
        start_node_index=path.start_node_index,
        end_node_index=path.end_node_index,
        start_state=start_state,
        end_state=end_state,
    )
