from __future__ import annotations

from collections.abc import Callable
from typing import Any


def _store_and_attach_index_candidate_cache(
    *,
    payload: dict[str, Any],
    index_candidate_cache_path: Any,
    index_candidate_cache_key: str,
    index_candidate_cache_ttl_seconds: int,
    index_candidate_cache_required_meta: dict[str, Any],
    index_candidate_cache: dict[str, Any],
    query: str,
    clone_index_candidate_payload_fn: Callable[..., dict[str, Any]],
    store_cached_index_candidates_fn: Callable[..., bool],
    attach_index_candidate_cache_info_fn: Callable[..., dict[str, Any]],
) -> dict[str, Any]:
    index_candidate_cache["store_written"] = bool(
        store_cached_index_candidates_fn(
            cache_path=index_candidate_cache_path,
            key=index_candidate_cache_key,
            payload=clone_index_candidate_payload_fn(payload),
            meta={
                **index_candidate_cache_required_meta,
                "query": query,
                "ttl_seconds": int(index_candidate_cache_ttl_seconds),
                "trust_class": "exact",
            },
        )
    )
    return attach_index_candidate_cache_info_fn(
        payload=payload,
        cache_info=index_candidate_cache,
    )


def finalize_index_stage_output(
    *,
    repo: str,
    root: str,
    terms: list[str],
    memory_paths: list[str],
    index_hash: str,
    index_data: dict[str, Any],
    cache_info: dict[str, Any],
    requested_ranker: str,
    selected_ranker: str,
    ranker_fallbacks: list[str],
    corpus_size: int,
    min_score_used: int,
    fusion_mode: str,
    hybrid_re2_rrf_k: int,
    top_k_files: int,
    candidate_relative_threshold: float,
    chunk_top_k: int,
    chunk_per_file_limit: int,
    chunk_token_budget: int,
    chunk_disclosure: str,
    candidates: list[dict[str, Any]],
    candidate_chunks: list[dict[str, Any]],
    chunk_metrics: dict[str, Any],
    exact_search_payload: dict[str, Any],
    docs_payload: dict[str, Any],
    worktree_prior: dict[str, Any],
    parallel_payload: dict[str, Any],
    prior_payload: dict[str, Any],
    graph_lookup_payload: dict[str, Any],
    cochange_payload: dict[str, Any],
    scip_payload: dict[str, Any],
    embeddings_payload: dict[str, Any],
    feedback_payload: dict[str, Any],
    multi_channel_fusion_payload: dict[str, Any],
    retrieval_refinement_payload: dict[str, Any],
    second_pass_payload: dict[str, Any],
    refine_pass_payload: dict[str, Any],
    chunk_semantic_rerank_payload: dict[str, Any],
    topological_shield_payload: dict[str, Any],
    chunk_guard_payload: dict[str, Any],
    adaptive_router_payload: dict[str, Any],
    policy_name: str,
    policy_version: str,
    timings_ms: dict[str, Any],
    benchmark_filter_payload: dict[str, Any],
    index_candidate_cache_path: str,
    index_candidate_cache_key: str,
    index_candidate_cache_ttl_seconds: int,
    index_candidate_cache_required_meta: dict[str, Any],
    index_candidate_cache: dict[str, Any],
    query: str,
    build_index_stage_output_fn: Callable[..., dict[str, Any]],
    clone_index_candidate_payload_fn: Callable[..., dict[str, Any]],
    store_cached_index_candidates_fn: Callable[..., bool],
    attach_index_candidate_cache_info_fn: Callable[..., dict[str, Any]],
) -> dict[str, Any]:
    payload = build_index_stage_output_fn(
        repo=repo,
        root=root,
        terms=terms,
        memory_paths=memory_paths,
        index_hash=index_hash,
        index_data=index_data,
        cache_info=cache_info,
        requested_ranker=requested_ranker,
        selected_ranker=selected_ranker,
        ranker_fallbacks=ranker_fallbacks,
        corpus_size=corpus_size,
        min_score_used=min_score_used,
        fusion_mode=fusion_mode,
        hybrid_re2_rrf_k=int(hybrid_re2_rrf_k),
        top_k_files=int(top_k_files),
        candidate_relative_threshold=float(candidate_relative_threshold),
        chunk_top_k=int(chunk_top_k),
        chunk_per_file_limit=int(chunk_per_file_limit),
        chunk_token_budget=int(chunk_token_budget),
        chunk_disclosure=str(chunk_disclosure),
        candidates=candidates,
        candidate_chunks=candidate_chunks,
        chunk_metrics=chunk_metrics,
        exact_search_payload=exact_search_payload,
        docs_payload=docs_payload,
        worktree_prior=worktree_prior,
        parallel_payload=parallel_payload,
        prior_payload=prior_payload,
        graph_lookup_payload=graph_lookup_payload,
        cochange_payload=cochange_payload,
        scip_payload=scip_payload,
        embeddings_payload=embeddings_payload,
        feedback_payload=feedback_payload,
        multi_channel_fusion_payload=multi_channel_fusion_payload,
        second_pass_payload=second_pass_payload,
        refine_pass_payload=refine_pass_payload,
        chunk_semantic_rerank_payload=chunk_semantic_rerank_payload,
        topological_shield_payload=topological_shield_payload,
        chunk_guard_payload=chunk_guard_payload,
        adaptive_router_payload=adaptive_router_payload,
        policy_name=str(policy_name),
        policy_version=str(policy_version),
        timings_ms=timings_ms,
    )
    if bool(retrieval_refinement_payload.get("enabled")):
        payload["retrieval_refinement"] = dict(retrieval_refinement_payload)
    if benchmark_filter_payload.get("requested"):
        payload["benchmark_filters"] = benchmark_filter_payload
    return _store_and_attach_index_candidate_cache(
        payload=payload,
        index_candidate_cache_path=index_candidate_cache_path,
        index_candidate_cache_key=index_candidate_cache_key,
        index_candidate_cache_ttl_seconds=index_candidate_cache_ttl_seconds,
        index_candidate_cache_required_meta=index_candidate_cache_required_meta,
        index_candidate_cache=index_candidate_cache,
        query=query,
        clone_index_candidate_payload_fn=clone_index_candidate_payload_fn,
        store_cached_index_candidates_fn=store_cached_index_candidates_fn,
        attach_index_candidate_cache_info_fn=attach_index_candidate_cache_info_fn,
    )


