/**
 * <Callout type="info">
 *   Usage of this hook assumes some additional setup in your application, for more information
 *   on that see the CoAgents <span className="text-blue-500">[getting started guide](/coagents/quickstart/langgraph)</span>.
 * </Callout>
 * <Frame className="my-12">
 *   <img
 *     src="https://cdn.vn.ai/docs/vn-sdk/images/coagents/SharedStateCoAgents.gif"
 *     alt="CoAgents demonstration"
 *     className="w-auto"
 *   />
 * </Frame>
 *
 * This hook is used to integrate an agent into your application. With its use, you can
 * render and update the state of an agent, allowing for a dynamic and interactive experience.
 * We call these shared state experiences agentic copilots, or CoAgents for short.
 *
 * ## Usage
 *
 * ### Simple Usage
 *
 * ```tsx
 * import { useAiAgent } from "@vn-sdk/react-core";
 *
 * type AgentState = {
 *   count: number;
 * }
 *
 * const agent = useAiAgent<AgentState>({
 *   name: "my-agent",
 *   initialState: {
 *     count: 0,
 *   },
 * });
 *
 * ```
 *
 * `useAiAgent` returns an object with the following properties:
 *
 * ```tsx
 * const {
 *   name,     // The name of the agent currently being used.
 *   nodeName, // The name of the current LangGraph node.
 *   state,    // The current state of the agent.
 *   setState, // A function to update the state of the agent.
 *   running,  // A boolean indicating if the agent is currently running.
 *   start,    // A function to start the agent.
 *   stop,     // A function to stop the agent.
 *   run,      // A function to re-run the agent. Takes a HintFunction to inform the agent why it is being re-run.
 * } = agent;
 * ```
 *
 * Finally we can leverage these properties to create reactive experiences with the agent!
 *
 * ```tsx
 * const { state, setState } = useAiAgent<AgentState>({
 *   name: "my-agent",
 *   initialState: {
 *     count: 0,
 *   },
 * });
 *
 * return (
 *   <div>
 *     <p>Count: {state.count}</p>
 *     <button onClick={() => setState({ count: state.count + 1 })}>Increment</button>
 *   </div>
 * );
 * ```
 *
 * This reactivity is bidirectional, meaning that changes to the state from the agent will be reflected in the UI and vice versa.
 *
 * ## Parameters
 * <PropertyReference name="options" type="UseCoagentOptions<T>" required>
 *   The options to use when creating the coagent.
 *   <PropertyReference name="name" type="string" required>
 *     The name of the agent to use.
 *   </PropertyReference>
 *   <PropertyReference name="initialState" type="T | any">
 *     The initial state of the agent.
 *   </PropertyReference>
 *   <PropertyReference name="state" type="T | any">
 *     State to manage externally if you are using this hook with external state management.
 *   </PropertyReference>
 *   <PropertyReference name="setState" type="(newState: T | ((prevState: T | undefined) => T)) => void">
 *     A function to update the state of the agent if you are using this hook with external state management.
 *   </PropertyReference>
 * </PropertyReference>
 */

import { useCallback, useEffect, useMemo, useRef } from "react";
import { AiContextParams, useAiContext } from "../context";
import { AiAgentState } from "../types/ai-agent-state";
import { useAiChat } from "./use-ai-chat-internal";
import { Message } from "@vn-sdk/shared";
import { useAsyncCallback } from "../components/error-boundary/error-utils";
import { useToast } from "../components/toast/toast-provider";
import { useAiRuntimeClient } from "./use-ai-runtime-client";
import { parseJson, AiSDKAgentDiscoveryError } from "@vn-sdk/shared";
import { useMessagesTap } from "../components/ai-provider/ai-messages";
import { Message as GqlMessage } from "@vn-sdk/runtime-client-gql";

interface UseCoagentOptionsBase {
  /**
   * The name of the agent being used.
   */
  name: string;
  /**
   * @deprecated - use "config.configurable"
   * Config to pass to a LangGraph Agent
   */
  configurable?: Record<string, any>;
  /**
   * Config to pass to a LangGraph Agent
   */
  config?: {
    configurable?: Record<string, any>;
    [key: string]: any;
  };
}

interface WithInternalStateManagementAndInitial<T> extends UseCoagentOptionsBase {
  /**
   * The initial state of the agent.
   */
  initialState: T;
}

interface WithInternalStateManagement extends UseCoagentOptionsBase {
  /**
   * Optional initialState with default type any
   */
  initialState?: any;
}

interface WithExternalStateManagement<T> extends UseCoagentOptionsBase {
  /**
   * The current state of the agent.
   */
  state: T;
  /**
   * A function to update the state of the agent.
   */
  setState: (newState: T | ((prevState: T | undefined) => T)) => void;
}

