Airline Management

  1  from __future__ import annotations
  2  from abc import ABC, abstractmethod
  3  from dataclasses import dataclass, field
  4  from datetime import datetime, timedelta
  5  from enum import Enum
  6  from threading import RLock
  7  from typing import Callable
  8  import uuid
  9  
 10  
 11  class SeatClass(Enum):
 12      ECONOMY = "economy"
 13      BUSINESS = "business"
 14      FIRST = "first"
 15  
 16  
 17  class FlightStatus(Enum):
 18      SCHEDULED = "scheduled"
 19      BOARDING = "boarding"
 20      DEPARTED = "departed"
 21      ARRIVED = "arrived"
 22      CANCELLED = "cancelled"
 23  
 24  
 25  class BookingStatus(Enum):
 26      HELD = "held"
 27      CONFIRMED = "confirmed"
 28      CHECKED_IN = "checked_in"
 29      CANCELLED = "cancelled"
 30      EXPIRED = "expired"
 31  
 32  
 33  @dataclass(frozen=True)
 34  class Seat:
 35      number: str
 36      seat_class: SeatClass
 37      is_exit_row: bool = False
 38      is_aisle: bool = False
 39  
 40  
 41  @dataclass(frozen=True)
 42  class Money:
 43      amount: int    # in cents; avoid float
 44      currency: str = "USD"
 45  
 46      def __add__(self, other: "Money") -> "Money":
 47          if self.currency != other.currency:
 48              raise ValueError("currency mismatch")
 49          return Money(self.amount + other.amount, self.currency)
 50  
 51  
 52  @dataclass
 53  class Aircraft:
 54      model: str
 55      seat_map: list[Seat]
 56  
 57  
 58  @dataclass
 59  class Passenger:
 60      id: str
 61      name: str
 62      document_id: str
 63  
 64  
 65  class Flight:
 66      def __init__(self, flight_id: str, origin: str, destination: str,
 67                   departure: datetime, arrival: datetime, aircraft: Aircraft):
 68          if arrival <= departure:
 69              raise ValueError("arrival must be after departure")
 70          self.id = flight_id
 71          self.origin = origin
 72          self.destination = destination
 73          self.departure = departure
 74          self.arrival = arrival
 75          self.aircraft = aircraft
 76          self.status = FlightStatus.SCHEDULED
 77          # seat_number -> booking_id (or "HELD:token" during hold)
 78          self._assignments: dict[str, str] = {}
 79          self._lock = RLock()
 80  
 81      def available_seats(self, seat_class: SeatClass | None = None) -> list[Seat]:
 82          with self._lock:
 83              return [
 84                  s for s in self.aircraft.seat_map
 85                  if s.number not in self._assignments
 86                  and (seat_class is None or s.seat_class == seat_class)
 87              ]
 88  
 89      def seat_by_number(self, number: str) -> Seat | None:
 90          return next((s for s in self.aircraft.seat_map if s.number == number), None)
 91  
 92      def try_hold(self, seat_number: str, token: str) -> bool:
 93          """Reserve a seat atomically with a hold marker. Returns False if taken."""
 94          with self._lock:
 95              if seat_number in self._assignments:
 96                  return False
 97              if self.seat_by_number(seat_number) is None:
 98                  return False
 99              self._assignments[seat_number] = f"HELD:{token}"
