from __future__ import annotations

import logging
from collections import defaultdict
from dataclasses import dataclass, field
from functools import singledispatchmethod
from typing import Dict, Iterator, List, Optional, TYPE_CHECKING, Type

from dcs.unittype import FlyingType, VehicleType

from gen.ato import Package
from gen.flights.flightplan import FlightPlanBuilder
from game.theater import ControlPoint, MissionTarget
from game.theater.supplyroutes import SupplyRoute
from gen.naming import namegen
from gen.flights.flight import Flight, FlightType

if TYPE_CHECKING:
    from game import Game
    from game.inventory import ControlPointAircraftInventory


# TODO: Remove base classes.
# Eventually we'll want multi-mode transfers (convoy from factory to port, onto a ship,
# then airlifted to the final destination, etc). To do this we'll need to make the
# transfer *order* represent the full journey and let classes like Convoy handle the
# individual hops.
@dataclass
class TransferOrder:
    """The base type of all transfer orders.

    A transfer order can transfer multiple units of multiple types.
    """

    #: The location the units are transferring from.
    origin: ControlPoint

    #: The location the units are transferring to.
    destination: ControlPoint

    #: True if the transfer order belongs to the player.
    player: bool

    #: The units being transferred.
    units: Dict[Type[VehicleType], int]

    @property
    def description(self) -> str:
        raise NotImplementedError


@dataclass
class RoadTransferOrder(TransferOrder):
    """A transfer order that moves units by road."""

    #: The current position of the group being transferred. Groups move one control
    #: point a turn through the supply line.
    position: ControlPoint = field(init=False)

    def __post_init__(self) -> None:
        self.position = self.origin

    def path(self) -> List[ControlPoint]:
        supply_route = SupplyRoute.for_control_point(self.position)
        return supply_route.shortest_path_between(self.position, self.destination)

    def next_stop(self) -> ControlPoint:
        return self.path()[0]

    @property
    def description(self) -> str:
        path = self.path()
        if len(path) == 1:
            turns = "1 turn"
        else:
            turns = f"{len(path)} turns"
        return f"Currently at {self.position}. Arrives at destination in {turns}."


@dataclass
class AirliftOrder(TransferOrder):
    """A transfer order that moves units by cargo planes and helicopters."""

    flight: Flight

    @property
    def description(self) -> str:
        return "Airlift"

    def iter_units(self) -> Iterator[Type[VehicleType]]:
        for unit_type, count in self.units.items():
            for _ in range(count):
                yield unit_type

    def kill_unit(self, unit_type: Type[VehicleType]) -> None:
        if unit_type in self.units:
            self.units[unit_type] -= 1
            return
        raise KeyError


class AirliftPlanner:
    def __init__(
        self,
        game: Game,
        pickup: ControlPoint,
        drop_off: ControlPoint,
        units: Dict[Type[VehicleType], int],
    ) -> None:
        self.game = game
        self.pickup = pickup
        self.drop_off = drop_off
        self.units = units
        self.for_player = drop_off.captured
        self.package = Package(target=drop_off, auto_asap=True)

    def create_package_for_airlift(self) -> Dict[Type[VehicleType], int]:
        for cp in self.game.theater.player_points():
            inventory = self.game.aircraft_inventory.for_control_point(cp)
            for unit_type, available in inventory.all_aircraft:
                if unit_type.helicopter:
                    while available and self.needed_capacity:
                        flight_size = self.create_airlift_flight(unit_type, inventory)
                        available -= flight_size
        self.game.ato_for(self.for_player).add_package(self.package)
        return self.units

    def take_units(self, count: int) -> Dict[Type[VehicleType], int]:
        taken = {}
        for unit_type, remaining in self.units.items():
            take = min(remaining, count)
            count -= take
            self.units[unit_type] -= take
            taken[unit_type] = take
            if not count:
                break
        return taken

    @property
    def needed_capacity(self) -> int:
        return sum(c for c in self.units.values())

    def create_airlift_flight(
        self, unit_type: Type[FlyingType], inventory: ControlPointAircraftInventory
    ) -> int:
        available = inventory.available(unit_type)
        # 4 is the max flight size in DCS.
        flight_size = min(self.needed_capacity, available, 4)
        flight = Flight(
            self.package,
            self.game.player_country,
            unit_type,
            flight_size,
            FlightType.TRANSPORT,
            self.game.settings.default_start_type,
            departure=inventory.control_point,
            arrival=inventory.control_point,
            divert=None,
        )

        transfer = AirliftOrder(
            player=True,
            origin=self.pickup,
            destination=self.drop_off,
            units=self.take_units(flight_size),
            flight=flight,
        )
        flight.cargo = transfer

        self.package.add_flight(flight)
        planner = FlightPlanBuilder(self.game, self.package, self.for_player)
        planner.populate_flight_plan(flight)
        self.game.aircraft_inventory.claim_for_flight(flight)
        self.game.transfers.new_transfer(transfer)
        return flight_size


