"""DeliveryPlanner - Plans which trucks deliver which packages."""

from typing import List, Dict, Tuple, Optional
from .delivery_search import DeliverySearch
from ..models.grid import Grid
from ..models.entities import Store, Destination, Tunnel
from ..models.state import PathResult, PlanResult, DeliveryAssignment, SearchStep


class DeliveryPlanner:
    """
    Plans the assignment of destinations to stores/trucks.

    For each destination, finds the optimal store to deliver from and computes the delivery path.
    """

    def __init__(
        self,
        grid: Grid,
        stores: List[Store],
        destinations: List[Destination],
        tunnels: Optional[List[Tunnel]] = None,
    ):
        """
        Initialize the delivery planner.

        Args:
            grid: The city grid with traffic information
            stores: List of store locations (each has a truck)
            destinations: List of customer destinations
            tunnels: Available tunnels
        """
        self.grid = grid
        self.stores = stores
        self.destinations = destinations
        self.tunnels = tunnels or []

    def plan(
        self, strategy: str, visualize: bool = False
    ) -> Tuple[PlanResult, Optional[Dict[int, List[SearchStep]]]]:
        """
        Create delivery plan assigning destinations to stores.

        Algorithm:
        1. For each destination, compute path cost from each store
        2. Assign each destination to the store with minimum path cost
        3. Compile results

        Args:
            strategy: Search strategy to use
            visualize: If True, collect visualization steps

        Returns:
            Tuple of (PlanResult, Optional visualization data)
        """
        assignments: List[DeliveryAssignment] = []
        all_steps: Dict[int, List[SearchStep]] = {} if visualize else None
        total_cost = 0.0
        total_nodes = 0

        # For each destination, find best store
        for dest in self.destinations:
            best_store: Optional[Store] = None
            best_result: Optional[PathResult] = None
            best_steps: Optional[List[SearchStep]] = None
            best_cost = float("inf")

            # Try each store
            for store in self.stores:
                result, steps = DeliverySearch.path(
                    self.grid,
                    store.position,
                    dest.position,
                    self.tunnels,
                    strategy,
                    visualize,
                )

                # Track nodes expanded
                total_nodes += result.nodes_expanded

                # Check if this is better
                if result.cost < best_cost:
                    best_cost = result.cost
                    best_store = store
                    best_result = result
                    best_steps = steps

            # Create assignment
            if best_store and best_result:
                assignment = DeliveryAssignment(
                    store_id=best_store.id,
                    destination_id=dest.id,
                    path_result=best_result,
                )
                assignments.append(assignment)
                total_cost += best_result.cost

                if visualize and best_steps:
                    all_steps[dest.id] = best_steps

        return (
            PlanResult(
                assignments=assignments,
                total_cost=total_cost,
                total_nodes_expanded=total_nodes,
            ),
            all_steps,
        )

    @staticmethod
    def plan_from_state(
        grid: Grid,
        stores: List[Store],
        destinations: List[Destination],
        tunnels: List[Tunnel],
        strategy: str,
        visualize: bool = False,
    ) -> Tuple[PlanResult, Optional[Dict]]:
        """
        Static method to create and run planner.

        Args:
            grid: The city grid
            stores: List of stores
            destinations: List of destinations
            tunnels: Available tunnels
            strategy: Search strategy
            visualize: If True, collect visualization steps

        Returns:
            Tuple of (PlanResult, Optional visualization data)
        """
        planner = DeliveryPlanner(grid, stores, destinations, tunnels)
        return planner.plan(strategy, visualize)