import {
  AiCloudConfig,
  FunctionCallHandler,
  AiErrorHandler,
  AiSDKError,
} from "@vn-sdk/shared";
import {
  ActionRenderProps,
  CatchAllActionRenderProps,
  FrontendAction,
} from "../types/frontend-action";
import React from "react";
import { TreeNodeId, Tree } from "../hooks/use-tree";
import { DocumentPointer } from "../types";
import { AiChatSuggestionConfiguration } from "../types/chat-suggestion-configuration";
import { AiAgentStateRender, AiAgentStateRenderProps } from "../types/ai-agent-action";
import { AiAgentState } from "../types/ai-agent-state";
import {
  AiRuntimeClient,
  ExtensionsInput,
  ForwardedParametersInput,
} from "@vn-sdk/runtime-client-gql";
import { Agent } from "@vn-sdk/runtime-client-gql";
import {
  LangGraphInterruptAction,
  LangGraphInterruptActionSetter,
} from "../types/interrupt-action";
import { SuggestionItem } from "../utils/suggestions";

/**
 * Interface for the configuration of the AI API.
 */
export interface AiApiConfig {
  /**
   * The public API key for VN Cloud.
   */
  publicApiKey?: string;

  /**
   * The configuration for VN Cloud.
   */
  cloud?: AiCloudConfig;

  /**
   * The endpoint for the chat API.
   */
  chatApiEndpoint: string;

  /**
   * The endpoint for the AI transcribe audio service.
   */
  transcribeAudioUrl?: string;

  /**
   * The endpoint for the AI text to speech service.
   */
  textToSpeechUrl?: string;

  /**
   * additional headers to be sent with the request
   * @default {}
   * @example
   * ```
   * {
   *   'Authorization': 'Bearer your_token_here'
   * }
   * ```
   */
  headers: Record<string, string>;

  /**
   * Custom properties to be sent with the request
   * @default {}
   * @example
   * ```
   * {
   *   'user_id': 'user_id'
   * }
   * ```
   */
  properties?: Record<string, any>;

  /**
   * Indicates whether the user agent should send or receive cookies from the other domain
   * in the case of cross-origin requests.
   */
  credentials?: RequestCredentials;

  /**
   * Optional configuration for connecting to Model Context Protocol (MCP) servers.
   * This is typically derived from the AiProps and used internally.
   * @experimental
   */
  mcpServers?: Array<{ endpoint: string; apiKey?: string }>;
}

export type InChatRenderFunction<TProps = ActionRenderProps<any> | CatchAllActionRenderProps<any>> =
  (props: TProps) => string | JSX.Element;
export type AiAgentInChatRenderFunction = (
  props: AiAgentStateRenderProps<any>,
) => string | JSX.Element | undefined | null;

export interface ChatComponentsCache {
  actions: Record<string, InChatRenderFunction | string>;
  aiAgentStateRenders: Record<string, AiAgentInChatRenderFunction | string>;
}

export interface AgentSession {
  agentName: string;
  threadId?: string;
  nodeName?: string;
}

export interface AuthState {
  status: "authenticated" | "unauthenticated";
  authHeaders: Record<string, string>;
  userId?: string;
  metadata?: Record<string, any>;
}

export type ActionName = string;
export type ContextTree = Tree;

export interface AiContextParams {
  // function-calling
  actions: Record<string, FrontendAction<any>>;
  setAction: (id: string, action: FrontendAction<any>) => void;
  removeAction: (id: string) => void;

  // coagent actions
  aiAgentStateRenders: Record<string, AiAgentStateRender<any>>;
  setAiAgentStateRender: (id: string, stateRender: AiAgentStateRender<any>) => void;
  removeAiAgentStateRender: (id: string) => void;

  chatComponentsCache: React.RefObject<ChatComponentsCache>;

  getFunctionCallHandler: (
    customEntryPoints?: Record<string, FrontendAction<any>>,
  ) => FunctionCallHandler;

  // text context
  addContext: (context: string, parentId?: string, categories?: string[]) => TreeNodeId;
  removeContext: (id: TreeNodeId) => void;
  getAllContext: () => Tree;
  getContextString: (documents: DocumentPointer[], categories: string[]) => string;

  // document context
  addDocumentContext: (documentPointer: DocumentPointer, categories?: string[]) => TreeNodeId;
  removeDocumentContext: (documentId: string) => void;
  getDocumentsContext: (categories: string[]) => DocumentPointer[];

  isLoading: boolean;
  setIsLoading: React.Dispatch<React.SetStateAction<boolean>>;

  chatSuggestionConfiguration: { [key: string]: AiChatSuggestionConfiguration };
  addChatSuggestionConfiguration: (
    id: string,
    suggestion: AiChatSuggestionConfiguration,
  ) => void;
  removeChatSuggestionConfiguration: (id: string) => void;

  chatInstructions: string;
  setChatInstructions: React.Dispatch<React.SetStateAction<string>>;

  additionalInstructions?: string[];
  setAdditionalInstructions: React.Dispatch<React.SetStateAction<string[]>>;

  // api endpoints
  aiApiConfig: AiApiConfig;

  showDevConsole: boolean;

