from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING

from blends.stack.bridge_paths import index_bridge_paths_for_file
from blends.stack.criteria import (
    is_definition_endpoint,
    is_reference_endpoint,
)
from blends.stack.multi_file_stitcher import MultiFileForwardStitcher
from blends.stack.partial_path_db import PartialPathDb
from blends.stack.selection import edge_list_shadows

if TYPE_CHECKING:
    from blends.stack.forward_stitcher import ForwardStitcherConfig
    from blends.stack.partial_path import (
        FileHandle,
        PartialPath,
        PartialPathDatabase,
        PartialSymbolStack,
    )
    from blends.stack.selection import PartialPathEdge
    from blends.stack.view import StackGraphView


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


@dataclass(frozen=True, slots=True)
class MultiFileResolutionRequest:
    file_handle: FileHandle
    ref_node_index: int
    config: ForwardStitcherConfig | None = None
    include_stats: bool = False


@dataclass(frozen=True, slots=True)
class _MultiFileDefinitionCandidate:
    definition_node_index: int
    file_handle: FileHandle
    edges: tuple[PartialPathEdge, ...]


def _build_source_only_db(
    database: PartialPathDatabase,
    view: StackGraphView,
    source_fh: FileHandle,
) -> PartialPathDb:
    db = PartialPathDb(symbols=view.symbols)
    record = database.get_file(source_fh)
    if record is not None:
        db.add_file(record)
    return db


def _build_stem_index(database: PartialPathDatabase) -> dict[str, FileHandle]:
    return {
        Path(record.file_path).stem: fh
        for fh in database.list_files()
        if (record := database.get_file(fh)) is not None
    }


def _extract_module_stem(
    symbol_stack: PartialSymbolStack,
    db: PartialPathDb,
) -> str | None:
    if not symbol_stack.symbols:
        return None
    return db.symbols.id_to_symbol[symbol_stack.symbols[0].symbol_id]


def resolve_definitions_multi_file(
    view: StackGraphView,
    database: PartialPathDatabase,
    *,
    request: MultiFileResolutionRequest,
    views: dict[FileHandle, StackGraphView] | None = None,
) -> list[tuple[FileHandle, int]] | tuple[list[tuple[FileHandle, int]], MultiFileResolutionStats]:
    if not is_reference_endpoint(view, request.ref_node_index):
        return _empty_multi_result(include_stats=request.include_stats)

    if request.file_handle not in database.list_files():
        return _empty_multi_result(include_stats=request.include_stats)
    if not _ref_record_matches_view(view, request, database):
        return _empty_multi_result(include_stats=request.include_stats)

    db = _build_source_only_db(database, view, request.file_handle)
    all_views: dict[FileHandle, StackGraphView] = (
        views if views is not None else {request.file_handle: view}
    )
    index_bridge_paths_for_file(view, file_handle=request.file_handle, db=db)

    stitcher = MultiFileForwardStitcher.from_references(
        view=view,
        file_handle=request.file_handle,
        db=db,
        reference_nodes=(request.ref_node_index,),
        config=request.config,
        views=all_views,
    )

    pairs, latest_stats, files_loaded = _run_resolution_multi_file(
        view, stitcher, database, all_views, db
    )
    return _finalize_multi_result(
        pairs, latest_stats, files_loaded, include_stats=request.include_stats
    )


def _empty_multi_result(
    *,
    include_stats: bool,
) -> list[tuple[FileHandle, int]] | tuple[list[tuple[FileHandle, int]], MultiFileResolutionStats]:
    if include_stats:
        return [], MultiFileResolutionStats(
            phases_executed=0,
            candidates_considered=0,
            concatenations_succeeded=0,
            complete_paths_found=0,
            cancelled=False,
            cycle_guard_hits=0,
            queue_pruned_count=0,
            files_considered=0,
        )
    return []


def _collect_multi_file_candidates(
    views: dict[FileHandle, StackGraphView],
    complete_paths: tuple[tuple[FileHandle, PartialPath], ...],
    ref_view: StackGraphView,
) -> list[_MultiFileDefinitionCandidate]:
    new_candidates: list[_MultiFileDefinitionCandidate] = []
    for def_fh, path in complete_paths:
        def_view = views.get(def_fh, ref_view)
        if not is_definition_endpoint(def_view, path.end_node_index):
            continue
        new_candidates.append(
            _MultiFileDefinitionCandidate(
                definition_node_index=path.end_node_index,
                file_handle=def_fh,
                edges=path.edges,
            )
        )
    return new_candidates


