import React, { useEffect, useMemo, useRef, useState } from "react";
import {
  useCurrentConversationState,
  useCurrentPartialContent,
} from "../../../contexts/ChatContext";
import { ParsedBlock, parseMarkdownAndCode } from "../../../utils/parseMarkdownAndCode";
import MarkdownBlock from "../Blocks/MarkdownBlock";
import CodeBlock from "../Blocks/CodeBlock";
import AssistantThinkingSummary from "./AssistantThinkingSummary";

// Component to render a list of completed blocks.
const CompletedBlocks: React.FC<{ blocks: ParsedBlock[] }> = React.memo(({ blocks }) => {
  if (blocks.length === 0) return null;
  return (
    <>
      {blocks.map((block) => {
        if (block.type === "markdown") {
          return (
            <MarkdownBlock
              key={`${block.startIndex}-${block.endIndex}`}
              content={block.content}
            />
          );
        }
        if (block.type === "code") {
          return (
            <CodeBlock
              key={`${block.startIndex}-${block.endIndex}`}
              content={block.content}
              language={block.language}
            />
          );
        }
        return null;
      })}
    </>
  );
});

const StreamingAssistantMessage: React.FC = () => {
  const conversationState = useCurrentConversationState();
  // All content received from the server so far
  const fullContent = useCurrentPartialContent();
  // How many characters make up completed blocks
  const finalizedLengthRef = useRef<number>(0);
  // The completed blocks (don't need to be re-parsed or rendered)
  const [completedBlocks, setCompletedBlocks] = useState<ParsedBlock[]>([]);
  // The latest block (is re-parsed until completed)
  const [latestBlock, setLatestBlock] = useState<ParsedBlock | null>(null);

  // Each time the fullContent changes, we update our incremental parser.
  useEffect(() => {
    if (!fullContent) {
      finalizedLengthRef.current = 0;
      setCompletedBlocks([]);
      setLatestBlock(null);
      return;
    }
    // Compute the string making up the latest block
    const textToParse = fullContent.slice(finalizedLengthRef.current);
    // Parse the text
    const parsed = parseMarkdownAndCode(textToParse);

    // If nothing is parsed, do nothing.
    if (parsed.length === 0) return;

    // If parser returns >1 block, mark all but the last block as complete
    if (parsed.length > 1) {
      const completeBlocks = parsed.slice(0, parsed.length - 1);
      const lastCompleteBlock = completeBlocks[completeBlocks.length - 1];

      // Update completedBlocks with these new complete blocks.
      setCompletedBlocks((prev) => [...prev, ...completeBlocks]);

      // Update the finalized length: add the length (endIndex) of the last complete block
      // Note: since parseMarkdownAndCode works on the newSegment, its indices are relative to newSegment.
      // To get the absolute index, add the current finalizedLength.
      const newFinalized = finalizedLengthRef.current + lastCompleteBlock.endIndex;
      finalizedLengthRef.current = newFinalized;

      // The current (latest) block is the last block returned.
      setLatestBlock(parsed[parsed.length - 1]);
    } else {
      // Only one block was parsed: update latestBlock.
      setLatestBlock(parsed[0]);
    }
  }, [fullContent]);

  const timeline = conversationState?.timeline ?? [];
  const hasVisibleTimeline = timeline.some(entry => entry.kind !== 'reasoning' || entry.text.trim());
  const showExpandedStatus = !conversationState?.hasStartedStreaming;
  const thoughtDurationMs = conversationState?.thoughtDurationMs ?? (
    conversationState?.startedAt ? Math.max(1, Date.now() - conversationState.startedAt) : 0
  );

  // Memoize the rendering of completed blocks.
  const renderedCompletedBlocks = useMemo(() => {
    return <CompletedBlocks blocks={completedBlocks} />;
  }, [completedBlocks]);

  // Memoize rendering for the current (partial) block.
  const renderedLatestBlock = useMemo(() => {
    if (!latestBlock) return null;
    if (latestBlock.type === "markdown") {
      return (
        <MarkdownBlock
          key={`${latestBlock.startIndex}-partial`}
          content={latestBlock.content}
        />
      );
    }
    if (latestBlock.type === "code") {
      return (
        <CodeBlock
          key={`${latestBlock.startIndex}-partial`}
          content={latestBlock.content}
          language={latestBlock.language}
        />
      );
    }
    return null;
  }, [latestBlock]);

  return (
    <div className="chat-output assistant-message px-4 py-3 my-2 rounded-[1.3rem] overflow-hidden max-w-full text-gray-700 dark:text-gray-300">
      {(showExpandedStatus || hasVisibleTimeline) && (
        <AssistantThinkingSummary
          label={
            showExpandedStatus
              ? (conversationState?.statusLabel || 'Thinking')
              : `Thought for ${Math.max(1, Math.round(thoughtDurationMs / 1000))}s`
          }
          timeline={timeline}
          isActive={showExpandedStatus}
        />
      )}

      <div className="chat-content max-w-none">
        {renderedCompletedBlocks}
        {renderedLatestBlock}
      </div>
    </div>
  );
};

export default React.memo(StreamingAssistantMessage);