  // agents
  aiAgentStates: Record<string, AiAgentState>;
  setAiAgentStates: React.Dispatch<React.SetStateAction<Record<string, AiAgentState>>>;
  aiAgentStatesRef: React.RefObject<Record<string, AiAgentState>>;
  setAiAgentStatesWithRef: (
    value:
      | Record<string, AiAgentState>
      | ((prev: Record<string, AiAgentState>) => Record<string, AiAgentState>),
  ) => void;

  agentSession: AgentSession | null;
  setAgentSession: React.Dispatch<React.SetStateAction<AgentSession | null>>;

  agentLock: string | null;

  threadId: string;
  setThreadId: React.Dispatch<React.SetStateAction<string>>;

  runId: string | null;
  setRunId: React.Dispatch<React.SetStateAction<string | null>>;

  // The chat abort controller can be used to stop generation globally,
  // i.e. when using `stop()` from `useChat`
  chatAbortControllerRef: React.MutableRefObject<AbortController | null>;

  // runtime
  runtimeClient: AiRuntimeClient;

  /**
   * The forwarded parameters to use for the task.
   */
  forwardedParameters?: Partial<Pick<ForwardedParametersInput, "temperature">>;
  availableAgents: Agent[];

  /**
   * The auth states for the VN SDK.
   */
  authStates_c?: Record<ActionName, AuthState>;
  setAuthStates_c?: React.Dispatch<React.SetStateAction<Record<ActionName, AuthState>>>;

  /**
   * The auth config for the VN SDK.
   */
  authConfig_c?: {
    SignInComponent: React.ComponentType<{
      onSignInComplete: (authState: AuthState) => void;
    }>;
  };

  extensions: ExtensionsInput;
  setExtensions: React.Dispatch<React.SetStateAction<ExtensionsInput>>;
  langGraphInterruptAction: LangGraphInterruptAction | null;
  setLangGraphInterruptAction: LangGraphInterruptActionSetter;
  removeLangGraphInterruptAction: (threadId: string) => void;

  /**
   * Optional trace handler for comprehensive debugging and observability.
   */
  onError: AiErrorHandler;

  // banner error state
  bannerError: AiSDKError | null;
  setBannerError: React.Dispatch<React.SetStateAction<AiSDKError | null>>;
  // Internal error handlers
  // These are used to handle errors that occur during the execution of the chat.
  // They are not intended for use by the developer. A component can register itself an error listener to be activated somewhere else as needed
  internalErrorHandlers: Record<string, AiErrorHandler>;
  setInternalErrorHandler: (handler: Record<string, AiErrorHandler>) => void;
  removeInternalErrorHandler: (id: string) => void;
}

const emptyAiContext: AiContextParams = {
  actions: {},
  setAction: () => {},
  removeAction: () => {},

  aiAgentStateRenders: {},
  setAiAgentStateRender: () => {},
  removeAiAgentStateRender: () => {},

  chatComponentsCache: { current: { actions: {}, aiAgentStateRenders: {} } },
  getContextString: (documents: DocumentPointer[], categories: string[]) =>
    returnAndThrowInDebug(""),
  addContext: () => "",
  removeContext: () => {},
  getAllContext: () => [],

  getFunctionCallHandler: () => returnAndThrowInDebug(async () => {}),

  isLoading: false,
  setIsLoading: () => returnAndThrowInDebug(false),

  chatInstructions: "",
  setChatInstructions: () => returnAndThrowInDebug(""),

  additionalInstructions: [],
  setAdditionalInstructions: () => returnAndThrowInDebug([]),

  getDocumentsContext: (categories: string[]) => returnAndThrowInDebug([]),
  addDocumentContext: () => returnAndThrowInDebug(""),
  removeDocumentContext: () => {},
  runtimeClient: {} as any,

  aiApiConfig: new (class implements AiApiConfig {
    get chatApiEndpoint(): string {
      throw new Error("Remember to wrap your app in a `<AiProvider> {...} </AiProvider>` !!!");
    }

    get headers(): Record<string, string> {
      return {};
    }
    get body(): Record<string, any> {
      return {};
    }
  })(),

  chatSuggestionConfiguration: {},
  addChatSuggestionConfiguration: () => {},
  removeChatSuggestionConfiguration: () => {},
  showDevConsole: false,
  aiAgentStates: {},
  setAiAgentStates: () => {},
  aiAgentStatesRef: { current: {} },
  setAiAgentStatesWithRef: () => {},
  agentSession: null,
  setAgentSession: () => {},
  forwardedParameters: {},
  agentLock: null,
  threadId: "",
  setThreadId: () => {},
  runId: null,
  setRunId: () => {},
  chatAbortControllerRef: { current: null },
  availableAgents: [],
  extensions: {},
  setExtensions: () => {},
  langGraphInterruptAction: null,
  setLangGraphInterruptAction: () => {},
  removeLangGraphInterruptAction: () => {},
  onError: () => {},
  bannerError: null,
  setBannerError: () => {},
  internalErrorHandlers: {},
  setInternalErrorHandler: () => {},
  removeInternalErrorHandler: () => {},
};

export const AiContext = React.createContext<AiContextParams>(emptyAiContext);

export function useAiContext(): AiContextParams {
  const context = React.useContext(AiContext);
  if (context === emptyAiContext) {
    throw new Error("Remember to wrap your app in a `<AiProvider> {...} </AiProvider>` !!!");
  }
  return context;
}

function returnAndThrowInDebug<T>(_value: T): T {
  throw new Error("Remember to wrap your app in a `<AiProvider> {...} </AiProvider>` !!!");
}