def _prune_multi_shadowed(
    candidates: list[_MultiFileDefinitionCandidate],
) -> list[_MultiFileDefinitionCandidate]:
    kept: list[_MultiFileDefinitionCandidate] = []
    for candidate in sorted(candidates, key=lambda c: (c.file_handle, c.definition_node_index)):
        is_dominated = any(edge_list_shadows(k.edges, candidate.edges) for k in kept)
        if is_dominated:
            continue
        kept = [k for k in kept if not edge_list_shadows(candidate.edges, k.edges)]
        kept.append(candidate)
    return kept


def _sort_multi_deterministically(
    candidates: list[_MultiFileDefinitionCandidate],
) -> list[_MultiFileDefinitionCandidate]:
    def sort_key(
        candidate: _MultiFileDefinitionCandidate,
    ) -> tuple[tuple[int, ...], tuple[int, ...], int, int]:
        edge_sources = tuple(edge.source_node_index for edge in candidate.edges)
        edge_precedences = tuple(-edge.precedence for edge in candidate.edges)
        return (
            edge_sources,
            edge_precedences,
            candidate.file_handle,
            candidate.definition_node_index,
        )

    return sorted(candidates, key=sort_key)


def _select_multi_file_definitions(
    candidates: list[_MultiFileDefinitionCandidate],
) -> list[tuple[FileHandle, int]]:
    if not candidates:
        return []
    pruned = _prune_multi_shadowed(candidates)
    ordered = _sort_multi_deterministically(pruned)
    seen: set[tuple[int, int]] = set()
    pairs: list[tuple[FileHandle, int]] = []
    for candidate in ordered:
        key = (candidate.file_handle, candidate.definition_node_index)
        if key in seen:
            continue
        seen.add(key)
        pairs.append(key)
    return pairs


def _run_resolution_multi_file(
    view: StackGraphView,
    stitcher: MultiFileForwardStitcher,
    database: PartialPathDatabase,
    all_views: dict[FileHandle, StackGraphView],
    db: PartialPathDb,
) -> tuple[list[tuple[FileHandle, int]], MultiFileResolutionStats | None, int]:
    candidates: list[_MultiFileDefinitionCandidate] = []
    loaded: set[FileHandle] = {stitcher.file_handle}
    stem_index = _build_stem_index(database)
    result = None

    while not stitcher.is_complete():
        result = stitcher.process_next_phase()
        for _fh, sym_stack in result.load_requests.root_requests:
            if (stem := _extract_module_stem(sym_stack, db)) is None:
                continue
            if (target_fh := stem_index.get(stem)) is None or target_fh in loaded:
                continue
            loaded.add(target_fh)
            if (rec := database.get_file(target_fh)) is not None:
                db.add_file(rec)
            if (fh_view := all_views.get(target_fh)) is not None:
                index_bridge_paths_for_file(fh_view, file_handle=target_fh, db=db)
                stitcher.views[target_fh] = fh_view
        candidates.extend(_collect_multi_file_candidates(all_views, result.complete_paths, view))

    pairs = _select_multi_file_definitions(candidates)
    if result is None:
        return pairs, None, len(loaded)
    latest_stats = MultiFileResolutionStats(
        phases_executed=result.stats.phases_executed,
        candidates_considered=result.stats.candidates_considered,
        concatenations_succeeded=result.stats.concatenations_succeeded,
        complete_paths_found=result.stats.complete_paths_found,
        cancelled=result.stats.cancelled,
        cycle_guard_hits=result.stats.cycle_guard_hits,
        queue_pruned_count=result.stats.queue_pruned_count,
        files_considered=len(loaded),
    )
    return pairs, latest_stats, len(loaded)


def _finalize_multi_result(
    pairs: list[tuple[FileHandle, int]],
    latest_stats: MultiFileResolutionStats | None,
    files_considered: int,
    *,
    include_stats: bool,
) -> list[tuple[FileHandle, int]] | tuple[list[tuple[FileHandle, int]], MultiFileResolutionStats]:
    if include_stats:
        if latest_stats is None:
            latest_stats = MultiFileResolutionStats(
                phases_executed=0,
                candidates_considered=0,
                concatenations_succeeded=0,
                complete_paths_found=0,
                cancelled=False,
                cycle_guard_hits=0,
                queue_pruned_count=0,
                files_considered=files_considered,
            )
        return pairs, latest_stats
    return pairs


def _ref_record_matches_view(
    view: StackGraphView,
    request: MultiFileResolutionRequest,
    database: PartialPathDatabase,
) -> bool:
    record = database.get_file(request.file_handle)
    if record is None:
        return False
    if record.file_path and view.file_path:
        return record.file_path == view.file_path
    return True
