import { Dispatch, MutableRefObject, SetStateAction, useCallback, useEffect, useId, useRef, useState } from 'react';
import useSWR, { KeyedMutator } from 'swr';
import { callChatApi } from '../shared/call-chat-api';
import { generateId as generateIdFunc } from '../shared/generate-id';
import { processChatStream } from '../shared/process-chat-stream';
import type {
  ChatRequest,
  ChatRequestOptions,
  CreateMessage,
  IdGenerator,
  JSONValue,
  Message,
  UseChatOptions,
  WebSocketMessage,
} from '../shared/types';
import type {
  ReactResponseRow,
  experimental_StreamingReactResponse,
} from '../streams/streaming-react-response';
import { readDataStream } from '../shared/read-data-stream';
export type { CreateMessage, Message, UseChatOptions };

export const ENDTXT = '|=EOF=|'

export type UseISDKHelpers = {
  /** Current messages in the chat */
  messages: Message[];
  /** The error object of the API request */
  error: undefined | Error;
  /**
   * Append a user message to the chat list. This triggers the API call to fetch
   * the assistant's response.
   * @param message The message to append
   * @param options Additional options to pass to the API call
   */
  append: (
    message: Message | CreateMessage,
    chatRequestOptions?: ChatRequestOptions,
  ) => Promise<string | null | undefined>;
  /**
   * Reload the last AI chat response for the given chat history. If the last
   * message isn't from the assistant, it will request the API to generate a
   * new response.
   */
  reload: (
    chatRequestOptions?: ChatRequestOptions,
  ) => Promise<string | null | undefined>;
  /**
   * Abort the current request immediately, keep the generated tokens if any.
   */
  stop: () => void;
  /**
   * Update the `messages` state locally. This is useful when you want to
   * edit the messages on the client, and then trigger the `reload` method
   * manually to regenerate the AI response.
   */
  setMessages: (messages: Message[]) => void;
  /** The current value of the input */
  input: string;
  /** setState-powered method to update the input value */
  setInput: React.Dispatch<React.SetStateAction<string>>;
  /** An input/textarea-ready onChange handler to control the value of the input */
  handleInputChange: (
    e:
      | React.ChangeEvent<HTMLInputElement>
      | React.ChangeEvent<HTMLTextAreaElement>,
  ) => void;
  /** Form submission handler to automatically reset input and append a user message */
  handleSubmit: (
    e: React.FormEvent<HTMLFormElement>,
    chatRequestOptions?: ChatRequestOptions,
  ) => void;
  metadata?: Object;
  /** Whether the API request is in progress */
  isLoading: boolean;
  isSocket: boolean;


  /** Additional data added on the server via StreamData */
  data?: JSONValue[];
};

export type IMessages = {
  action: string;
  type: string; //text: 文本聊天会话 audio: 语音聊天会话
  content: string;
}


type StreamingReactResponseAction = (payload: {
  messages: Message[];
  data?: Record<string, string>;
}) => Promise<experimental_StreamingReactResponse>;




