agents/master/graph.py

from __future__ import annotations

from collections import defaultdict

from .schema import DoorNode, Edge, NpcTrade, ReadableNode, UseEffect, WorldDefinition


def readable_clue_mapping(world: WorldDefinition) -> dict[str, str]:
    return {node.id: node.clue_id for node in world.nodes if isinstance(node, ReadableNode)}


def clue_source_mapping(world: WorldDefinition) -> dict[str, str]:
    mapping = {node.clue_id: node.id for node in world.nodes if isinstance(node, ReadableNode)}
    for node in world.nodes:
        if node.type == "npc" and node.gives_clue_id:
            mapping[node.gives_clue_id] = node.id
    return mapping


def npc_trade_mapping(world: WorldDefinition) -> dict[str, NpcTrade]:
    trades: dict[str, NpcTrade] = {}
    for node in world.nodes:
        if node.type != "npc" or node.id == world.meta.win_condition.target_npc_id:
            continue
        trades[node.id] = NpcTrade(
            required_item_id=node.requires_item_id or "",
            gives_item_id=node.gives_item_id,
            gives_clue_id=node.gives_clue_id,
        )
    return trades


def use_effect_mapping(world: WorldDefinition) -> dict[str, UseEffect]:
    effects: dict[str, UseEffect] = {}
    for node in world.nodes:
        if node.type == "readable" and node.requires_item_id:
            effects[node.id] = UseEffect(
                required_item_id=node.requires_item_id,
                clue_id=node.clue_id,
                consumes_item=node.consumes_item,
            )
        elif node.type == "fixture":
            effects[node.id] = UseEffect(
                required_item_id=node.requires_item_id,
                reveals_item_id=node.reveals_item_id,
                reveals_readable_id=node.reveals_readable_id,
                consumes_item=node.consumes_item,
            )
    return effects


def recipe_mapping(world: WorldDefinition) -> dict[frozenset[str], str]:
    return {frozenset(recipe.input_item_ids): recipe.output_item_id for recipe in world.recipes}


def produced_item_ids(world: WorldDefinition) -> set[str]:
    produced = {recipe.output_item_id for recipe in world.recipes}
    for node in world.nodes:
        if node.type == "npc" and node.gives_item_id:
            produced.add(node.gives_item_id)
        if node.type == "fixture" and node.reveals_item_id:
            produced.add(node.reveals_item_id)
    return produced


def hidden_readable_ids(world: WorldDefinition) -> set[str]:
    return {node.reveals_readable_id for node in world.nodes if node.type == "fixture" and node.reveals_readable_id}


def door_room_mapping(world: WorldDefinition) -> dict[str, frozenset[str]]:
    mapping: dict[str, set[str]] = defaultdict(set)
    for edge in world.edges:
        if edge.door_node_id:
            mapping[edge.door_node_id].add(edge.from_node_id)
            mapping[edge.door_node_id].add(edge.to_node_id)
    return {door_id: frozenset(rooms) for door_id, rooms in mapping.items()}


def edge_for_door(world: WorldDefinition, door_id: str) -> Edge | None:
    for edge in world.edges:
        if edge.door_node_id == door_id:
            return edge
    return None


def door_nodes(world: WorldDefinition) -> dict[str, DoorNode]:
    return {node.id: node for node in world.nodes if isinstance(node, DoorNode)}