from datetime import datetime

from pydantic import BaseModel, validator


class Flight(BaseModel):
    id: int
    flight_code: str
    status: str
    origin: str
    destination: str
    departure_time: str
    arrival_time: str
    gate: str = None
    # last_updated: str

    # @validator("departure_time", "arrival_time", "last_updated", pre=True, always=True)
    @validator("departure_time", "arrival_time", pre=True, always=True)
    def parse_datetime(cls, value):
        if isinstance(value, datetime):
            return value.strftime("%Y-%m-%d %I:%M %p")
        return value


class FlightCreate(BaseModel):
    flight_code: str
    status: str
    origin: str
    destination: str
    departure_time: str
    arrival_time: str
    gate: str = None


class FlightStatusUpdate(BaseModel):
    status: str