"""Unified graph routing helpers."""
import heapq

from .arango_client import ensure_graph


def _allowed_next_phase(current_phase, transport_type):
    """
    Phase-based routing: auto → rail* → auto.
    - end_auto: allow one auto, rail, or offer
    - end_auto_done: auto used — rail or offer
    - rail: any number of rail, then one auto or offer
    - start_auto_done: auto used — only offer
    """
    if current_phase == 'end_auto':
        if transport_type == 'offer':
            return 'offer'
        if transport_type == 'auto':
            return 'end_auto_done'
        if transport_type == 'rail':
            return 'rail'
        return None
    if current_phase == 'end_auto_done':
        if transport_type == 'offer':
            return 'offer'
        if transport_type == 'rail':
            return 'rail'
        return None
    if current_phase == 'rail':
        if transport_type == 'offer':
            return 'offer'
        if transport_type == 'rail':
            return 'rail'
        if transport_type == 'auto':
            return 'start_auto_done'
        return None
    if current_phase == 'start_auto_done':
        if transport_type == 'offer':
            return 'offer'
        return None
    return None


def _allowed_types_for_phase(phase):
    if phase == 'end_auto':
        return ['auto', 'rail', 'offer']
    if phase == 'end_auto_done':
        return ['rail', 'offer']
    if phase == 'rail':
        return ['rail', 'auto', 'offer']
    if phase == 'start_auto_done':
        return ['offer']
    return ['offer']


def _fetch_neighbors(db, node_key, allowed_types):
    aql = """
    FOR edge IN edges
        FILTER edge.transport_type IN @types
        FILTER edge._from == @node_id OR edge._to == @node_id
        LET neighbor_id = edge._from == @node_id ? edge._to : edge._from
        LET neighbor = DOCUMENT(neighbor_id)
        FILTER neighbor != null
        RETURN {
            neighbor_key: neighbor._key,
            neighbor_doc: neighbor,
            from_id: edge._from,
            to_id: edge._to,
            transport_type: edge.transport_type,
            distance_km: edge.distance_km,
            travel_time_seconds: edge.travel_time_seconds
        }
    """
    cursor = db.aql.execute(
        aql,
        bind_vars={'node_id': f"nodes/{node_key}", 'types': allowed_types},
    )
    return list(cursor)


def graph_find_targets(db, start_uuid, target_predicate, route_builder, limit=10, max_expansions=20000):
    """Unified graph traversal: auto → rail* → auto, returns routes for target nodes."""
    ensure_graph()

    nodes_col = db.collection('nodes')
    start = nodes_col.get(start_uuid)
    if not start:
        return []

    queue = []
    counter = 0
    heapq.heappush(queue, (0, counter, start_uuid, 'end_auto'))

    visited = {}
    predecessors = {}
    node_docs = {start_uuid: start}
    found = []
    expansions = 0

    while queue and len(found) < limit and expansions < max_expansions:
        cost, _, node_key, phase = heapq.heappop(queue)

        if (node_key, phase) in visited and cost > visited[(node_key, phase)]:
            continue
        visited[(node_key, phase)] = cost

        node_doc = node_docs.get(node_key)
        if node_doc and target_predicate(node_doc):
            path_edges = []
            state = (node_key, phase)
            current_key = node_key
            while state in predecessors:
                prev_state, edge_info = predecessors[state]
                prev_key = prev_state[0]
                path_edges.append((current_key, prev_key, edge_info))
                state = prev_state
                current_key = prev_key

            route = route_builder(path_edges, node_docs) if route_builder else None
            distance_km = route.total_distance_km if route else None

            found.append({
                'node': node_doc,
                'route': route,
                'distance_km': distance_km,
                'cost': cost,
            })
            continue

        neighbors = _fetch_neighbors(db, node_key, _allowed_types_for_phase(phase))
        expansions += 1

        for neighbor in neighbors:
            transport_type = neighbor.get('transport_type')
            next_phase = _allowed_next_phase(phase, transport_type)
            if next_phase is None:
                continue

            travel_time = neighbor.get('travel_time_seconds')
            distance_km = neighbor.get('distance_km')
            neighbor_key = neighbor.get('neighbor_key')
            if not neighbor_key:
                continue

            node_docs[neighbor_key] = neighbor.get('neighbor_doc')
            step_cost = travel_time if travel_time is not None else (distance_km or 0)
            new_cost = cost + step_cost

            state_key = (neighbor_key, next_phase)
            if state_key in visited and new_cost >= visited[state_key]:
                continue

            counter += 1
            heapq.heappush(queue, (new_cost, counter, neighbor_key, next_phase))
            predecessors[state_key] = ((node_key, phase), neighbor)

    return found


def snap_to_nearest_hub(db, lat, lon):
    aql = """
    FOR hub IN nodes
        FILTER hub.node_type == 'logistics' OR hub.node_type == null
        FILTER hub.product_uuid == null
        LET types = hub.transport_types != null ? hub.transport_types : []
        FILTER ('rail' IN types) OR ('sea' IN types)
        FILTER hub.latitude != null AND hub.longitude != null
        LET dist = DISTANCE(hub.latitude, hub.longitude, @lat, @lon) / 1000
        SORT dist ASC
        LIMIT 1
        RETURN hub
    """
    cursor = db.aql.execute(aql, bind_vars={'lat': lat, 'lon': lon})
    hubs = list(cursor)
    return hubs[0] if hubs else None


def resolve_start_hub(db, source_uuid=None, lat=None, lon=None):
    nodes_col = db.collection('nodes')

    if source_uuid:
        node = nodes_col.get(source_uuid)
        if not node:
            return None

        if node.get('node_type') in ('logistics', None):
            types = node.get('transport_types') or []
            if ('rail' in types) or ('sea' in types):
                return node

        node_lat = node.get('latitude')
        node_lon = node.get('longitude')
        if node_lat is None or node_lon is None:
            return None
        return snap_to_nearest_hub(db, node_lat, node_lon)

    if lat is None or lon is None:
        return None

    return snap_to_nearest_hub(db, lat, lon)