import { z } from "zod"

import { discriminatedUnionArray } from "@/lib/discriminatedUnion"

import {
  cardGridRefsSchema,
  cardsGridSchema,
} from "../schemas/blocks/cardsGrid"
import {
  contentRefsSchema as blockContentRefsSchema,
  contentSchema as blockContentSchema,
} from "../schemas/blocks/content"
import {
  dynamicContentRefsSchema,
  dynamicContentSchema as blockDynamicContentSchema,
} from "../schemas/blocks/dynamicContent"
import {
  shortcutsRefsSchema,
  shortcutsSchema,
} from "../schemas/blocks/shortcuts"
import { tempImageVaultAssetSchema } from "../schemas/imageVault"
import {
  contentRefsSchema as sidebarContentRefsSchema,
  contentSchema as sidebarContentSchema,
} from "../schemas/sidebar/content"
import { dynamicContentSchema as sidebarDynamicContentSchema } from "../schemas/sidebar/dynamicContent"
import {
  joinLoyaltyContactRefsSchema,
  joinLoyaltyContactSchema,
} from "../schemas/sidebar/joinLoyaltyContact"
import { systemSchema } from "../schemas/system"

import { LoyaltyPageEnum } from "@/types/enums/loyaltyPage"

// LoyaltyPage Refs
const extendedCardGridRefsSchema = z
  .object({
    __typename: z.literal(LoyaltyPageEnum.ContentStack.blocks.CardsGrid),
  })
  .merge(cardGridRefsSchema)

const extendedContentRefsSchema = z
  .object({
    __typename: z.literal(LoyaltyPageEnum.ContentStack.blocks.Content),
  })
  .merge(blockContentRefsSchema)

const extendedDynamicContentRefsSchema = z
  .object({
    __typename: z.literal(LoyaltyPageEnum.ContentStack.blocks.DynamicContent),
  })
  .merge(dynamicContentRefsSchema)

const extendedShortcutsRefsSchema = z
  .object({
    __typename: z.literal(LoyaltyPageEnum.ContentStack.blocks.Shortcuts),
  })
  .merge(shortcutsRefsSchema)

const blocksRefsSchema = z.discriminatedUnion("__typename", [
  extendedCardGridRefsSchema,
  extendedContentRefsSchema,
  extendedDynamicContentRefsSchema,
  extendedShortcutsRefsSchema,
])

const contentSidebarRefsSchema = z
  .object({
    __typename: z.literal(LoyaltyPageEnum.ContentStack.sidebar.Content),
  })
  .merge(sidebarContentRefsSchema)

const extendedJoinLoyaltyContactRefsSchema = z
  .object({
    __typename: z.literal(
      LoyaltyPageEnum.ContentStack.sidebar.JoinLoyaltyContact
    ),
  })
  .merge(joinLoyaltyContactRefsSchema)

const sidebarRefsSchema = z.discriminatedUnion("__typename", [
  contentSidebarRefsSchema,
  extendedJoinLoyaltyContactRefsSchema,
  z.object({
    __typename: z.literal(LoyaltyPageEnum.ContentStack.sidebar.DynamicContent),
  }),
])

export const loyaltyPageRefsSchema = z.object({
  loyalty_page: z.object({
    blocks: discriminatedUnionArray(blocksRefsSchema.options).optional(),
    sidebar: discriminatedUnionArray(sidebarRefsSchema.options)
      .optional()
      .transform((data) => {
        if (data) {
          return data.filter(
            (block) =>
              block.__typename !==
              LoyaltyPageEnum.ContentStack.sidebar.DynamicContent
          )
        }
        return data
      }),
    system: systemSchema,
  }),
})

// LoyaltyPage
export const extendedCardsGridSchema = z
  .object({
    __typename: z.literal(LoyaltyPageEnum.ContentStack.blocks.CardsGrid),
  })
  .merge(cardsGridSchema)

export const extendedContentSchema = z
  .object({
    __typename: z.literal(LoyaltyPageEnum.ContentStack.blocks.Content),
  })
  .merge(blockContentSchema)

export const extendedDynamicContentSchema = z
  .object({
    __typename: z.literal(LoyaltyPageEnum.ContentStack.blocks.DynamicContent),
  })
  .merge(blockDynamicContentSchema)

export const extendedShortcutsSchema = z
  .object({
    __typename: z.literal(LoyaltyPageEnum.ContentStack.blocks.Shortcuts),
  })
  .merge(shortcutsSchema)

export const blocksSchema = z.discriminatedUnion("__typename", [
  extendedCardsGridSchema,
  extendedContentSchema,
  extendedDynamicContentSchema,
  extendedShortcutsSchema,
])

const contentSidebarSchema = z
  .object({
    __typename: z.literal(LoyaltyPageEnum.ContentStack.sidebar.Content),
  })
  .merge(sidebarContentSchema)

const dynamicContentSidebarSchema = z
  .object({
    __typename: z.literal(LoyaltyPageEnum.ContentStack.sidebar.DynamicContent),
  })
  .merge(sidebarDynamicContentSchema)

export const joinLoyaltyContactSidebarSchema = z
  .object({
    __typename: z.literal(
      LoyaltyPageEnum.ContentStack.sidebar.JoinLoyaltyContact
    ),
  })
  .merge(joinLoyaltyContactSchema)

export const sidebarSchema = z.discriminatedUnion("__typename", [
  contentSidebarSchema,
  dynamicContentSidebarSchema,
  joinLoyaltyContactSidebarSchema,
])

export const loyaltyPageSchema = z.object({
  loyalty_page: z
    .object({
      blocks: discriminatedUnionArray(blocksSchema.options).nullable(),
      heading: z.string().optional(),
      hero_image: tempImageVaultAssetSchema,
      preamble: z.string().optional(),
      sidebar: discriminatedUnionArray(sidebarSchema.options).nullable(),
      title: z.string().optional(),
      system: systemSchema.merge(
        z.object({
          created_at: z.string(),
          updated_at: z.string(),
        })
      ),
    })
    .transform((data) => {
      return {
        blocks: data.blocks ? data.blocks : [],
        heading: data.heading,
        heroImage: data.hero_image,
        preamble: data.preamble,
        sidebar: data.sidebar ? data.sidebar : [],
        system: data.system,
      }
    }),
})