type UseCoagentOptions<T> =
  | WithInternalStateManagementAndInitial<T>
  | WithInternalStateManagement
  | WithExternalStateManagement<T>;

export interface UseCoagentReturnType<T> {
  /**
   * The name of the agent being used.
   */
  name: string;
  /**
   * The name of the current LangGraph node.
   */
  nodeName?: string;
  /**
   * The ID of the thread the agent is running in.
   */
  threadId?: string;
  /**
   * A boolean indicating if the agent is currently running.
   */
  running: boolean;
  /**
   * The current state of the agent.
   */
  state: T;
  /**
   * A function to update the state of the agent.
   */
  setState: (newState: T | ((prevState: T | undefined) => T)) => void;
  /**
   * A function to start the agent.
   */
  start: () => void;
  /**
   * A function to stop the agent.
   */
  stop: () => void;
  /**
   * A function to re-run the agent. The hint function can be used to provide a hint to the agent
   * about why it is being re-run again.
   */
  run: (hint?: HintFunction) => Promise<void>;
}

export interface HintFunctionParams {
  /**
   * The previous state of the agent.
   */
  previousState: any;
  /**
   * The current state of the agent.
   */
  currentState: any;
}

export type HintFunction = (params: HintFunctionParams) => Message | undefined;

/**
 * This hook is used to integrate an agent into your application. With its use, you can
 * render and update the state of the agent, allowing for a dynamic and interactive experience.
 * We call these shared state experiences "agentic copilots". To get started using agentic copilots, which
 * we refer to as CoAgents, checkout the documentation at https://docs.vn.ai/coagents/quickstart/langgraph.
 */
export function useAiAgent<T = any>(options: UseCoagentOptions<T>): UseCoagentReturnType<T> {
  const context = useAiContext();
  const { availableAgents, onError } = context;
  const { setBannerError } = useToast();
  const lastLoadedThreadId = useRef<string>();
  const lastLoadedState = useRef<any>();

  const { name } = options;
  useEffect(() => {
    if (availableAgents?.length && !availableAgents.some((a) => a.name === name)) {
      const message = `(useAiAgent): Agent "${name}" not found. Make sure the agent exists and is properly configured.`;
      console.warn(message);

      // Route to banner instead of toast for consistency
      const agentError = new AiSDKAgentDiscoveryError({
        agentName: name,
        availableAgents: availableAgents.map((a) => ({ name: a.name, id: a.id })),
      });
      setBannerError(agentError);
    }
  }, [availableAgents]);

  const { getMessagesFromTap } = useMessagesTap();

  const { aiAgentStates, aiAgentStatesRef, setAiAgentStatesWithRef, threadId, aiApiConfig } =
    context;
  const { sendMessage, runChatCompletion } = useAiChat();
  const headers = {
    ...(aiApiConfig.headers || {}),
  };

  const runtimeClient = useAiRuntimeClient({
    url: aiApiConfig.chatApiEndpoint,
    publicApiKey: aiApiConfig.publicApiKey,
    headers,
    credentials: aiApiConfig.credentials,
    showDevConsole: context.showDevConsole,
    onError,
  });

  // if we manage state internally, we need to provide a function to set the state
  const setState = useCallback(
    (newState: T | ((prevState: T | undefined) => T)) => {
      // aiAgentStatesRef.current || {}
      let coagentState: AiAgentState = getAiAgentState({ aiAgentStates, name, options });
      const updatedState =
        typeof newState === "function" ? (newState as Function)(coagentState.state) : newState;

      setAiAgentStatesWithRef({
        ...aiAgentStatesRef.current,
        [name]: {
          ...coagentState,
          state: updatedState,
        },
      });
    },
    [aiAgentStates, name],
  );

  useEffect(() => {
    const fetchAgentState = async () => {
      if (!threadId || threadId === lastLoadedThreadId.current) return;

      const result = await runtimeClient.loadAgentState({
        threadId,
        agentName: name,
      });

      // Runtime client handles errors automatically via handleGQLErrors
      if (result.error) {
        return; // Don't process data on error
      }

      const newState = result.data?.loadAgentState?.state;
      if (newState === lastLoadedState.current) return;

      if (result.data?.loadAgentState?.threadExists && newState && newState != "{}") {
        lastLoadedState.current = newState;
        lastLoadedThreadId.current = threadId;
        const fetchedState = parseJson(newState, {});
        isExternalStateManagement(options)
          ? options.setState(fetchedState)
          : setState(fetchedState);
      }
    };
    void fetchAgentState();
  }, [threadId]);

  // Sync internal state with external state if state management is external
  useEffect(() => {
    if (isExternalStateManagement(options)) {
      setState(options.state);
    } else if (aiAgentStates[name] === undefined) {
      setState(options.initialState === undefined ? {} : options.initialState);
    }
  }, [
    isExternalStateManagement(options) ? JSON.stringify(options.state) : undefined,
    // reset initialstate on reset
    aiAgentStates[name] === undefined,
  ]);

  // Sync config when runtime configuration changes
  useEffect(() => {
    const newConfig = options.config
      ? options.config
      : options.configurable
        ? { configurable: options.configurable }
        : undefined;

    if (newConfig === undefined) return;

    setAiAgentStatesWithRef((prev) => {
      const existing = prev[name] ?? {
        name,
        state: isInternalStateManagementWithInitial(options) ? options.initialState : {},
        config: {},
        running: false,
        active: false,
        threadId: undefined,
        nodeName: undefined,
        runId: undefined,
      };

      if (JSON.stringify(existing.config) === JSON.stringify(newConfig)) {
        return prev;
      }

      return {
        ...prev,
        [name]: {
          ...existing,
          config: newConfig,
        },
      };
    });
  }, [JSON.stringify(options.config), JSON.stringify(options.configurable)]);

  const runAgentCallback = useAsyncCallback(
    async (hint?: HintFunction) => {
      await runAgent(name, context, getMessagesFromTap(), sendMessage, runChatCompletion, hint);
    },
    [name, context, sendMessage, runChatCompletion],
  );

  // Return the state and setState function
  return useMemo(() => {
    const coagentState = getAiAgentState({ aiAgentStates, name, options });
    return {
      name,
      nodeName: coagentState.nodeName,
      threadId: coagentState.threadId,
      running: coagentState.running,
      state: coagentState.state,
      setState: isExternalStateManagement(options) ? options.setState : setState,
      start: () => startAgent(name, context),
      stop: () => stopAgent(name, context),
      run: runAgentCallback,
    };
  }, [name, aiAgentStates, options, setState, runAgentCallback]);
}

