import heapq
import logging
from collections import defaultdict
from typing import Optional

from tqdm import tqdm

from graph import NavigationGraph as Graph, LegMetrics
from pathfinder import Ship
from arbitrage import Trade
from typing import Any

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

MAX_TRADES_TO_CONSIDER = 20_000
MAX_CAPITAL = 100_000_000

Route = list[dict[str, Any]]
RoutePlan = tuple[Optional[Route], Optional[dict], list[Trade], list[int], list[int]]
LegMetricLookup = dict[tuple[int, int], Optional[LegMetrics]]


def outstanding_trades_for_locations(location_ids: list[int], trades: list[Trade]) -> list[Trade]:
    remaining_locations = set(location_ids)
    return [trade for trade in trades if trade.to_station in remaining_locations]


def calculate_peak_capital(location_ids: list[int], trades: list[Trade], ship_cost: float) -> float:
    outstanding_trades = outstanding_trades_for_locations(location_ids, trades)
    if not outstanding_trades:
        return ship_cost

    remaining_locations = set(location_ids)
    committed_capital = ship_cost + sum(
        trade.from_price * trade.quantity
        for trade in outstanding_trades
        if trade.from_station not in remaining_locations
    )
    peak_capital = committed_capital

    for location_id in location_ids:
        committed_capital -= sum(
            trade.from_price * trade.quantity
            for trade in outstanding_trades
            if trade.to_station == location_id
        )
        committed_capital = max(ship_cost, committed_capital)
        committed_capital += sum(
            trade.from_price * trade.quantity
            for trade in outstanding_trades
            if trade.from_station == location_id
        )
        peak_capital = max(peak_capital, committed_capital)

    return peak_capital


def summarize_graph_path(graph: Graph, path: list[int]) -> LegMetrics:
    weight = 0.0
    transport_time = 0.0
    risk = 0.0

    for source, target in zip(path, path[1:]):
        edge = graph.graph[source][target]
        weight += edge["weight"]
        transport_time += edge["time"]
        risk += edge["risk"]

    return LegMetrics(weight=weight, time=transport_time, risk=risk)


def summarize_location_sequence(graph: Graph, location_ids: list[int]) -> Optional[tuple[float, float]]:
    if len(location_ids) < 2:
        return 0.0, 0.0

    transport_time = 0.0
    risk = 0.0

    for source, target in zip(location_ids, location_ids[1:]):
        if source == target:
            continue

        if graph.graph.has_edge(source, target):
            leg_metrics = LegMetrics(
                weight=graph.graph[source][target]["weight"],
                time=graph.graph[source][target]["time"],
                risk=graph.graph[source][target]["risk"],
            )
        else:
            path = graph.shortest_path(source, target)
            if path is None:
                return None
            leg_metrics = summarize_graph_path(graph, path)

        transport_time += leg_metrics.time
        risk += leg_metrics.risk

    return transport_time, risk


def build_route_info(
    location_ids: list[int],
    trades: list[Trade],
    ship: Ship,
    transport_time: float,
    risk: float,
) -> Optional[dict]:
    outstanding_trades = outstanding_trades_for_locations(location_ids, trades)
    if not outstanding_trades:
        return None

    capital = calculate_peak_capital(location_ids, outstanding_trades, ship.get("ship_cost", 0.0))
    gross_profit = sum(trade.gross_profit for trade in outstanding_trades)
    net_profit = gross_profit - risk * capital
    profit_rate = net_profit / transport_time if transport_time > 0 else 0.0

    return {
        "profit_rate": profit_rate,
        "risk": risk,
        "capital": capital,
        "transport_time": transport_time,
        "gross_profit": gross_profit,
        "net_profit": net_profit,
    }


def extract_route_location_ids(route_steps: Route) -> list[int]:
    return [
        step["location_id"]
        for step in route_steps
        if isinstance(step, dict) and isinstance(step.get("location_id"), int)
    ]


def calculate_route_info(
    route_steps: Route,
    trades: list[Trade],
    graph: Graph,
    ship: Ship,
    metric_path: Optional[list[int]] = None,
) -> Optional[dict]:
    location_ids = [
        location_id
        for location_id in extract_route_location_ids(route_steps)
    ]
    if not location_ids:
        return None

    if metric_path is not None:
        if len(metric_path) < 2:
            transport_time = 0.0
            risk = 0.0
        else:
            path_metrics = summarize_graph_path(graph, metric_path)
            transport_time = path_metrics.time
            risk = path_metrics.risk
    else:
        path_metrics = summarize_location_sequence(graph, location_ids)
        if path_metrics is None:
            return None

        transport_time, risk = path_metrics

    return build_route_info(location_ids, trades, ship, transport_time, risk)


