import { z } from "zod"

import { citySchema } from "@/server/routers/hotels/schemas/city"
import { nearbyHotelsSchema } from "@/server/routers/hotels/schemas/hotel/include/nearbyHotels"
import { restaurantsSchema } from "@/server/routers/hotels/schemas/hotel/include/restaurants"
import {
  roomCategoriesSchema,
  transformRoomCategories,
} from "@/server/routers/hotels/schemas/hotel/include/roomCategories"

import { additionalDataSchema, transformAdditionalData } from "./additionalData"

export const includeSchema = z
  .union([
    z.null(),
    z.undefined(),
    z.discriminatedUnion("type", [
      additionalDataSchema,
      citySchema,
      nearbyHotelsSchema,
      restaurantsSchema,
      roomCategoriesSchema,
    ]),
  ])
  .transform((data) => {
    switch (data?.type) {
      case "additionalData":
        return transformAdditionalData(data)
      case "cities":
      case "hotels":
      case "restaurants":
        return {
          ...data.attributes,
          id: data.id,
          type: data.type,
        }
      case "roomcategories":
        return transformRoomCategories(data)
      default:
        return null
    }
  })
