from abc import ABC
import copy
from enum import Enum
from typing import List, Tuple
import gymnasium as gym

class MyEnum(Enum):
    @classmethod
    def from_value(cls, value):
        """根据数值获取相应的枚举成员"""
        for member in cls:
            if member.value == value:
                return member
        raise ValueError(f"{value} is not a valid value for {cls.__name__}")

    @classmethod
    def from_string(cls, name):
        """根据字符串获取相应的枚举成员"""
        try:
            return cls[name]
        except KeyError:
            raise ValueError(f"{name} is not a valid {cls.__name__}")

    @classmethod
    def from_input(cls, input):
        """根据输入（字符串或数值）获取相应的枚举成员"""
        if isinstance(input, str):
            return cls.from_string(input)
        elif isinstance(input, int):
            return cls.from_value(input)
        else:
            raise ValueError("Input must be either a string or an integer.")

class MoveType(MyEnum):
    AIR = 0
    GROUND = 1
    SEA = 2
    SUBWATER = 3

class TerrainType(MyEnum):
    PLAIN = 0
    ROAD = 1
    SEA = 2
    SUBWATER = 3

class ActionType(MyEnum):
    ILLEGAL = -1
    END_OF_TURN = 0
    MOVE = 1
    ATTACK = 2
    INTERACT = 3
    RELEASE = 4
    SWITCH_WEAPON = 5
    SUPPLY = 6

class ActionStatus(MyEnum):
    VALID = 0
    OCCUPIED_DES = 1
    OUT_OF_RANGE = 2

class AgentType(MyEnum):
    Infantry = 0 # 步兵
    Tank = 1 # 装甲单位
    AntiAir = 2 # 自行防空炮
    MobilizedInfantry = 3 # 机械化步兵
    Helicopter = 4 # 直升机
    Fighter = 5 # 战斗机
    UAV = 6 # 无人机
    CombatShip = 7 # 战舰
    Carrier = 8 # 航母
    Artillery = 9 # 火炮
    Construction = 10 # 建筑
    MissileLauncher = 11 # 导弹发射车
    TransportHelicopter = 12 # 运输直升机
    SupplyTruck = 13 # 移动补给车
    Submarine = 14 # 潜艇
    AdvancedTank = 15 # 先进坦克
    Airport = 16 # 机场
    SupplyStation = 17 # 补给站
    RadarStation = 18 # 雷达站
    Bomber = 19 # 轰炸机
    AWACS = 20 # 预警机
    TransportShip = 21 # 运输船
    CommandPost = 22 # 指挥部

class ModuleType(MyEnum):
    HANGAR = 0
    SUPPLY = 1
    TRANSPORT = 2   

class Weapon:
    def __init__(self, name: str, attack_range: int, damage: float, max_ammo: int, ammo: int = None, strike_types: List[MoveType] = []):
        self.name = name
        self.attack_range = attack_range 
        self.damage = damage
        self.max_ammo = max_ammo
        self.ammo = ammo if ammo is not None else max_ammo
        self.strike_types = strike_types

    def reset(self):
        self.ammo = self.max_ammo

    def to_dict(self):
        return {
            "name": self.name,
            "attack_range": self.attack_range,
            "damage": self.damage,
            "max_ammo": self.max_ammo,
            "ammo": self.ammo,
            "strike_types": [strike_type.name for strike_type in self.strike_types]
        }
    
    def __deepcopy__(self, memo):
        copied_strike_types = copy.deepcopy(self.strike_types, memo)
        return Weapon(self.name, self.attack_range, self.damage, self.max_ammo, self.ammo, copied_strike_types)