class Convoy(MissionTarget):
    def __init__(self, origin: ControlPoint, destination: ControlPoint) -> None:
        super().__init__(namegen.next_convoy_name(), origin.position)
        self.origin = origin
        self.destination = destination
        self.transfers: List[RoadTransferOrder] = []

    def mission_types(self, for_player: bool) -> Iterator[FlightType]:
        if self.is_friendly(for_player):
            return

        yield FlightType.BAI
        yield from super().mission_types(for_player)

    def is_friendly(self, to_player: bool) -> bool:
        return self.origin.captured

    def add_units(self, transfer: RoadTransferOrder) -> None:
        self.transfers.append(transfer)

    def remove_units(self, transfer: RoadTransferOrder) -> None:
        self.transfers.remove(transfer)

    def kill_unit(self, unit_type: Type[VehicleType]) -> None:
        for transfer in self.transfers:
            if unit_type in transfer.units:
                transfer.units[unit_type] -= 1
                return
        raise KeyError

    @property
    def size(self) -> int:
        return sum(sum(t.units.values()) for t in self.transfers)

    @property
    def units(self) -> Dict[Type[VehicleType], int]:
        units: Dict[Type[VehicleType], int] = defaultdict(int)
        for transfer in self.transfers:
            for unit_type, count in transfer.units.items():
                units[unit_type] += count
        return units

    @property
    def player_owned(self) -> bool:
        return self.origin.captured


class ConvoyMap:
    def __init__(self) -> None:
        # Dict of origin -> destination -> convoy.
        self.convoys: Dict[ControlPoint, Dict[ControlPoint, Convoy]] = defaultdict(dict)

    def convoy_exists(self, origin: ControlPoint, destination: ControlPoint) -> bool:
        return destination in self.convoys[origin]

    def find_convoy(
        self, origin: ControlPoint, destination: ControlPoint
    ) -> Optional[Convoy]:
        return self.convoys[origin].get(destination)

    def find_or_create_convoy(
        self, origin: ControlPoint, destination: ControlPoint
    ) -> Convoy:
        convoy = self.find_convoy(origin, destination)
        if convoy is None:
            convoy = Convoy(origin, destination)
            self.convoys[origin][destination] = convoy
        return convoy

    def departing_from(self, origin: ControlPoint) -> Iterator[Convoy]:
        yield from self.convoys[origin].values()

    def travelling_to(self, destination: ControlPoint) -> Iterator[Convoy]:
        for destination_dict in self.convoys.values():
            if destination in destination_dict:
                yield destination_dict[destination]

    def disband_convoy(self, convoy: Convoy) -> None:
        del self.convoys[convoy.origin][convoy.destination]

    def add(self, transfer: RoadTransferOrder) -> None:
        next_stop = transfer.next_stop()
        self.find_or_create_convoy(transfer.position, next_stop).add_units(transfer)

    def remove(self, transfer: RoadTransferOrder) -> None:
        next_stop = transfer.next_stop()
        convoy = self.find_convoy(transfer.position, next_stop)
        if convoy is None:
            logging.error(
                f"Attempting to remove {transfer} from convoy but it is in no convoy."
            )
            return
        convoy.remove_units(transfer)
        if not convoy.transfers:
            self.disband_convoy(convoy)

    def disband_all(self) -> None:
        self.convoys = defaultdict(dict)

    def __iter__(self) -> Iterator[Convoy]:
        for destination_dict in self.convoys.values():
            yield from destination_dict.values()


