import type {
  LogEntry,
  LSPCodeMessage,
  LSPMarkdownMessage,
  LSPTextMessage,
  QueueTaskItem,
  ViewPatchMessage,
} from "@codeflash/types";
import styles from "./taskLogging.module.css";
import MarkdownBlock from "./markdown/markdown";
import CollapsibleCodeBlock from "./codeBlock/collapsibleCodeBlock";
import DemoSuccess from "./demoSuccess";
import { memo, useEffect, useRef } from "react";
import { vscode } from "../utils/vscode";
import {
  Virtuoso,
  type FollowOutputScalarType,
  type VirtuosoHandle,
} from "react-virtuoso";
import { useMemoFunc } from "../utils/useMemoFunc";
import { isTaskRunning } from "@codeflash/shared";

const MemoMarkdownBlock = memo(MarkdownBlock);
const MemoCollapsibleCodeBlock = memo(CollapsibleCodeBlock);
const footerContainerStyles = {
  paddingBottom: "1rem",
  display: "flex",
  justifyContent: "space-between",
  alignItems: "center",
  padding: "0 1rem",
  margin: "1rem 0",
  background: "var(--vscode-editor-background)",
  color: "var(--vscode-editor-foreground)",
  borderRadius: "7px",
};
const LogItem = memo(
  ({
    log,
    isLast,
    isStillRunning,
    functionName,
  }: {
    log: LogEntry;
    isLast: boolean;
    isStillRunning: boolean;
    functionName: string;
    filePath: string | undefined;
  }) => {
    const isLoading = isStillRunning && log.takes_time && isLast;
    const longTaskFinished =
      (log.takes_time && !isLast) || (log.takes_time && !isStillRunning);

    const contentClassName =
      styles.logContent + " " + (isLoading ? styles.textShimmer : "");
    switch (log.type) {
      case "text":
        const textLog = log as LSPTextMessage;
        return (
          <div className={styles.logEntry}>
            {longTaskFinished && (
              <span className="codicon codicon-pass-filled" />
            )}
            <div className={contentClassName}>
              <MemoMarkdownBlock
                markdown={textLog.text}
                currentOptimizedFunction={functionName}
              />
            </div>
          </div>
        );
      case "code":
        const codeLog = log as LSPCodeMessage;
        return (
          <div className={styles.logEntry}>
            <div className={contentClassName}>
              {codeLog.collapsed ? (
                <MemoCollapsibleCodeBlock
                  source={codeLog.code}
                  language="python"
                />
              ) : (
                <MemoCollapsibleCodeBlock
                  source={codeLog.code}
                  language="python"
                />
              )}
            </div>
          </div>
        );
      case "markdown":
        const markdownLog = log as LSPMarkdownMessage;
        return (
          <div className={styles.logEntry}>
            {longTaskFinished && (
              <span className="codicon codicon-pass-filled" />
            )}
            <div className={contentClassName}>
              <MemoMarkdownBlock
                markdown={markdownLog.markdown}
                currentOptimizedFunction={functionName}
              />
            </div>
          </div>
        );
      default:
        return (
          <div className={styles.logEntry}>
            <div className={contentClassName}>
              <MemoMarkdownBlock
                markdown={`| Not supported log type: ${log.type}`}
                currentOptimizedFunction={functionName}
              />
            </div>
          </div>
        );
    }
  },
);

const TaskLogging = ({
  task,
  onDemoSuccess,
}: {
  task: QueueTaskItem;
  onDemoSuccess?: () => void;
}) => {
  const listRef = useRef<VirtuosoHandle>(null);
  const scrollerRef = useRef<HTMLElement>(null);

  const followOutput = useMemoFunc(
    (isAtBottom: boolean): FollowOutputScalarType => {
      if (isAtBottom) {
        setTimeout(() => {
          listRef.current?.autoscrollToBottom();
        }, 20);
      }

      return isAtBottom ? "smooth" : false;
    },
  );

  useEffect(() => {
    listRef.current?.scrollToIndex({ index: task.logs.length - 1 });
  }, []);

  const getFooterOnSuccess = () => {
    return onDemoSuccess ? (
      // completed the demo optimization
      <div className={styles.logEntry}>
        <DemoSuccess onGetStarted={onDemoSuccess} />
      </div>
    ) : (
      <div className={styles.logEntry}>
        <div className={styles.logContent}>
          <div
            style={{
              ...footerContainerStyles,
              color: "var(--vscode-foreground)",
              background: "rgba(34, 134, 58, 0.15)",
            }}
          >
            <MemoMarkdownBlock
              markdown={
                task.optimizationReview
                  ? `✅ Optimization completed successfully. \n **🎯Optimization Quality:** ${task.optimizationReview}`
                  : `✅ Optimization completed successfully.`
              }
            />
            <button
              className={styles.queueTaskPatch}
              onClick={() => {
                const msg: ViewPatchMessage = {
                  type: "viewPatch",
                  payload: {
                    id: task.id,
                    functionName: task.functionName,
                    patchFile: task.patchFile!,
                    explanation: task.explanation || "",
                    speedupStr: task.speedupStr || "",
                    optimizationReview: task.optimizationReview || "",
                    originalRuntime: task.originalRuntime || 0,
                    optimizedRuntime: task.optimizedRuntime || 0,
                    originalLineProfiler: task.originalLineProfiler || "",
                    optimizedLineProfiler: task.optimizedLineProfiler || "",
                  },
                };
                vscode.postMessage(msg);
              }}
            >
              <span className="codicon codicon-eye"></span>
              <span>View Optimization</span>
            </button>
          </div>
        </div>
      </div>
    );
  };
  return (
    <div className={styles.messagesContainer}>
      <Virtuoso
        ref={listRef}
        style={{ height: "100%" }}
        data={task.logs}
        computeItemKey={(_, log) => log.id}
        itemContent={(index) => (
          <LogItem
            log={task.logs[index]!}
            isLast={index === task.logs.length - 1}
            isStillRunning={isTaskRunning(task)}
            functionName={task.functionName}
            filePath={task.filepath}
          />
        )}
        scrollerRef={(ref) => (scrollerRef.current = ref as HTMLElement)}
        // initialTopMostItemIndex={task.logs.length - 1}
        totalCount={task.logs.length}
        overscan={{
          main: 2000,
          reverse: 1500,
        }}
        // atBottomStateChange={atBottomStateChange} // TODO: show scroll to bottom button
        followOutput={followOutput}
        defaultItemHeight={300}
        atTopThreshold={100}
        atBottomThreshold={40}
        useWindowScroll={false}
      />

      {task.error && (
        <div className={styles.logEntry}>
          <div className={styles.logContent} style={footerContainerStyles}>
            <MemoMarkdownBlock
              markdown={`**Couldn't find a useful optimization for this function.**\n${task.error}`}
            />
            <button
              className={styles.ghostActionButton}
              title="Retry"
              onClick={() => {
                vscode.postMessage({
                  type: "optimizeFunction",
                  payload: {
                    functionName: task.functionName || "",
                    filePath: task.filepath || "",
                  },
                });
              }}
            >
              <span>Retry optimization</span>
              <span className="codicon codicon-refresh"></span>
            </button>
          </div>
        </div>
      )}
      {task.status === "completed" && getFooterOnSuccess()}
    </div>
  );
};

export default TaskLogging;