const handleWebSocketMessage = async (

  event: MessageEvent<any>,
  messagesRef: React.MutableRefObject<Message[]>,
  isLoadingRef: MutableRefObject<boolean>,
  mutate: KeyedMutator<Message[]>,
  setInput: Dispatch<SetStateAction<string>>,
  mutateLoading: KeyedMutator<boolean>,
  generateId: IdGenerator,
  streamMode?: 'stream-data' | 'text',
  onFinish?: (message: Message) => void,
  onResponse?: (response: Response) => void | Promise<void>,
  onFunctionCall?: (
    functionCall: WebSocketMessage,
    // functionCall: FunctionCall,
    // onToolCall: (message: Message) => void
  ) => void | Promise<void>,


  // mutateStreamData: KeyedMutator<JSONValue[] | undefined>,
  // chatRequest: ChatRequest,
  // existingData: JSONValue[] | undefined,
  // extraMetadataRef: React.MutableRefObject<any>,

  // abortControllerRef: React.MutableRefObject<AbortController | null>,


) => {
  const receivedMessage = event.data;
  let parsedMessage = receivedMessage;


  // const reader = event.data.body.getReader(); // 获取数据流读取器

  mutate(messagesRef.current, false);
  const previousMessages = messagesRef.current;

  const length = previousMessages.length - 1
  // mutate(chatRequest.messages, false);

  try {
    parsedMessage = JSON.parse(receivedMessage);

    if (typeof parsedMessage == "object" && parsedMessage.type === 'function') {

      console.log("------请调用:", parsedMessage)


      if (onFunctionCall) {
        // // const functionCall = message.function_call;
        // // // Make sure functionCall is an object
        // // // If not, we got tool calls instead of function calls
        // // if (typeof functionCall !== 'object') {
        // //   console.warn(
        // //     'experimental_onFunctionCall should not be defined when using tools',
        // //   );
        // //   continue;
        // // }

        // // User handles the function call in their own functionCallHandler.
        // // The "arguments" key of the function call object will still be a string which will have to be parsed in the function handler.
        // // If the "arguments" JSON is malformed due to model error the user will have to handle that themselves.

        // const functionCallResponse: ChatRequest | void =
        //   await experimental_onFunctionCall(
        //     getCurrentMessages(),
        //     functionCall,
        //   );

        // // If the user does not return anything as a result of the function call, the loop will break.
        // if (functionCallResponse === undefined) {
        //   hasFollowingResponse = false;
        //   break;
        // }

        // // A function call response was returned.
        // // The updated chat with function call response will be sent to the API in the next iteration of the loop.
        // updateChatRequest(functionCallResponse);

        mutate([...previousMessages, {
          id: generateId(),
          content: "",
          role: 'assistant',
          createdAt: new Date(),
          function_call: parsedMessage.functionName
        }], false);

        await onFunctionCall(parsedMessage);


        if (onFinish) {
          onFinish(parsedMessage);
        }

      }

      return
    }
  } catch (error) {
    // console.error("解析接收到的消息时出错:", error);
  }

  let messageBuffer: Message = {
    id: generateId(),
    content: "",
    role: 'assistant',
    createdAt: new Date()
  };

  console.log("===================", isLoadingRef.current)

  if (streamMode === "text") {
    // V1 版本
    switch (parsedMessage.streamStatus) {
      case "start":
        mutateLoading(true);
        console.log("开始接收信息", parsedMessage);

        mutate([...previousMessages, messageBuffer], false);

        break;
      case "end":
        console.log("结束接收信息", parsedMessage);

        mutate([...previousMessages.slice(0, -1), {
          ...previousMessages[length],
          createdAt: new Date()
        }], false);
        setInput("")
        mutateLoading(false);
        break;
      default:


        // let lastMessage = previousMessages[length];
        // console.log("中间接收信息进行消息输出拼接", parsedMessage, previousMessages[length]);

        previousMessages[length].content = previousMessages[length].content + parsedMessage
        previousMessages[length].createdAt = new Date()
        // setMessages([...messagesRef.current.slice(0, -1), lastMessage])

        mutate([...previousMessages.slice(0, -1), ...[previousMessages[length]]], false);


        break;
    }
  } else {
    switch (parsedMessage) {

      case ENDTXT:
        console.log("结束接收信息", parsedMessage);

        mutate([...previousMessages.slice(0, -1), {
          ...previousMessages[length],
          createdAt: new Date()
        }], false);
        setInput("")
        isLoadingRef.current = false
        mutateLoading(false);
        break;
      default:
        if (!isLoadingRef.current) {
          isLoadingRef.current = true
          mutateLoading(true);
          console.log("开始接收信息", parsedMessage);

          messageBuffer.content = parsedMessage

          mutate([...previousMessages, messageBuffer], false);

        } else {
          // let lastMessage = previousMessages[length];
          // console.log("中间接收信息进行消息输出拼接", parsedMessage, previousMessages[length]);

          previousMessages[length].content = previousMessages[length].content + parsedMessage
          previousMessages[length].createdAt = new Date()
          // setMessages([...messagesRef.current.slice(0, -1), lastMessage])

          mutate([...previousMessages.slice(0, -1), ...[previousMessages[length]]], false);
        }

        break;
    }
  }






};