def trade_capital_efficiency(trade: Trade) -> float:
    invested_capital = trade.from_price * trade.quantity
    return trade.gross_profit / invested_capital if invested_capital > 0 else 0.0


def trade_profit_density(trade: Trade) -> float:
    return trade.gross_profit / trade.cargo if trade.cargo > 0 else 0.0


def count_unique_triplets(trades: list[Trade], ship_location: Optional[int]) -> int:
    if ship_location is None:
        return 0

    return len({(ship_location, trade.from_station, trade.to_station) for trade in trades})


def build_ranked_triplets(trades: list[Trade], ship_location: int) -> list[tuple[int, int, int]]:
    ranked_triplets: list[tuple[int, int, int]] = []
    seen_triplets: set[tuple[int, int, int]] = set()

    for trade in sorted(trades, key=trade_capital_efficiency, reverse=True):
        triplet = (ship_location, trade.from_station, trade.to_station)
        if triplet in seen_triplets:
            continue

        seen_triplets.add(triplet)
        ranked_triplets.append(triplet)

    return ranked_triplets


def build_station_pair_index(trades: list[Trade]) -> dict[tuple[int, int], list[Trade]]:
    trades_by_station_pair: dict[tuple[int, int], list[Trade]] = defaultdict(list)

    for trade in trades:
        trades_by_station_pair[(trade.from_station, trade.to_station)].append(trade)

    for station_pair_trades in trades_by_station_pair.values():
        station_pair_trades.sort(key=trade_profit_density, reverse=True)

    return dict(trades_by_station_pair)


def route(trades: list[Trade], graph: Graph, ship: Ship, leg_metrics: Optional[LegMetricLookup] = None) -> RoutePlan:
    """
    Given a list of trades, find the best route starting from the ship's location.

    Args:
        trades: List of Trade instances.
        graph: Navigation graph.
        ship: Ship parameters.

    Returns:
        A tuple containing the best route and route information.
    """
    if not trades:
        return None, None, [], [], []

    ranked_triplets = build_ranked_triplets(trades, ship["location"])
    trades_by_station_pair = build_station_pair_index(trades)

    if leg_metrics is None:
        leg_metrics = graph.build_trade_leg_metrics(ship["location"], trades)

    best_profit_rate = 0.0
    best_route: Optional[Route] = None
    best_route_info: Optional[dict] = None
    best_route_stations: Optional[tuple[int, int, int]] = None
    best_optimized_trades: Optional[list[Trade]] = None

    for route_stations in tqdm(ranked_triplets[:MAX_TRADES_TO_CONSIDER], desc="Evaluating triplets"):
        optimized_trades = select_trades(route_stations, trades_by_station_pair, ship)
        if not optimized_trades:
            continue

        source_leg = leg_metrics.get((route_stations[0], route_stations[1]))
        destination_leg = leg_metrics.get((route_stations[1], route_stations[2]))

        if source_leg is None or destination_leg is None:
            continue

        risk = source_leg.risk + destination_leg.risk
        transport_time = source_leg.time + destination_leg.time

        candidate_route_info = build_route_info(list(route_stations), optimized_trades, ship, transport_time, risk)
        if candidate_route_info is None:
            continue

        profit_rate = candidate_route_info["profit_rate"]

        if profit_rate > best_profit_rate:
            logger.info(
                f"New best profit rate: {profit_rate * 3600:.2f} per hour, "
                f"Risk: {risk:.4f}, Capital: {candidate_route_info['capital']:.2f}, "
                f"Transport time: {transport_time:.2f} sec, Gross profit: {candidate_route_info['gross_profit']:.2f}"
            )
            best_route_info = candidate_route_info
            best_profit_rate = profit_rate
            best_route_stations = route_stations
            best_optimized_trades = optimized_trades

    if best_route_stations is None or best_optimized_trades is None:
        return None, None, [], [], []

    path = reconstruct_route_path(graph, best_route_stations)
    if path is None:
        logger.error("Unable to reconstruct best route path for stations %s", best_route_stations)
        return None, None, [], [], []

    formatted_route, raw_offsets = graph.formatted_route_with_offsets(path)
    best_route = set_actions(formatted_route, best_optimized_trades)
    best_route_info = calculate_route_info(best_route, best_optimized_trades, graph, ship, metric_path=path) or best_route_info

    return best_route, best_route_info, best_optimized_trades, path, raw_offsets

