import { useMutation } from '@tanstack/react-query';
import { omit } from 'lodash-es';
import {
  createContext,
  useCallback,
  useEffect,
  useMemo,
  useState,
} from 'react';

import useClip from '@/hooks/useClip';
import useGenerate, {
  GenerateParams,
  ModelTier,
  getClipModelTier,
} from '@/hooks/useGenerate';
import { TransactionLogger } from '@/logging/logWebUserEvent';

export const useGenerateContext = (clipId: string) => {
  const [generateCount, setGenerateCount] = useState(0);

  const { clip } = useClip(clipId);

  const [modelOverride, setModelOverride] = useState<string | null>(null);
  const [modelTier, setModelTier] = useState<ModelTier>(() =>
    clip ? getClipModelTier(clip, false, true) : ModelTier.V4_5_PLUS
  );

  useEffect(() => {
    if (clip) {
      setModelTier(getClipModelTier(clip, false));
    }
  }, [clip?.major_model_version]);

  const incrementGenerateCount = useCallback(() => {
    setGenerateCount((prev) => prev + 1);
  }, [setGenerateCount]);

  const generate = useGenerate();

  const generateAndIncrement = useCallback(
    async (
      params: GenerateParams & { transactionLogger: TransactionLogger }
    ) => {
      const result = await generate(
        params.transactionLogger,
        omit(params, 'transactionLogger')
      );
      incrementGenerateCount();
      return result;
    },
    [incrementGenerateCount, generate]
  );

  const generateMutation = useMutation({
    mutationFn: generateAndIncrement,
  });

  return useMemo(
    () => ({
      generateCount,
      incrementGenerateCount,
      generateMutation,
      modelTier,
      setModelTier,
      modelOverride,
      setModelOverride,
    }),
    [
      generateCount,
      generateMutation,
      modelTier,
      setModelTier,
      modelOverride,
      setModelOverride,
    ]
  );
};

const GenerateContext = createContext<ReturnType<typeof useGenerateContext>>(
  undefined as never
);

export default GenerateContext;