# Unit-Type Mapping für C&C Remastered
# Konvertiert interne Unit-IDs zu lesbaren Namen und Fraktionen

UNIT_MAPPING = {
    # Allies Units
    "rifle": {
        "display_name": "Rifle Infantry",
        "faction": "Allies",
        "category": "infantry"
    },
    "medic": {
        "display_name": "Medic",
        "faction": "Allies",
        "category": "infantry"
    },
    "engineer": {
        "display_name": "Engineer",
        "faction": "Allies",
        "category": "infantry"
    },
    "spy": {
        "display_name": "Spy",
        "faction": "Allies",
        "category": "infantry"
    },
    "thief": {
        "display_name": "Thief",
        "faction": "Allies",
        "category": "infantry"
    },
    "light_tank": {
        "display_name": "Light Tank",
        "faction": "Allies",
        "category": "vehicle"
    },
    "medium_tank": {
        "display_name": "Medium Tank",
        "faction": "Allies",
        "category": "vehicle"
    },
    "apc": {
        "display_name": "APC",
        "faction": "Allies",
        "category": "vehicle"
    },
    "artillery": {
        "display_name": "Artillery",
        "faction": "Allies",
        "category": "vehicle"
    },
    "mobile_artillery": {
        "display_name": "Mobile Artillery",
        "faction": "Allies",
        "category": "vehicle"
    },
    "ore_truck": {
        "display_name": "Ore Truck",
        "faction": "Allies",
        "category": "vehicle"
    },
    "ranger": {
        "display_name": "Ranger",
        "faction": "Allies",
        "category": "vehicle"
    },
    "mobile_gap": {
        "display_name": "Mobile Gap Generator",
        "faction": "Allies",
        "category": "vehicle"
    },
    
    # Soviet Units
    "flamethrower": {
        "display_name": "Flamethrower Infantry",
        "faction": "Soviet",
        "category": "infantry"
    },
    "grenadier": {
        "display_name": "Grenadier",
        "faction": "Soviet",
        "category": "infantry"
    },
    "heavy_tank": {
        "display_name": "Heavy Tank",
        "faction": "Soviet",
        "category": "vehicle"
    },
    "mammoth_tank": {
        "display_name": "Mammoth Tank",
        "faction": "Soviet",
        "category": "vehicle"
    },
    "v2_rocket": {
        "display_name": "V2 Rocket Launcher",
        "faction": "Soviet",
        "category": "vehicle"
    },
    "tesla_tank": {
        "display_name": "Tesla Tank",
        "faction": "Soviet",
        "category": "vehicle"
    },
    
    # Air Units
    "longbow": {
        "display_name": "Longbow Helicopter",
        "faction": "Allies",
        "category": "air"
    },
    "hind": {
        "display_name": "Hind Helicopter",
        "faction": "Soviet",
        "category": "air"
    },
    "yak": {
        "display_name": "Yak Aircraft",
        "faction": "Soviet",
        "category": "air"
    },
    "mig": {
        "display_name": "MiG Aircraft",
        "faction": "Soviet",
        "category": "air"
    },
    
    # Naval Units
    "gunboat": {
        "display_name": "Gunboat",
        "faction": "Allies",
        "category": "naval"
    },
    "destroyer": {
        "display_name": "Destroyer",
        "faction": "Allies",
        "category": "naval"
    },
    "cruiser": {
        "display_name": "Cruiser",
        "faction": "Allies",
        "category": "naval"
    },
    "transport": {
        "display_name": "Transport",
        "faction": "Allies",
        "category": "naval"
    },
    "submarine": {
        "display_name": "Submarine",
        "faction": "Soviet",
        "category": "naval"
    }
}


def get_unit_info(unit_type: str) -> dict:
    """Gibt Unit-Info zurück, oder Default-Werte wenn unbekannt"""
    return UNIT_MAPPING.get(unit_type, {
        "display_name": unit_type.title(),
        "faction": "Unknown",
        "category": "unknown"
    })


def get_units_by_faction(faction: str) -> list:
    """Gibt alle Units einer Fraktion zurück"""
    return [unit_id for unit_id, info in UNIT_MAPPING.items() 
            if info["faction"] == faction]


def get_units_by_category(category: str) -> list:
    """Gibt alle Units einer Kategorie zurück"""
    return [unit_id for unit_id, info in UNIT_MAPPING.items() 
            if info["category"] == category]


if __name__ == "__main__":
    print("=== Unit Mapping Test ===")
    print(f"Anzahl bekannte Units: {len(UNIT_MAPPING)}")
    print(f"Allies Units: {len(get_units_by_faction('Allies'))}")
    print(f"Soviet Units: {len(get_units_by_faction('Soviet'))}")
    print(f"Vehicle Units: {len(get_units_by_category('vehicle'))}")
    
    # Test
    print("\nBeispiel:")
    print(get_unit_info("medium_tank"))
    print(get_unit_info("unknown_unit"))