components/main/hooks/useInputGeneration.ts

import { useState } from "react"
import { useMutation, useQuery, useQueryClient } from "@tanstack/react-query"
import { useLocalStorage } from 'react-use';

import { Collection } from "@/type"
import list_styles from "@/assets/list_styles.json"

export const useInputGeneration = () => {
  const [myGenerationsId, setGenerationsId] = useLocalStorage<any>('my-own-generations', []);
  const [style, setStyle] = useState<string>(list_styles[0].name)

  const { data: prompt } = useQuery(["prompt"], () => {
    return ''
  }, {
    refetchOnWindowFocus: false,
    refetchOnMount: false,
    refetchOnReconnect: false,
    initialData: ''
  })
  const setPrompt = (str:string) => client.setQueryData(["prompt"], () => str)

  const client = useQueryClient()

  const { mutate: submit, isLoading: loading } = useMutation(
    ["generation"],
    async () => {
      if (!hasMadeFirstGeneration) setFirstGenerationDone()
      client.setQueryData(["collections"], (old) => {
        return [{
          id: -1,
          blob: {
            type: "image/png",
            data: new ArrayBuffer(0),
          },
          prompt
        }, ...old as Collection[]]
      })

      const findStyle = list_styles.find((item) => item.name === style)

      const response = await fetch("/api", {
        method: "POST",
        body: JSON.stringify({
          inputs: findStyle?.prompt.replace("{prompt}", prompt) ?? prompt,
          negative_prompt: findStyle?.negative_prompt ?? "",
        }),
      })
      const data = await response.json()

      if (!response.ok) {
        throw new Error(data.message)
      }

      client.setQueryData(["collections"], (old) => {
        const newArray = [...old as Collection[]]
        const index = newArray.findIndex((item: Collection) => item.id === -1)

        newArray[index] = data?.blob as Collection

        return newArray
      })

      setGenerationsId(myGenerationsId?.length ? [...myGenerationsId, data?.blob?.id] : [data?.blob?.id])

      return data ?? {}
    },
  )

  const { data: hasMadeFirstGeneration } = useQuery(["firstGenerationDone"], () => {
    return false
  }, {
    refetchOnWindowFocus: false,
    refetchOnMount: false,
    refetchOnReconnect: false,
    initialData: false
  })
  const setFirstGenerationDone = () => client.setQueryData(["firstGenerationDone"], () => true)

  return {
    prompt,
    setPrompt,
    loading,
    submit,
    hasMadeFirstGeneration,
    list_styles,
    style,
    setStyle
  }

}