from __future__ import annotations

from dataclasses import asdict, dataclass
from typing import Any

from router.query_classifier import RouteDecisionFlags
from router.store_sufficiency import StoreSufficiencyResult


ALL_ROUTES = [
    "live_mcp_drilldown",
    "store_canonical",
    "store_feature_risk",
    "hybrid_store_plus_live",
    "batch_refresh_then_store",
]


@dataclass
class RouteSelectionResult:
    chosen_route: str
    candidate_routes: list[str]
    rejected_routes: dict[str, str]

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)


def choose_route(
    flags: RouteDecisionFlags,
    suff: StoreSufficiencyResult,
    *,
    parsed_as_trend_or_risk: bool,
) -> RouteSelectionResult:
    rejected: dict[str, str] = {}

    if flags.needs_exact_object_trace:
        rejected["store_canonical"] = "exact_trace_requires_live"
        rejected["store_feature_risk"] = "exact_trace_requires_live"
        rejected["hybrid_store_plus_live"] = "exact_trace_prefers_direct_live"
        rejected["batch_refresh_then_store"] = "exact_trace_not_batch"
        return RouteSelectionResult(
            chosen_route="live_mcp_drilldown",
            candidate_routes=list(ALL_ROUTES),
            rejected_routes=rejected,
        )

    heavy_shape = (
        flags.needs_full_period_aggregation
        or flags.needs_ranking
        or (flags.needs_anomaly_summary and not parsed_as_trend_or_risk)
    )
    if heavy_shape:
        aggregate_ok = (
            flags.precomputed_aggregate_available
            and suff.freshness_ok
            and suff.aggregate_level_ok
            and (not flags.needs_ranking or suff.ranking_ready)
        )
        if not aggregate_ok:
            rejected["store_feature_risk"] = "aggregate_not_sufficient"
            rejected["store_canonical"] = "wrong_query_shape"
            rejected["live_mcp_drilldown"] = "heavy_query_not_live"
            return RouteSelectionResult(
                chosen_route="batch_refresh_then_store",
                candidate_routes=list(ALL_ROUTES),
                rejected_routes=rejected,
            )

    if flags.needs_cross_entity_join and flags.needs_causal_chain:
        if not suff.explanation_ready:
            rejected["store_canonical"] = "cross_entity_causal_needs_live_stitching"
            return RouteSelectionResult(
                chosen_route="hybrid_store_plus_live",
                candidate_routes=list(ALL_ROUTES),
                rejected_routes=rejected,
            )

    if parsed_as_trend_or_risk:
        if suff.feature_sufficient and not flags.needs_runtime_truth:
            rejected["store_canonical"] = "trend_risk_query_prefers_feature"
            return RouteSelectionResult(
                chosen_route="store_feature_risk",
                candidate_routes=list(ALL_ROUTES),
                rejected_routes=rejected,
            )
        rejected["store_feature_risk"] = "feature_or_freshness_not_sufficient"

    if (
        suff.canonical_sufficient
        and not flags.needs_causal_chain
        and not flags.ambiguous_object_scope
        and not flags.needs_runtime_truth
    ):
        if flags.needs_cross_entity_join and not suff.explanation_ready:
            rejected["store_canonical"] = "cross_entity_needs_explanation_stitching"
        else:
            return RouteSelectionResult(
                chosen_route="store_canonical",
                candidate_routes=list(ALL_ROUTES),
                rejected_routes=rejected,
            )

    if not suff.canonical_sufficient:
        rejected["store_canonical"] = "canonical_not_sufficient"
    if flags.ambiguous_object_scope:
        rejected["store_canonical"] = "ambiguous_scope_requires_hybrid_or_feature"

    return RouteSelectionResult(
        chosen_route="hybrid_store_plus_live",
        candidate_routes=list(ALL_ROUTES),
        rejected_routes=rejected,
    )