from datetime import datetime
from typing import Optional

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
    user_id: int

    @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
    user_id: int


class FlightStatusUpdate(BaseModel):
    status: str
    user_id: int


class FlightUpdate(BaseModel):
    flight_code: Optional[str] = None
    status: Optional[str] = None
    origin: Optional[str] = None
    destination: Optional[str] = None
    departure_time: Optional[str] = None
    arrival_time: Optional[str] = None
    gate: Optional[str] = None
    user_id: int