100              return True
101  
102      def release_hold(self, seat_number: str, token: str) -> None:
103          with self._lock:
104              current = self._assignments.get(seat_number)
105              if current == f"HELD:{token}":
106                  del self._assignments[seat_number]
107  
108      def confirm_hold(self, seat_number: str, token: str, booking_id: str) -> bool:
109          with self._lock:
110              if self._assignments.get(seat_number) == f"HELD:{token}":
111                  self._assignments[seat_number] = booking_id
112                  return True
113              return False
114  
115      def free_seat(self, seat_number: str) -> None:
116          with self._lock:
117              self._assignments.pop(seat_number, None)
118  
119  
120  class FareStrategy(ABC):
121      @abstractmethod
122      def price(self, flight: Flight, seat: Seat) -> Money: ...
123  
124  
125  class ClassBasedFare(FareStrategy):
126      DEFAULTS = {
127          SeatClass.ECONOMY: 15_000,
128          SeatClass.BUSINESS: 50_000,
129          SeatClass.FIRST: 120_000,
130      }
131  
132      def price(self, flight: Flight, seat: Seat) -> Money:
133          base = self.DEFAULTS[seat.seat_class]
134          if seat.is_exit_row:
135              base += 3_000   # extra legroom surcharge
136          return Money(base)
137  
138  
139  @dataclass
140  class Hold:
141      token: str
142      flight_id: str
143      seat_number: str
144      expires_at: datetime
145  
146  
147  class SeatHoldManager:
148      """Short-lived holds during payment. Expired holds are reclaimed lazily."""
149  
150      def __init__(self, default_ttl: timedelta = timedelta(minutes=10),
151                   clock: Callable[[], datetime] = datetime.utcnow):
152          self._holds: dict[str, Hold] = {}
153          self._default_ttl = default_ttl
154          self._now = clock
155          self._lock = RLock()
156  
157      def hold(self, flight: Flight, seat_number: str,
158               ttl: timedelta | None = None) -> str:
159          token = str(uuid.uuid4())[:8]
160          if not flight.try_hold(seat_number, token):
161              raise RuntimeError(f"Seat {seat_number} unavailable on flight {flight.id}")
162          expires = self._now() + (ttl or self._default_ttl)
163          with self._lock:
164              self._holds[token] = Hold(token, flight.id, seat_number, expires)
165          return token
166  
167      def release(self, token: str, flight: Flight) -> None:
168          with self._lock:
169              hold = self._holds.pop(token, None)
170          if hold:
171              flight.release_hold(hold.seat_number, token)
172  
173      def consume(self, token: str, flight: Flight, booking_id: str) -> bool:
174          with self._lock:
175              hold = self._holds.get(token)
176              if hold is None:
177                  return False
178              if self._now() >= hold.expires_at:
179                  flight.release_hold(hold.seat_number, token)
180                  del self._holds[token]
181                  return False
182              ok = flight.confirm_hold(hold.seat_number, token, booking_id)
183              if ok:
184                  del self._holds[token]
185              return ok
186  
187      def sweep_expired(self, flights_by_id: dict[str, Flight]) -> int:
188          now = self._now()
189          expired_tokens = []
190          with self._lock:
191              for t, h in list(self._holds.items()):
192                  if now >= h.expires_at:
193                      expired_tokens.append(h)
194                      del self._holds[t]
195          for hold in expired_tokens:
196              f = flights_by_id.get(hold.flight_id)
197              if f:
198                  f.release_hold(hold.seat_number, hold.token)
199          return len(expired_tokens)
200  
201  
202  @dataclass
203  class Booking:
204      id: str
205      flight_id: str
206      passenger: Passenger
207      seat_number: str
208      fare: Money
209      status: BookingStatus
210  
211  
212  @dataclass
213  class BookingRequest:
214      flight: Flight
215      passenger: Passenger
216      seat_number: str
217  
218  
219  class AirlineService:
220      def __init__(self, fare: FareStrategy | None = None):
221          self._flights: dict[str, Flight] = {}
222          self._bookings: dict[str, Booking] = {}
223          self._fare = fare or ClassBasedFare()
224          self._holds = SeatHoldManager()
225  
226      def register_flight(self, flight: Flight) -> None:
227          self._flights[flight.id] = flight
228  
229      def search(self, origin: str, destination: str, on_date: datetime) -> list[Flight]:
230          return [
231              f for f in self._flights.values()
232              if f.origin == origin and f.destination == destination
233              and f.departure.date() == on_date.date()
234              and f.status == FlightStatus.SCHEDULED
235          ]
236  
237      def hold_seat(self, flight: Flight, seat_number: str, ttl: timedelta | None = None) -> str:
238          if flight.status != FlightStatus.SCHEDULED:
239              raise RuntimeError("Flight not open for booking")
240          return self._holds.hold(flight, seat_number, ttl=ttl)
241  
242      def sweep_expired_holds(self) -> int:
243          """Run on a short cadence (seconds) to release TTL-expired holds."""
244          return self._holds.sweep_expired(self._flights)
245  
246      def confirm_booking(self, token: str, flight: Flight, passenger: Passenger,
247                          payment_succeeded: bool) -> Booking:
248          if not payment_succeeded:
249              self._holds.release(token, flight)
250              raise RuntimeError("Payment failed; hold released")
251  
252          seat = None
253          for s in flight.aircraft.seat_map:
254              if flight._assignments.get(s.number) == f"HELD:{token}":
255                  seat = s
256                  break
257          if seat is None:
258              raise RuntimeError("Hold expired or invalid")
259  
260          booking_id = str(uuid.uuid4())[:8]
261          if not self._holds.consume(token, flight, booking_id):
262              raise RuntimeError("Hold no longer valid")
263  
264          fare = self._fare.price(flight, seat)
265          booking = Booking(
266              id=booking_id,
267              flight_id=flight.id,
268              passenger=passenger,
269              seat_number=seat.number,
270              fare=fare,
271              status=BookingStatus.CONFIRMED,
272          )
273          self._bookings[booking_id] = booking
274          return booking
275  
276      def cancel(self, booking_id: str) -> None:
277          booking = self._bookings.get(booking_id)
278          if booking is None:
279              raise ValueError("Unknown booking")
280          if booking.status in (BookingStatus.CANCELLED, BookingStatus.EXPIRED):
281              return
282          booking.status = BookingStatus.CANCELLED
283          flight = self._flights[booking.flight_id]
284          flight.free_seat(booking.seat_number)
285  
286      def check_in(self, booking_id: str) -> None:
287          booking = self._bookings.get(booking_id)
288          if booking is None:
289              raise ValueError("Unknown booking")
290          if booking.status != BookingStatus.CONFIRMED:
291              raise RuntimeError(f"Cannot check in a {booking.status.value} booking")
292          booking.status = BookingStatus.CHECKED_IN
293  
294  
295  def _make_737_seatmap() -> list[Seat]:
296      seats = []
297      for row in range(1, 5):
298          for col in "ABCDEF":
299              seats.append(Seat(f"{row}{col}", SeatClass.BUSINESS, is_aisle=(col in "CD")))
300      for row in range(5, 31):
301          for col in "ABCDEF":
302              seats.append(Seat(f"{row}{col}", SeatClass.ECONOMY,
303                                is_exit_row=(row in (12, 13)),
304                                is_aisle=(col in "CD")))
305      return seats
306  
307  
308  if __name__ == "__main__":
309      airline = AirlineService()
310      aircraft = Aircraft("Boeing 737", _make_737_seatmap())
311      flight = Flight("UA-123", "SFO", "JFK",
312                      datetime(2026, 5, 1, 7, 0),
313                      datetime(2026, 5, 1, 15, 30),
314                      aircraft)
315      airline.register_flight(flight)
316  
317      print("Business class available:", len(flight.available_seats(SeatClass.BUSINESS)))
318  
319      passenger = Passenger("p1", "Rakesh Rathi", "X123456")
320      token = airline.hold_seat(flight, "2A")
321      print(f"Hold token: {token}")
322  
323      # Any concurrent hold attempt on the same seat must fail.
324      try:
325          airline.hold_seat(flight, "2A")
326      except RuntimeError as e:
327          print(f"Second hold correctly rejected: {e}")
328  
329      booking = airline.confirm_booking(token, flight, passenger, payment_succeeded=True)
330      print(f"Booked {booking.id} -> seat {booking.seat_number}, fare ${booking.fare.amount / 100:.2f}")
331  
332      airline.check_in(booking.id)
333      print(f"Status: {booking.status.value}")
334  
335      airline.cancel(booking.id)
336      print(f"Cancelled. Seat 2A available again: {'2A' in [s.number for s in flight.available_seats()]}")
337  
338      # Short-TTL hold gets reclaimed by the sweeper.
339      import time as _time
340      short_token = airline.hold_seat(flight, "3A", ttl=timedelta(milliseconds=5))
341      _time.sleep(0.02)
342      reclaimed = airline.sweep_expired_holds()
343      print(f"Sweeper reclaimed {reclaimed} hold(s); 3A available: "
344            f"{'3A' in [s.number for s in flight.available_seats()]}")