export function startAgent(name: string, context: AiContextParams) {
  const { setAgentSession } = context;
  setAgentSession({
    agentName: name,
  });
}

export function stopAgent(name: string, context: AiContextParams) {
  const { agentSession, setAgentSession } = context;
  if (agentSession && agentSession.agentName === name) {
    setAgentSession(null);
    context.setAiAgentStates((prevAgentStates) => {
      return {
        ...prevAgentStates,
        [name]: {
          ...prevAgentStates[name],
          running: false,
          active: false,
          threadId: undefined,
          nodeName: undefined,
          runId: undefined,
        },
      };
    });
  } else {
    console.warn(`No agent session found for ${name}`);
  }
}

export async function runAgent(
  name: string,
  context: AiContextParams,
  messages: GqlMessage[],
  sendMessage: (message: Message) => Promise<void>,
  runChatCompletion: () => Promise<Message[]>,
  hint?: HintFunction,
) {
  const { agentSession, setAgentSession } = context;
  if (!agentSession || agentSession.agentName !== name) {
    setAgentSession({
      agentName: name,
    });
  }

  let previousState: any = null;
  for (let i = messages.length - 1; i >= 0; i--) {
    const message = messages[i];
    if (message.isAgentStateMessage() && message.agentName === name) {
      previousState = message.state;
    }
  }

  let state = context.aiAgentStatesRef.current?.[name]?.state || {};

  if (hint) {
    const hintMessage = hint({ previousState, currentState: state });
    if (hintMessage) {
      await sendMessage(hintMessage);
    } else {
      await runChatCompletion();
    }
  } else {
    await runChatCompletion();
  }
}

const isExternalStateManagement = <T>(
  options: UseCoagentOptions<T>,
): options is WithExternalStateManagement<T> => {
  return "state" in options && "setState" in options;
};

const isInternalStateManagementWithInitial = <T>(
  options: UseCoagentOptions<T>,
): options is WithInternalStateManagementAndInitial<T> => {
  return "initialState" in options;
};

const getAiAgentState = <T>({
  aiAgentStates,
  name,
  options,
}: {
  aiAgentStates: Record<string, AiAgentState>;
  name: string;
  options: UseCoagentOptions<T>;
}) => {
  if (aiAgentStates[name]) {
    return aiAgentStates[name];
  } else {
    return {
      name,
      state: isInternalStateManagementWithInitial<T>(options) ? options.initialState : {},
      config: options.config
        ? options.config
        : options.configurable
          ? { configurable: options.configurable }
          : {},
      running: false,
      active: false,
      threadId: undefined,
      nodeName: undefined,
      runId: undefined,
    };
  }
};