from __future__ import annotations

from enum import Enum

from planemapper.models import Aircraft


class AircraftType(Enum):
    GA_LIGHT = "ga_light"
    COMMERCIAL = "commercial"
    PRIVATE_JET = "private_jet"
    AIRLINER = "airliner"
    HELICOPTER = "helicopter"
    MILITARY = "military"
    UNKNOWN = "unknown"


_AIRLINE_PREFIXES: frozenset[str] = frozenset(
    {
        "BAW",
        "EIN",
        "RYR",
        "EZY",
        "THY",
        "DLH",
        "AFR",
        "IBE",
        "KLM",
        "UAE",
        "SWR",
        "AAL",
        "UAL",
        "DAL",
        "SAS",
        "TAP",
        "VLG",
        "NOS",
        "WZZ",
        "AEA",
        "NAX",
        "FIN",
        "CSN",
        "CCA",
    }
)

_CATEGORY_MAP: dict[str, AircraftType] = {
    "A1": AircraftType.GA_LIGHT,
    "A2": AircraftType.GA_LIGHT,
    "A3": AircraftType.COMMERCIAL,
    "A4": AircraftType.COMMERCIAL,
    "A5": AircraftType.COMMERCIAL,
    "A7": AircraftType.HELICOPTER,
    "B1": AircraftType.MILITARY,
    "B2": AircraftType.MILITARY,
    "B3": AircraftType.MILITARY,
    "B4": AircraftType.MILITARY,
}


def classify_aircraft_type(aircraft: Aircraft) -> AircraftType:
    if aircraft.category and aircraft.category in _CATEGORY_MAP:
        return _CATEGORY_MAP[aircraft.category]
    callsign = aircraft.callsign.strip()
    if len(callsign) >= 3 and callsign[:3].upper() in _AIRLINE_PREFIXES:
        return AircraftType.COMMERCIAL
    if aircraft.altitude_ft < 10000:
        return AircraftType.GA_LIGHT
    if aircraft.altitude_ft < 30000:
        return AircraftType.PRIVATE_JET
    return AircraftType.AIRLINER