class PendingTransfers:
    def __init__(self, game: Game) -> None:
        self.game = game
        self.convoys = ConvoyMap()
        self.pending_transfers: List[TransferOrder] = []

    def __iter__(self) -> Iterator[TransferOrder]:
        yield from self.pending_transfers

    @property
    def pending_transfer_count(self) -> int:
        return len(self.pending_transfers)

    def transfer_at_index(self, index: int) -> TransferOrder:
        return self.pending_transfers[index]

    def index_of_transfer(self, transfer: TransferOrder) -> int:
        return self.pending_transfers.index(transfer)

    # TODO: Move airlift arrangements here?
    @singledispatchmethod
    def arrange_transport(self, transfer) -> None:
        pass

    @arrange_transport.register
    def _arrange_transport_road(self, transfer: RoadTransferOrder) -> None:
        self.convoys.add(transfer)

    def new_transfer(self, transfer: TransferOrder) -> None:
        transfer.origin.base.commit_losses(transfer.units)
        self.pending_transfers.append(transfer)
        self.arrange_transport(transfer)

    @singledispatchmethod
    def cancel_transport(self, transfer) -> None:
        pass

    @cancel_transport.register
    def _cancel_transport_air(self, transfer: AirliftOrder) -> None:
        flight = transfer.flight
        flight.package.remove_flight(flight)
        self.game.aircraft_inventory.return_from_flight(flight)

    @cancel_transport.register
    def _cancel_transport_road(self, transfer: RoadTransferOrder) -> None:
        self.convoys.remove(transfer)

    def cancel_transfer(self, transfer: TransferOrder) -> None:
        self.cancel_transport(transfer)
        self.pending_transfers.remove(transfer)
        transfer.origin.base.commision_units(transfer.units)

    def perform_transfers(self) -> None:
        incomplete = []
        for transfer in self.pending_transfers:
            if not self.perform_transfer(transfer):
                incomplete.append(transfer)
        self.pending_transfers = incomplete
        self.rebuild_convoys()

    def rebuild_convoys(self) -> None:
        self.convoys.disband_all()
        for transfer in self.pending_transfers:
            self.arrange_transport(transfer)

    @singledispatchmethod
    def perform_transfer(self, transfer) -> bool:
        raise NotImplementedError

    @perform_transfer.register
    def _perform_transfer_air(self, transfer: AirliftOrder) -> bool:
        if transfer.player != transfer.destination.captured:
            logging.info(
                f"Transfer destination {transfer.destination} was captured. Cancelling "
                "transport."
            )
            transfer.origin.base.commision_units(transfer.units)
            return True

        transfer.destination.base.commision_units(transfer.units)
        return True

    @perform_transfer.register
    def _perform_transfer_road(self, transfer: RoadTransferOrder) -> bool:
        # TODO: Can be improved to use the convoy map.
        # The convoy map already has a lot of the data that we're recomputing here.
        if transfer.player != transfer.destination.captured:
            logging.info(
                f"Transfer destination {transfer.destination.name} was captured."
            )
            self.handle_route_interrupted(transfer)
            return True

        supply_route = SupplyRoute.for_control_point(transfer.destination)
        if transfer.position not in supply_route:
            logging.info(
                f"Route from {transfer.position.name} to {transfer.destination.name} "
                "was cut off."
            )
            self.handle_route_interrupted(transfer)
            return True

        path = transfer.path()
        next_hop = path[0]
        if next_hop == transfer.destination:
            logging.info(
                f"Units transferred from {transfer.origin.name} to "
                f"{transfer.destination.name}"
            )
            transfer.destination.base.commision_units(transfer.units)
            return True

        logging.info(
            f"Units transferring from {transfer.origin.name} to "
            f"{transfer.destination.name} arrived at {next_hop.name}. {len(path) - 1} "
            "turns remaining."
        )
        transfer.position = next_hop
        return False

    @staticmethod
    def handle_route_interrupted(transfer: RoadTransferOrder):
        # Halt the transfer in place if safe.
        if transfer.player == transfer.position.captured:
            logging.info(f"Transferring units are halting at {transfer.position.name}.")
            transfer.position.base.commision_units(transfer.units)
            return

        # If the current position was captured attempt to divert to a neighboring
        # friendly CP.
        for connected in transfer.position.connected_points:
            if connected.captured == transfer.player:
                logging.info(f"Transferring units are re-routing to {connected.name}.")
                connected.base.commision_units(transfer.units)
                return

        # If the units are cutoff they are destroyed.
        logging.info(
            f"Both {transfer.position.name} and {transfer.destination.name} were "
            "captured. Units were surrounded and destroyed during transfer."
        )