class Action:
    def __init__(self, 
                 action_type: ActionType=ActionType.ILLEGAL, 
                 agent_id: str = "", 
                 des: Tuple[int, int] = (-1, -1), 
                 target_id: str = "",
                 weapon_id: int = 0,
                 weapon_name: str = "",
                 start_time: int = 0,
                 end_time: int = -1,
                 state: 'AgentState' = None,
                 **kwargs):
        self.agent_id = agent_id
        self.target_id = target_id
        self.des = des
        self.action_type = action_type
        self.weapon_id = weapon_id
        self.weapon_name = weapon_name
        self.start_time = start_time
        self.end_time = start_time if end_time == -1 else end_time
        self.end_type = None
        self.state = state
        for key, value in kwargs.items():
            setattr(self, key, value)

    def json(self) -> dict:
        from .utils import axial_to_cube_dict
        return {
            "curOrderedUnitID": self.agent_id,
            "targetUnitID": self.target_id,
            "destination": axial_to_cube_dict((self.des[0]-1000, self.des[1]-1000)),
            "commandType": self.action_type.value
        }
    
    def __str__(self):
        if self.action_type == ActionType.MOVE:
            return f"{self.agent_id} moves to {self.des}"
        elif self.action_type == ActionType.ATTACK:
            return f"{self.agent_id} attacks {self.target_id}"
        elif self.action_type == ActionType.INTERACT:
            return f"Interact {self.agent_id} with {self.target_id}"
        elif self.action_type == ActionType.RELEASE:
            return f"Release {self.target_id} from {self.agent_id} to {self.des}"
        elif self.action_type == ActionType.END_OF_TURN:
            return f"End of turn"
        elif self.action_type == ActionType.SWITCH_WEAPON:
            return f"{self.agent_id} switches to {self.weapon_name}"
        else:
            return "Unknown action"
        
    def to_dict(self):
        return {
            "agent_id": self.agent_id,
            "target_id": self.target_id,
            "des": self.des,
            "action_type": self.action_type.name,
            "weapon_id": self.weapon_id,
            "weapon_name": self.weapon_name,
            "start_time": self.start_time,
            "end_type": self.end_type
        }
        
class Command(Action):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)

    def description(self) -> List[str]:
        if self.command.action_type == ActionType.MOVE:
            return [
                f"Agent {self.agent_id} started to move to {self.command.des} at time {self.start_time}.",
                f"Agent {self.agent_id} arrived at {self.command.des} at time {self.end_time}."
            ]
        elif self.command.action_type == ActionType.ATTACK:
            return [
                f"Agent {self.agent_id} started to attack {self.command.des} at time {self.start_time}.",
                f"Agent {self.agent_id} attacked {self.command.des} at time {self.end_time}."
            ]
        elif self.command.action_type == ActionType.SUPPLY:
            return [
                f"Agent {self.agent_id} started to get supply from {self.command.des} at time {self.start_time}.",
                f"Agent {self.agent_id} got supply from {self.command.des} at time {self.end_time}."
            ]
        elif self.command.action_type == ActionType.SWITCH_WEAPON:
            return [
                f"Agent {self.agent_id} started to switch weapon to {self.command.weapon_name} at time {self.start_time}.",
                f"Agent {self.agent_id} switched weapon to {self.command.weapon_name} at time {self.end_time}."
            ]
        else:
            raise NotImplementedError(f"Unsupported command type: {self.command.action_type}")