def reconstruct_route_path(graph: Graph, route_stations: tuple[int, int, int]) -> Optional[list[int]]:
    source_path = graph.shortest_path(route_stations[0], route_stations[1])
    if source_path is None:
        return None

    destination_path = graph.shortest_path(route_stations[1], route_stations[2])
    if destination_path is None:
        return None

    path = source_path[:-1]
    path.extend(destination_path)
    return path


def select_trades(route_stations: tuple[int, int, int], trades_by_station_pair: dict[tuple[int, int], list[Trade]], ship: Ship) -> list[Trade]:
    """
    Given a triplet of stations, find the optimal trades to add.

    Args:
        route_stations: Ordered station triplet for the route.
        trades_by_station_pair: Trade instances grouped and pre-sorted by ordered station pair.
        ship: Ship parameters.

    Returns:
        A list of optimized Trade instances.
    """
    candidate_pairs: list[tuple[int, int]] = []
    for index, from_station in enumerate(route_stations):
        for to_station in route_stations[index + 1:]:
            station_pair = (from_station, to_station)
            if station_pair not in candidate_pairs:
                candidate_pairs.append(station_pair)

    pair_groups = [
        trades_by_station_pair[station_pair]
        for station_pair in candidate_pairs
        if station_pair in trades_by_station_pair
    ]

    if not pair_groups:
        return []

    filtered_trades = heapq.merge(*pair_groups, key=trade_profit_density, reverse=True)

    used_orders = {}
    optimized_trades = []
    cargo = 0.0
    capital = 0.0

    for trade in filtered_trades:
        quantity = trade.quantity
        from_order, to_order = trade.from_order_id, trade.to_order_id

        from_remaining = used_orders.get(from_order, quantity)
        to_remaining = used_orders.get(to_order, quantity)

        max_quantity_by_cargo = int((ship["cargo"] - cargo) // trade.item_cargo_volume)
        max_quantity_by_capital = int((MAX_CAPITAL - capital) // trade.from_price)

        max_quantity = min(quantity, from_remaining, to_remaining, max_quantity_by_cargo, max_quantity_by_capital)

        if max_quantity <= 0:
            continue

        cargo += trade.item_cargo_volume * max_quantity
        capital += trade.from_price * max_quantity
        used_orders[from_order] = from_remaining - max_quantity
        used_orders[to_order] = to_remaining - max_quantity

        # Create a new Trade instance with updated quantities
        updated_trade = Trade(
            from_station=trade.from_station,
            to_station=trade.to_station,
            from_order_id=trade.from_order_id,
            to_order_id=trade.to_order_id,
            item_name=trade.item_name,
            type_id=trade.type_id,
            item_cargo_volume=trade.item_cargo_volume,
            from_price=trade.from_price,
            to_price=trade.to_price,
            quantity=max_quantity,
            cargo=trade.item_cargo_volume * max_quantity,
            gross_profit=trade.gross_profit * (max_quantity / trade.quantity) if trade.quantity > 0 else 0.0,
            net_profit=None  # Will be calculated later if needed
        )

        optimized_trades.append(updated_trade)

        if cargo >= ship["cargo"] or capital >= MAX_CAPITAL:
            break

    return optimized_trades

def set_actions(route: Route, trades: list[Trade]) -> Route:
    """
    Sets the actions (buy/sell) at each step in the route based on the trades.

    Args:
        route: The route steps.
        trades: List of Trade instances.

    Returns:
        The route with actions set.
    """
    bought_items = set()
    for trade in trades:
        for step in route:
            if trade.from_station == step["location_id"]:
                step["actions"].append({
                    "action_type": "buy",
                    "item": trade.item_name,
                    "type_id": trade.type_id,
                    "quantity": trade.quantity,
                    "price": trade.from_price,
                })
                bought_items.add(trade.type_id)
            if trade.to_station == step["location_id"] and trade.type_id in bought_items:
                step["actions"].append({
                    "action_type": "sell",
                    "item": trade.item_name,
                    "type_id": trade.type_id,
                    "quantity": trade.quantity,
                    "price": trade.to_price,
                })

    for step in route:
        step["actions"].sort(
            key=lambda action: ({"sell": 0, "buy": 1}[action["action_type"]], action["item"], action["price"])
        )

    return route