const getStreamedResponse = async (
  api: string | StreamingReactResponseAction,
  chatRequest: ChatRequest,
  mutate: KeyedMutator<Message[]>,
  mutateStreamData: KeyedMutator<JSONValue[] | undefined>,
  existingData: JSONValue[] | undefined,
  extraMetadataRef: React.MutableRefObject<any>,
  messagesRef: React.MutableRefObject<Message[]>,
  abortControllerRef: React.MutableRefObject<AbortController | null>,
  generateId: IdGenerator,
  onFinish?: (message: Message) => void,
  onResponse?: (response: Response) => void | Promise<void>,
  sendExtraMessageFields?: boolean,
) => {
  // Do an optimistic update to the chat state to show the updated messages
  // immediately.
  const previousMessages = messagesRef.current;
  mutate(chatRequest.messages, false);

  const constructedMessagesPayload = sendExtraMessageFields
    ? chatRequest.messages
    : chatRequest.messages.map(
      ({ role, content, name, function_call, tool_calls, tool_call_id }) => ({
        role,
        content,
        tool_call_id,
        ...(name !== undefined && { name }),
        ...(function_call !== undefined && {
          function_call: function_call,
        }),
        ...(tool_calls !== undefined && {
          tool_calls: tool_calls,
        }),
      }),
    );

  if (typeof api !== 'string') {
    // In this case, we are handling a Server Action. No complex mode handling needed.

    const replyId = generateId();
    const createdAt = new Date();
    let responseMessage: Message = {
      id: replyId,
      createdAt,
      content: '',
      role: 'assistant',
    };

    async function readRow(promise: Promise<ReactResponseRow>) {


      console.log("-----promise--------", promise)
      const { content, ui, next } = await promise;

      // TODO: Handle function calls.
      responseMessage['content'] = content;
      responseMessage['ui'] = await ui;

      mutate([...chatRequest.messages, { ...responseMessage }], false);

      if (next) {
        await readRow(next);
      }
    }

    try {
      const promise = api({
        messages: constructedMessagesPayload as Message[],
        data: chatRequest.data,
      }) as Promise<ReactResponseRow>;
      await readRow(promise);
    } catch (e) {
      // Restore the previous messages if the request fails.
      mutate(previousMessages, false);
      throw e;
    }

    if (onFinish) {
      onFinish(responseMessage);
    }

    return responseMessage;
  }

  return await callChatApi({
    api,
    messages: constructedMessagesPayload,
    body: {
      data: chatRequest.data,
      ...extraMetadataRef.current.body,
      ...chatRequest.options?.body,
      ...(chatRequest.functions !== undefined && {
        functions: chatRequest.functions,
      }),
      ...(chatRequest.function_call !== undefined && {
        function_call: chatRequest.function_call,
      }),
      ...(chatRequest.tools !== undefined && {
        tools: chatRequest.tools,
      }),
      ...(chatRequest.tool_choice !== undefined && {
        tool_choice: chatRequest.tool_choice,
      }),
    },
    credentials: extraMetadataRef.current.credentials,
    headers: {
      ...extraMetadataRef.current.headers,
      ...chatRequest.options?.headers,
    },
    abortController: () => abortControllerRef.current,
    restoreMessagesOnFailure() {
      mutate(previousMessages, false);
    },
    onResponse,
    onUpdate(merged, data) {
      console.log("------onUpdate----------", chatRequest.messages, merged, data)
      mutate([...chatRequest.messages, ...merged], false);
      mutateStreamData([...(existingData || []), ...(data || [])], false);
    },
    onFinish,
    generateId,
  });
};