class Agent:
    def __init__(self, 
                 agent_id: str, 
                 agent_type: AgentType, 
                 team_id: int, 
                 faction_id: int,
                 move_type: MoveType, 
                 pos: Tuple[int, int], 
                 info_level: int, stealth_level: int, 
                 max_endurance: int, defense: int, 
                 max_fuel: float, mobility: float, 
                 switchable_weapons = [], 
                 modules = [],
                 is_key: bool = False):
        self.agent_id = agent_id
        self.agent_type = agent_type
        self.team_id = team_id
        self.faction_id = faction_id
        self.move_type = MoveType(move_type)
        self.init_pos = pos
        self.info_level = info_level
        self.stealth_level = stealth_level
        self.max_endurance = max_endurance
        self.defense = defense
        self.endurance = max_endurance
        self.max_fuel = max_fuel
        self.mobility = mobility
        self.switchable_weapons = switchable_weapons
        self.modules = modules
        self.is_key = is_key

        # Get action space
        from .utils import range_to_count
        max_attack_range = max(weapon.attack_range for weapon in self.switchable_weapons) if self.switchable_weapons else 0
        max_capacity = max(module.capacity if hasattr(module, "capacity") else 0 for module in self.modules) if self.modules else 0
        # max_capacity = 6 if self.modules and any(hasattr(module, "parked_agents") for module in self.modules) else 0
        self._action_space = gym.spaces.Dict(
            {
                action_type: space
                for action_type, space in {
                    ActionType.MOVE: gym.spaces.Discrete(range_to_count(int(self.mobility))) \
                        if self.mobility > 0 else None,
                    ActionType.ATTACK: gym.spaces.Discrete(range_to_count(max_attack_range)) \
                        if max_attack_range > 0 else None,
                    ActionType.SWITCH_WEAPON: gym.spaces.Discrete(len(self.switchable_weapons)) \
                        if len(self.switchable_weapons) >= 1 else None,
                    ActionType.INTERACT: gym.spaces.Discrete(range_to_count(1)),
                    ActionType.RELEASE: gym.spaces.Discrete(max_capacity) \
                        if max_capacity > 0 else None,
                }.items()
                if space is not None  # 过滤掉值为 None 的键值对
            }
        )

        self._has_supply = False
        self._available_types = []
        self._parked_agents = []
        self._capacity = max_capacity # 实际上应该最多只有一个运输单元

        for module in self.modules:
            if module.module_type == ModuleType.SUPPLY:
                self._has_supply = True
                self.supply = module
            elif module.module_type == ModuleType.HANGAR:
                self._parked_agents = module.parked_agents
            elif module.module_type == ModuleType.TRANSPORT:
                self._parked_agents = module.parked_agents
            self._available_types.extend(module.available_types)

        self.reset()

    def reset(self):
        self.pos = self.init_pos
        self.endurance = self.max_endurance
        self.fuel = self.max_fuel
        if self.switchable_weapons:
            for weapon in self.switchable_weapons:
                weapon.reset()
            self.weapon = self.switchable_weapons[0]
        else:
            self.weapon = None
        self.commenced_action = False
        self.cmd_todo = []
        self.todo = []
        self.is_carried = False
        
    def __lt__(self, other):
        return self.agent_id < other.agent_id
    
    def __eq__(self, other):
        return self.agent_id == other.agent_id
    
    def __hash__(self):
        return hash(self.agent_id)
    
    def __str__(self):
        return f"{self.agent_id} ({self.move_type.name}) at {self.pos}"
    
    def __repr__(self):
        return f"{self.agent_id} ({self.move_type.name}) at {self.pos}"
    
    def to_dict(self):
        return {
            "agent_id": self.agent_id,
            "agent_type": self.agent_type.name,
            "team_id": self.team_id,
            "faction_id": self.faction_id,
            "move_type": self.move_type.name,
            "info_level": self.info_level,
            "stealth_level": self.stealth_level,
            "defense": self.defense,
            "max_endurance": self.max_endurance,
            "max_fuel": self.max_fuel,
            "switchable_weapons": [weapon.to_dict() for weapon in self.switchable_weapons],
            "modules": [module.to_dict() for module in self.modules],
            "is_key": self.is_key,
            # 以下为可变属性
            "pos": {"q": self.pos[0], "r": self.pos[1]},
            "endurance": self.endurance,
            "fuel": self.fuel,
            "mobility": self.mobility,
            "weapon": self.weapon.to_dict() if self.weapon is not None else None
        }
    
    def plan_to_dict(self):
        state = self.todo[-1].state if self.todo else self.state
        return {
            "agent_id": self.agent_id,
            "agent_type": self.agent_type.name,
            "move_type": self.move_type.name,
            "defense": self.defense,
            "max_endurance": self.max_endurance,
            "max_fuel": self.max_fuel,
            "switchable_weapons": [weapon.to_dict() for weapon in self.switchable_weapons],
            "modules": [module.to_dict() for module in self.modules],
            "pos": list(state.pos),
            "endurance": self.endurance,
            "fuel": self.fuel,
            "mobility": self.mobility,
            "weapon": self.weapon.to_dict() if self.weapon is not None else None,
        }

    @property
    def alive(self):
        return self.endurance > 0
    
    @property
    def attack_range(self):
        return self.weapon.attack_range if self.weapon is not None else 0
    
    @property
    def strike_types(self):
        return self.weapon.strike_types if self.weapon is not None else []
    
    @property
    def damage(self):
        return self.weapon.damage if self.weapon is not None else 0
    
    @property
    def ammo(self):
        return self.weapon.ammo if self.weapon is not None else 0
    
    @property
    def action_space(self):
        return self._action_space
    
    @property
    def has_supply(self):
        return self._has_supply

    @property
    def available_types(self):
        return self._available_types
    
    @property
    def parked_agents(self):
        return self._parked_agents
    
    @property
    def capacity(self):
        return self._capacity
    
    @property
    def state(self):
        return AgentState(self.pos, self.endurance, self.fuel, self.weapon, self.commenced_action, self.cmd_todo, self.todo)
    
    def update(self, state: 'AgentState'):
        self.pos = state.pos
        self.endurance = state.endurance
        self.fuel = state.fuel
        self.commenced_action = state.commenced_action
        self.cmd_todo = copy.deepcopy(state.cmd_todo)
        self.todo = copy.deepcopy(state.todo)
        self.weapon = copy.deepcopy(state.weapon)
    
