import { useState, useEffect } from "react";
import { FeedbackButton } from "@banorte/chat-ui";
import Markdown from "react-markdown";
import rehypeRaw from "rehype-raw";
import { SSE } from "sse.js";

export { ChatMessage };

interface ChatMessageProps {
  isUser: boolean;
  content: string;
  event: CallableFunction;
  conversationId: string;
  withDeepResearch: boolean;
  setReceivingMsg: (receiving: boolean) => void;
  userAvatar: string;
  botAvatar: string;
  onFeedback?: (key: string, rating: string) => Promise<void>;
}

function ChatMessage({
  isUser,
  content,
  event,
  conversationId,
  withDeepResearch,
  setReceivingMsg,
  userAvatar,
  botAvatar,
  onFeedback,
}: ChatMessageProps) {
  const [buff, setBuff] = useState("");
  const [responseId, setResponseId] = useState("");
  const [loading, setLoading] = useState(false);

  const [images, setImages] = useState<string[]>([]);
  const [currentImageIndex, setCurrentImageIndex] = useState(0);
  const [acceptFeedback, setAcceptFeedback] = useState(false);
  const [streamIndex, setStreamIndex] = useState(0);
  const [fullResponse, setFullResponse] = useState("");
  const [pendingReferences, setPendingReferences] = useState<Array<string>>([]);
  const [streamingComplete, setStreamingComplete] = useState(false);

  const nextImage = () => {
    if (currentImageIndex < images.length - 1) {
      setCurrentImageIndex((prev) => prev + 1);
    }
  };

  const prevImage = () => {
    if (currentImageIndex > 0) {
      setCurrentImageIndex((prev) => prev - 1);
    }
  };

  function setReferences(buff: string, references: Array<string>) {
    const citations = buff.match(/\[(\d+)\]/g);
    let newText = buff;
    if (citations) {
      citations.forEach((citation) => {
        const citationNumber = parseInt(citation.replace(/\[|\]/g, "")) - 1;
        const reference = references[citationNumber];
        const anchorTag = `<a class="text-blue-700 underline" href="${reference}"  target="_blank" rel="noopener noreferrer">${citation}</a>`;
        newText = newText.replaceAll(citation, anchorTag);
      });
    }
    return newText;
  }

  useEffect(() => {
    if (fullResponse && streamIndex < fullResponse.length) {
      setLoading(false);

      const timer = setTimeout(() => {
        setBuff((prev) => prev + fullResponse[streamIndex]);
        setStreamIndex((prev) => prev + 1);
        event();
      }, 3);

      return () => clearTimeout(timer);
    } else if (fullResponse && streamIndex === fullResponse.length) {
      setReceivingMsg(false);
      setStreamingComplete(true);
      // Apply references after streaming is complete
      if (pendingReferences.length > 0) {
        const referencedText = setReferences(fullResponse, pendingReferences);
        setBuff(referencedText);
        setPendingReferences([]);
      }
    }
  }, [fullResponse, streamIndex, pendingReferences]);

  async function getStream() {
    const payload = JSON.stringify({
      prompt: content,
      conversation_id: conversationId,
      with_deep_research: withDeepResearch,
    });

    const url = "/api/v1/message?stream=True";
    const eventSource = new SSE(url, {
      withCredentials: true,
      headers: { "Content-Type": "application/json" },
      payload: payload,
    });

    eventSource.onmessage = async (event) => {
      console.log(event.data);
      const ResponseChunk = JSON.parse(event.data);

      if (ResponseChunk["type"] === "text") {
        setFullResponse((prev) => prev + ResponseChunk["content"]);
      } else if (ResponseChunk["type"] === "reference") {
        setPendingReferences(ResponseChunk["content"]);
      } else if (ResponseChunk["type"] === "end") {
        setResponseId(ResponseChunk["content"]);
        eventSource.close();
      } else if (ResponseChunk["type"] === "image") {
        const newImages = ResponseChunk.content.slice(0, 3);
        setImages((prev) => {
          const combinedImages = [...prev, ...newImages];
          return combinedImages.slice(0, 3);
        });
      } else if (ResponseChunk["type"] == "tool") {
        setAcceptFeedback(true);
      } else if (ResponseChunk["type"] === "error") {
        setFullResponse((prev) => prev + "\n\n" + ResponseChunk["content"]);
        eventSource.close();
      }
    };
    eventSource.onerror = async (e) => {
      console.log("error" + e);
      setReceivingMsg(false);
      setLoading(false);
      eventSource.close();
    };
  }

  useEffect(() => {
    if (!isUser) {
      setLoading(true);
      setReceivingMsg(true);
      getStream();
    } else {
      setBuff(content);
      event();
    }
  }, []);

  const ImageViewer = () => {
    if (images.length === 0) return null;

    return (
      <div className="mt-5 space-y-4">
        <div className="relative">
          <img
            src={images[currentImageIndex]}
            alt={`Generated image ${currentImageIndex + 1}`}
            className="w-full h-auto rounded-lg"
          />
          <div className="flex justify-between items-center mt-4">
            <button
              onClick={prevImage}
              disabled={currentImageIndex === 0}
              className={`px-4 py-2 rounded ${
                currentImageIndex === 0
                  ? "text-gray-400 cursor-not-allowed"
                  : "text-gray-700 hover:bg-gray-100"
              }`}
            >
              ←
            </button>
            <button
              onClick={nextImage}
              disabled={currentImageIndex === images.length - 1}
              className={`px-4 py-2 rounded ${
                currentImageIndex === images.length - 1
                  ? "text-gray-400 cursor-not-allowed"
                  : "text-gray-700 hover:bg-gray-100"
              }`}
            >
              →
            </button>
          </div>
          <span className="text-sm text-gray-600 mt-2 block text-center">
            Imagen {currentImageIndex + 1} de {images.length}
          </span>
        </div>
      </div>
    );
  };

  return (
    <>
      {isUser ? (
        <div className="m-5 mr-5 flex flex-row-reverse items-start space-x-4">
          <div className="avatar placeholder mx-4 w-14 -mt-1">
            <img src={userAvatar} alt="user avatar icon" />
          </div>
          <div className="inline-block max-w-[82%] 2xl:max-w-[88%]">
            <div className="border border-slate-400 rounded-3xl bg-white p-4 text-gray-500">
              <div className="whitespace-pre-wrap text-left">
                {loading && (
                  <span className="loading loading-dots loading-md"></span>
                )}
                <Markdown rehypePlugins={[rehypeRaw]}>{buff}</Markdown>
              </div>
            </div>
          </div>
        </div>
      ) : (
        <div className="m-5 flex items-start space-x-4 w-full">
          <div className="avatar placeholder mx-4 w-14 -mt-1 mr-2">
            <img src={botAvatar} alt="bot avatar icon" />
          </div>
          <div className="inline-block max-w-[82%] 2xl:max-w-[88%]">
            <div className="border-2 border-red-500 rounded-3xl bg-white p-4 text-gray-500 pl-6">
              <div className="flex flex-col items-start">
                <div className="text-left w-full">
                  {loading && (
                    <>
                      {withDeepResearch ? (
                        <div className="flex items-center justify-center gap-2 w-full my-2">
                          <span className="loading loading-spinner loading-md"></span>
                          <span className="text-gray-800 text-m font-medium">
                            Pensamiento profundo...
                          </span>
                        </div>
                      ) : (
                        <span className="loading loading-dots loading-md"></span>
                      )}
                    </>
                  )}
                  <Markdown
                    rehypePlugins={[rehypeRaw]}
                    components={{
                      h1: ({ ...props }) => (
                        <h1 className="text-2xl font-bold mb-4" {...props} />
                      ),
                      h2: ({ ...props }) => (
                        <h2 className="text-xl font-bold mb-3" {...props} />
                      ),
                      h3: ({ ...props }) => (
                        <h3 className="text-lg font-bold mb-2" {...props} />
                      ),
                      p: ({ ...props }) => <p className="mb-4" {...props} />,
                      ul: ({ ...props }) => (
                        <ul
                          className="list-disc pl-6 mb-4 space-y-2"
                          {...props}
                        />
                      ),
                      ol: ({ ...props }) => (
                        <ol
                          className="list-decimal pl-6 mb-4 space-y-2"
                          {...props}
                        />
                      ),
                      li: ({ ...props }) => <li className="mb-1" {...props} />,
                      a: ({ ...props }) => (
                        <a
                          className="text-blue-600 underline hover:text-blue-800"
                          {...props}
                        />
                      ),
                      strong: ({ ...props }) => (
                        <strong className="font-bold" {...props} />
                      ),
                    }}
                  >
                    {buff}
                  </Markdown>
                  <ImageViewer />
                </div>
                {streamingComplete && acceptFeedback && onFeedback && (
                  <FeedbackButton
                    messageKey={responseId}
                    onFeedback={onFeedback}
                  />
                )}
              </div>
            </div>
          </div>
        </div>
      )}
    </>
  );
}