def finalize_index_stage_output_from_state(
    *,
    state: Any,
    repo: str,
    root: str,
    fusion_mode: str,
    hybrid_re2_rrf_k: int,
    top_k_files: int,
    candidate_relative_threshold: float,
    chunk_top_k: int,
    chunk_per_file_limit: int,
    chunk_token_budget: int,
    chunk_disclosure: str,
    timings_ms: dict[str, Any],
    query: str,
    policy_version: str,
    build_index_stage_output_fn: Callable[..., dict[str, Any]],
    clone_index_candidate_payload_fn: Callable[..., dict[str, Any]],
    store_cached_index_candidates_fn: Callable[..., bool],
    attach_index_candidate_cache_info_fn: Callable[..., dict[str, Any]],
) -> dict[str, Any]:
    effective_corpus_size = getattr(state, "effective_corpus_size", state.corpus_size)
    return finalize_index_stage_output(
        repo=repo,
        root=root,
        terms=state.terms,
        memory_paths=state.memory_paths,
        index_hash=state.index_hash,
        index_data=state.index_data,
        cache_info=state.cache_info,
        requested_ranker=state.requested_ranker,
        selected_ranker=state.selected_ranker,
        ranker_fallbacks=state.ranker_fallbacks,
        corpus_size=int(effective_corpus_size),
        min_score_used=state.min_score_used,
        fusion_mode=fusion_mode,
        hybrid_re2_rrf_k=int(hybrid_re2_rrf_k),
        top_k_files=int(top_k_files),
        candidate_relative_threshold=float(candidate_relative_threshold),
        chunk_top_k=int(chunk_top_k),
        chunk_per_file_limit=int(chunk_per_file_limit),
        chunk_token_budget=int(chunk_token_budget),
        chunk_disclosure=str(chunk_disclosure),
        candidates=state.candidates,
        candidate_chunks=state.candidate_chunks,
        chunk_metrics=state.chunk_metrics,
        exact_search_payload=state.exact_search_payload,
        docs_payload=state.docs_payload,
        worktree_prior=state.worktree_prior,
        parallel_payload=state.parallel_payload,
        prior_payload=state.prior_payload,
        graph_lookup_payload=state.graph_lookup_payload,
        cochange_payload=state.cochange_payload,
        scip_payload=state.scip_payload,
        embeddings_payload=state.embeddings_payload,
        feedback_payload=state.feedback_payload,
        multi_channel_fusion_payload=state.multi_channel_fusion_payload,
        retrieval_refinement_payload=state.retrieval_refinement_payload,
        second_pass_payload=state.second_pass_payload,
        refine_pass_payload=state.refine_pass_payload,
        chunk_semantic_rerank_payload=state.chunk_semantic_rerank_payload,
        topological_shield_payload=state.topological_shield_payload,
        chunk_guard_payload=state.chunk_guard_payload,
        adaptive_router_payload=state.adaptive_router_payload,
        policy_name=str(state.policy.get("name", "general")),
        policy_version=str(policy_version),
        timings_ms=timings_ms,
        benchmark_filter_payload=state.benchmark_filter_payload,
        index_candidate_cache_path=state.index_candidate_cache_path,
        index_candidate_cache_key=state.index_candidate_cache_key,
        index_candidate_cache_ttl_seconds=state.index_candidate_cache_ttl_seconds,
        index_candidate_cache_required_meta=state.index_candidate_cache_required_meta,
        index_candidate_cache=state.index_candidate_cache,
        query=query,
        build_index_stage_output_fn=build_index_stage_output_fn,
        clone_index_candidate_payload_fn=clone_index_candidate_payload_fn,
        store_cached_index_candidates_fn=store_cached_index_candidates_fn,
        attach_index_candidate_cache_info_fn=attach_index_candidate_cache_info_fn,
    )


__all__ = [
    "finalize_index_stage_output",
    "finalize_index_stage_output_from_state",
]