export function useISDK({
  api = '/api/chat',
  id,
  initialMessages,
  initialInput = '',
  sendExtraMessageFields,
  experimental_onFunctionCall,
  experimental_onFunctionCallV2,
  experimental_onToolCall,
  onResponse,
  onFinish,
  onError,
  credentials,
  headers,
  body,
  streamMode,
  // username,
  // password,

  generateId = generateIdFunc,
}: Omit<UseChatOptions, 'api'> & {
  api?: string | StreamingReactResponseAction;
  key?: string;
} = {}): UseISDKHelpers {
  // Generate a unique id for the chat if not provided.
  const hookId = useId();
  const idKey = id ?? hookId;
  const chatKey = typeof api === 'string' ? [api, idKey] : idKey;


  const socketRef = useRef<WebSocket | null>()
  const [isSocket, setIsSocket] = useState<boolean>(false);

  // Store a empty array as the initial messages
  // (instead of using a default parameter value that gets re-created each time)
  // to avoid re-renders:
  const [initialMessagesFallback] = useState([]);

  // Store the chat state in SWR, using the chatId as the key to share states.
  const { data: messages, mutate } = useSWR<Message[]>(
    [chatKey, 'messages'],
    null,
    { fallbackData: initialMessages ?? initialMessagesFallback },
  );

  console.log("-----useSWR--initialMessages---------", initialMessages)

  // We store loading state in another hook to sync loading states across hook invocations
  const { data: isLoading = false, mutate: mutateLoading } = useSWR<boolean>(
    [chatKey, 'loading'],
    null,
  );

  const { data: streamData, mutate: mutateStreamData } = useSWR<
    JSONValue[] | undefined
  >([chatKey, 'streamData'], null);

  const { data: error = undefined, mutate: setError } = useSWR<
    undefined | Error
  >([chatKey, 'error'], null);

  const isLoadingRef = useRef<boolean>(false);
  useEffect(() => {
    isLoadingRef.current = isLoading;
  }, [isLoading]);

  // Keep the latest messages in a ref.
  const messagesRef = useRef<Message[]>(messages || []);
  useEffect(() => {
    messagesRef.current = messages || [];
  }, [messages]);

  // Abort controller to cancel the current API call.
  const abortControllerRef = useRef<AbortController | null>(null);

  const extraMetadataRef = useRef({
    credentials,
    headers,
    body,
  });

  useEffect(() => {
    extraMetadataRef.current = {
      credentials,
      headers,
      body,
    };
  }, [credentials, headers, body]);


  const setupWebSocket = ((newURL: string) => {
    if (socketRef.current) {
      socketRef.current.close(); // 关闭旧连接
    }

    // if (!!username && !!password) {
    //   // 构建 Basic 认证字符串
    //   // const credentials = Buffer.from(`${username}:${password}`).toString('base64');
    //   // const authHeader = `Basic ${credentials}`;

    //   const authHeader = `${Buffer.from(`${username}:${password}`).toString('base64')}`;

    //   const credentialsV2 = btoa(`${username}:${password}`);
    //   // socket.send('Authorization: Basic ' + credentialsV2);

    //   // socket.addEventListener('open', (event) => {
    //   //   // 在握手阶段添加Authorization头
    //   //   socket.send('Authorization: Bearer ' + authHeader);
    //   // });
    // }

    // const webSocket = new WebSocket(newURL)
    const socket = new WebSocket(newURL, []);

    // const socket = new WebSocket("ws://JS8dw07Ld0hcCsVf:8jfPM53Sq84Etvc3J98kqlxtFwqPOvVMQxlGmc3ud7sy10ch@116.213.39.234:8084/ws")

    // if (socketRef.current) return;
    // setWs(socket);

    socket.onopen = () => {
      console.log('WebSocket 连接已建立',);

      setIsSocket(true)
      if (initialInput !== undefined) {
        // console.log("---useWebSocket------", options)
        // ws.send(options.initialInput as string); // 将 undefined 转换为 string
        // setInput("")
        // append({
        //   id: '',
        //   content: initialInput,
        //   role: 'function'
        // })
      }
    };

    socket.onclose = async () => {
      console.log('WebSocket 连接已关闭', socketRef?.current);
      // if (onFinish && !!messages) {
      //   onFinish(messages);
      // }

      await new Promise(r => setTimeout(r, 3000));

      if (socketRef?.current) {
        socketRef.current.close()
        setupWebSocket(newURL)
      } else {
        setupWebSocket(newURL)
      }
    };


    socket.onerror = (error) => {
      console.error('WebSocket 连接发生错误:', error);
      if (onError) {
        // onError(error);
      }
      mutateLoading(false);// 发生错误后，设置 isLoading 为 false
    };

    let loading = false
    socket.onmessage = (event) => {
      try {
        mutateLoading(true);
        setError(undefined);

        const abortController = new AbortController();
        abortControllerRef.current = abortController;

        handleWebSocketMessage(

          event,
          messagesRef,
          isLoadingRef,
          mutate,
          setInput,
          mutateLoading,
          generateId,
          streamMode,
          onFinish,
          onResponse,
          // experimental_onFunctionCallV2: 

          async (data: WebSocketMessage) => {
            if (experimental_onFunctionCallV2) {
              await experimental_onFunctionCallV2(data)
            }

            // toolCallsResult
            // 回调成功
            if (socketRef.current) {
              // socketRef.current.send("INTENTION_CLIENT_TOOL_CALL_FUNCTION_FINISHED")
            }

          }
        );

        abortControllerRef.current = null;
      } catch (err) {
        // Ignore abort errors as they are expected.
        if ((err as any).name === 'AbortError') {
          abortControllerRef.current = null;
          return null;
        }

        if (onError && err instanceof Error) {
          onError(err);
        }

        setError(err as Error);
      } finally {
        mutateLoading(false);
      }
    }

    socketRef.current = socket

  })

  useEffect(() => {

    const socket = socketRef.current

    if (!socket) {
      // let url = "ws://116.213.39.234:8083/ws"
      // let url = "wss://www.jellyai.xyz/ws"
      let url = api as string
      setupWebSocket(url)
    }
  }, [socketRef, onResponse, onFinish, onError])

  const triggerRequest = useCallback(
    async (chatRequest: ChatRequest) => {
      try {
        mutateLoading(true);
        setError(undefined);

        const abortController = new AbortController();
        abortControllerRef.current = abortController;


        await processChatStream({
          getStreamedResponse: () =>
            getStreamedResponse(
              api,
              chatRequest,
              mutate,
              mutateStreamData,
              streamData!,
              extraMetadataRef,
              messagesRef,
              abortControllerRef,
              generateId,
              onFinish,
              onResponse,
              sendExtraMessageFields,
            ),
          experimental_onFunctionCall,
          experimental_onToolCall,
          updateChatRequest: chatRequestParam => {
            chatRequest = chatRequestParam;
          },
          getCurrentMessages: () => messagesRef.current,
        });

        abortControllerRef.current = null;
      } catch (err) {
        // Ignore abort errors as they are expected.
        if ((err as any).name === 'AbortError') {
          abortControllerRef.current = null;
          return null;
        }

        if (onError && err instanceof Error) {
          onError(err);
        }

        setError(err as Error);
      } finally {
        mutateLoading(false);
      }
    },
    [
      mutate,
      mutateLoading,
      api,
      extraMetadataRef,
      onResponse,
      onFinish,
      onError,
      setError,
      mutateStreamData,
      streamData,
      sendExtraMessageFields,
      experimental_onFunctionCall,
      experimental_onToolCall,
      messagesRef,
      abortControllerRef,
      generateId,
    ],
  );

  const append = useCallback(
    async (
      message: Message | CreateMessage,
      {
        options,
        functions,
        function_call,
        tools,
        tool_choice,
        data,
      }: ChatRequestOptions = {},
    ) => {
      if (!message.id) {
        message.id = generateId();
      }

      const chatRequest: ChatRequest = {
        messages: messagesRef.current.concat(message as Message),
        options,
        data,
        ...(functions !== undefined && { functions }),
        ...(function_call !== undefined && { function_call }),
        ...(tools !== undefined && { tools }),
        ...(tool_choice !== undefined && { tool_choice }),
      };

      if (isSocket) {
        try {
          if (socketRef?.current && socketRef?.current.readyState === WebSocket.OPEN) {

            // setMessages([...messagesRef.current, {
            //   id: message.id,
            //   content: message.content,
            //   role: message.role
            // }])
            mutateLoading(true);

            mutate([...messagesRef.current, {
              id: message.id,
              content: message.content,
              role: message.role
            }], false)

            // 发送消息时，设置 isLoading 为 true
            // socketRef.current.send(message.content);
            socketRef.current.send(JSON.stringify(message));
          } else {
            console.error('WebSocket 连接未建立或已关闭');
          }
        } catch (error) {

        } finally {
          mutateLoading(false);
        }
      } else {
        return triggerRequest(chatRequest);
      }

      // return triggerRequest(chatRequest);

    },
    [triggerRequest, generateId],
  );

  const reload = useCallback(
    async ({
      options,
      functions,
      function_call,
      tools,
      tool_choice,
    }: ChatRequestOptions = {}) => {
      if (messagesRef.current.length === 0) return null;

      // if (isSocket) {
      //   try {
      //     if (socketRef?.current && socketRef?.current.readyState === WebSocket.OPEN) {
      //       mutateLoading(true); // 发送消息时，设置 isLoading 为 true
      //       socketRef.current.send(message.content);
      //     } else {
      //       console.error('WebSocket 连接未建立或已关闭');
      //     }
      //   } catch (error) {

      //   } finally {
      //     mutateLoading(false);
      //   }

      //   return
      // } 


      // Remove last assistant message and retry last user message.
      const lastMessage = messagesRef.current[messagesRef.current.length - 1];
      if (lastMessage.role === 'assistant') {
        const chatRequest: ChatRequest = {
          messages: messagesRef.current.slice(0, -1),
          options,
          ...(functions !== undefined && { functions }),
          ...(function_call !== undefined && { function_call }),
          ...(tools !== undefined && { tools }),
          ...(tool_choice !== undefined && { tool_choice }),
        };

        return triggerRequest(chatRequest);
      }

      const chatRequest: ChatRequest = {
        messages: messagesRef.current,
        options,
        ...(functions !== undefined && { functions }),
        ...(function_call !== undefined && { function_call }),
        ...(tools !== undefined && { tools }),
        ...(tool_choice !== undefined && { tool_choice }),
      };

      return triggerRequest(chatRequest);
    },
    [triggerRequest],
  );

  const stop = useCallback(() => {
    if (abortControllerRef.current) {
      abortControllerRef.current.abort();
      abortControllerRef.current = null;
    }
  }, []);

  const setMessages = useCallback(
    (messages: Message[]) => {
      mutate(messages, false);
      messagesRef.current = messages;
    },
    [mutate],
  );

  // Input state and handlers.
  const [input, setInput] = useState(initialInput);

  const handleSubmit = useCallback(
    (
      e: React.FormEvent<HTMLFormElement>,
      options: ChatRequestOptions = {},
      metadata?: Object,
    ) => {
      if (metadata) {
        extraMetadataRef.current = {
          ...extraMetadataRef.current,
          ...metadata,
        };
      }

      e.preventDefault();
      if (!input) return;

      append(
        {
          content: input,
          role: 'user',
          createdAt: new Date(),
          action: "chat",
          type: 'text',
        },
        options,
      );
      setInput('');
    },
    [input, append],
  );

  const handleInputChange = (e: any) => {
    setInput(e.target.value);
  };

  return {
    messages: messages || [],
    error,
    append,
    reload,
    stop,
    setMessages,
    input,
    setInput,
    handleInputChange,
    handleSubmit,
    isLoading,
    isSocket,
    data: streamData,
  };
}