class AgentState:
    def __init__(self, pos: Tuple[int, int], endurance: int, fuel: float, weapon: Weapon = None, commenced_action: bool = False, cmd_todo: List[Action] = [], todo: List[Action] = []):
        self.pos = pos
        self.endurance = endurance
        self.fuel = fuel
        self.commenced_action = commenced_action
        self.weapon = copy.deepcopy(weapon) if weapon is not None else None
        self.cmd_todo = copy.deepcopy(cmd_todo)
        self.todo = copy.deepcopy(todo)
    
class Team:
    def __init__(self, team_id: int, faction_id: int, agents: dict[str, Agent] = {}):
        self.team_id = team_id
        self.faction_id = faction_id
        self.agents = agents
        self.reset()
    
    def __lt__(self, other):
        return self.team_id < other.team_id
    
    def __eq__(self, other):
        return self.team_id == other.team_id
    
    def __hash__(self):
        return hash(self.team_id)
    
    def __str__(self):
        return f"Team {self.team_id}"
    
    def __repr__(self):
        return f"Team {self.team_id}"
    
    def add_agent(self, agent: Agent):
        self.agents.append(agent)
    
    def remove_agent(self, agent_id: str):
        self.agents.pop(agent_id)
    
    def reset(self):
        for agent in self.agents.values():
            agent.reset()

    @property
    def alive_count(self):
        return sum(1 for agent in self.agents.values() if agent.alive)
    
    @property
    def alive_ids(self):
        return [agent.agent_id for agent in self.agents.values() if agent.alive]
    
class TileNode:
    def __init__(self, pos: Tuple[int, int], terrain_type: TerrainType, agent_id: str = None, is_city: bool = False):
        self.pos = pos
        self.terrain_type = TerrainType(terrain_type)
        self.agent_id = agent_id
        self.is_city = is_city
        self.team_id = -1
        self.faction_id = -1
        
    def reset(self):
        self.chaotic_value = 0
        self.occupy_value = 0
        self.occupied_by = None
        self.agent_id = None
        self.team_id = -1
        self.faction_id = -1

    def __lt__(self, other):
        return self.pos < other.pos
    
    def __eq__(self, other):
        return self.pos == other.pos
    
    def __hash__(self):
        return hash(self.pos)
    
    def __str__(self):
        return f"{self.pos} ({self.terrain_type.name})"
    
    def __repr__(self):
        return f"{self.pos} ({self.terrain_type.name})"


class Map:
    def __init__(self, nodes: dict[Tuple[int, int], TileNode]):
        self.width = max(pos[0] for pos in nodes.keys()) + 1
        self.height = max(pos[1] for pos in nodes.keys()) + 1
        self.nodes = nodes

class Module(ABC):
    def __init__(self, module_type: ModuleType, add_endurance: int, add_ammo: int, add_fuel: float, available_types: List[AgentType] = None, capacity: int = 0):
        self.module_type = module_type
        self.add_endurance = add_endurance
        self.add_ammo = add_ammo
        self.add_fuel = add_fuel
        self.available_types = available_types if available_types is not None else []
        self.capacity = capacity

    def to_dict(self):
        return {
            "module_type": self.module_type.name,
            "add_endurance": self.add_endurance,
            "add_ammo": self.add_ammo,
            "add_fuel": self.add_fuel,
            "available_types": [agent_type.name for agent_type in self.available_types]
        }

class Hangar(Module):
    def __init__(self, available_types: List[AgentType] = None, capacity: int = 6):
        super().__init__(module_type=ModuleType.HANGAR, 
                         add_endurance=4, 
                         add_ammo=-1, 
                         add_fuel=-1, 
                         available_types=available_types,
                         capacity=capacity)
        self.parked_agents = []

    def reset(self):
        self.parked_agents = []

class Supply(Module):
    def __init__(self, available_types: List[AgentType] = None, capacity: int = 0):
        super().__init__(module_type=ModuleType.SUPPLY, 
                         add_endurance=2, 
                         add_ammo=2, 
                         add_fuel=-1, 
                         available_types=available_types,
                         capacity=capacity)

class Transport(Module):
    def __init__(self, available_types: List[AgentType] = None, capacity: int = 6):
        super().__init__(module_type=ModuleType.TRANSPORT, 
                         add_endurance=0, 
                         add_ammo=0, 
                         add_fuel=0, 
                         available_types=available_types,
                         capacity=capacity)
        self.parked_agents = []

    def reset(self):
        self.parked_agents = []