/* Copyright 2026 Marimo. All rights reserved. */

import { useChat } from "@ai-sdk/react";
import { EditorView } from "@codemirror/view";
import {
  CircleCheckIcon,
  Loader2Icon,
  SparklesIcon,
  XIcon,
} from "lucide-react";
import React, { useCallback, useEffect, useId, useState } from "react";
import CodeMirrorMerge from "react-codemirror-merge";
import { DefaultChatTransport } from "ai";
import { Button } from "@/components/ui/button";
import { customPythonLanguageSupport } from "@/core/codemirror/language/languages/python";

import "./merge-editor.css";
import { storePrompt } from "@marimo-team/codemirror-ai";
import type { ReactCodeMirrorRef } from "@uiw/react-codemirror";
import { useAtom } from "jotai";
import { AIModelDropdown } from "@/components/ai/ai-model-dropdown";
import {
  AddContextButton,
  SendButton,
} from "@/components/chat/chat-components";
import { Checkbox } from "@/components/ui/checkbox";
import { Label } from "@/components/ui/label";
import { Switch } from "@/components/ui/switch";
import { Tooltip } from "@/components/ui/tooltip";
import { toast } from "@/components/ui/use-toast";
import {
  CELL_COMPLETION_DATA_TYPE,
  cellCompletionSchema,
  type CompletionUIMessage,
} from "@/core/ai/completion-output";
import { AI_SDK_UI_THROTTLE_MS } from "@/core/ai/constants";
import { type AiCompletionCell, includeOtherCellsAtom } from "@/core/ai/state";
import type { CellId } from "@/core/cells/ids";
import { getCodes } from "@/core/codemirror/copilot/getCodes";
import type { LanguageAdapterType } from "@/core/codemirror/language/types";
import { selectAllText } from "@/core/codemirror/utils";
import { useRuntimeManager } from "@/core/runtime/config";
import type { AiCompletionRequest } from "@/core/network/types";
import { useTheme } from "@/theme/useTheme";
import { cn } from "@/utils/cn";
import { prettyError } from "@/utils/errors";
import { retryWithTimeout } from "@/utils/timeout";
import { PromptInput } from "./add-cell-with-ai";
import {
  AcceptCompletionButton,
  AI_EDIT_ACTION_LABELS,
  createAiCompletionOnKeydown,
  RejectCompletionButton,
} from "./completion-handlers";
import { addContextCompletion, getAICompletionBody } from "./completion-utils";
import { useStagedAICell } from "@/core/ai/staged-cells";

const Original = CodeMirrorMerge.Original;
const Modified = CodeMirrorMerge.Modified;

interface Props {
  cellId: CellId;
  aiCompletionCell: AiCompletionCell | null;
  className?: string;
  currentCode: string;
  currentLanguageAdapter: LanguageAdapterType | undefined;
  onChange: (code: string) => void;
  declineChange: () => void;
  acceptChange: (rightHandCode: string) => void;
  runCell: () => void;
  outputArea?: "above" | "below";
  /**
   * Children shown when there is no completion
   */
  children: React.ReactNode;
}

const baseExtensions = [customPythonLanguageSupport(), EditorView.lineWrapping];

/**
 * Editor for AI completions that goes above a cell to modify it.
 *
 * This shows a left/right split with the original and modified code.
 */
export const AiCompletionEditor: React.FC<Props> = ({
  cellId,
  aiCompletionCell,
  className,
  onChange,
  currentLanguageAdapter,
  currentCode,
  declineChange,
  acceptChange,
  runCell,
  outputArea,
  children,
}) => {
  const [showInputPrompt, setShowInputPrompt] = useState(false);

  const [includeOtherCells, setIncludeOtherCells] = useAtom(
    includeOtherCellsAtom,
  );
  const includeOtherCellsCheckboxId = useId();

  const runtimeManager = useRuntimeManager();

  const {
    initialPrompt,
    triggerImmediately,
    cellId: aiCellId,
  } = aiCompletionCell ?? {};
  const enabled = aiCellId === cellId;

  const updatedCell = useStagedAICell(cellId);
  let previousCellCode: string | undefined;
  if (updatedCell?.type === "update_cell") {
    previousCellCode = updatedCell.previousCode;
  }

  const [completion, setCompletion] = useState("");
  const [input, setInput] = useState(initialPrompt ?? "");
  const transport = React.useMemo(
    () =>
      new DefaultChatTransport<CompletionUIMessage>({
        api: runtimeManager.getAiURL("completion").toString(),
        headers: () => runtimeManager.headers(),
        prepareSendMessagesRequest: ({ body, id }) => ({
          body: { ...body, id },
        }),
      }),
    [runtimeManager],
  );
  const {
    sendMessage,
    stop: stopChat,
    status,
  } = useChat<CompletionUIMessage>({
    throttle: AI_SDK_UI_THROTTLE_MS,
    transport,
    onData: (part) => {
      if (part.type === CELL_COMPLETION_DATA_TYPE) {
        const completion = cellCompletionSchema.parse(part.data);
        setCompletion(completion.code.trimEnd());
      }
    },
    onError: (error) => {
      setCompletion("");
      toast({
        title: "Completion failed",
        description: prettyError(error),
      });
    },
    onFinish: ({ isAbort, isDisconnect, isError, finishReason }) => {
      if (isAbort || isDisconnect || isError || finishReason !== "stop") {
        setCompletion("");
      }
    },
  });
  const isLoading = status === "submitted" || status === "streaming";

  const stop = useCallback(() => {
    setCompletion("");
    void stopChat();
  }, [stopChat]);

  const complete = useCallback(
    async (prompt: string) => {
      await stopChat();
      setCompletion("");

      const body = {
        ...getAICompletionBody({ input: prompt }),
        prompt,
        includeOtherCode: includeOtherCells ? getCodes(currentCode) : "",
        code: currentCode,
        language: currentLanguageAdapter ?? "python",
      } satisfies AiCompletionRequest;

      await sendMessage({ text: prompt }, { body });
    },
    [
      currentCode,
      currentLanguageAdapter,
      includeOtherCells,
      sendMessage,
      stopChat,
    ],
  );

  const handleSubmit = useCallback(() => {
    void complete(input);
  }, [complete, input]);

  const inputRef = React.useRef<ReactCodeMirrorRef>(null);

  const initialSubmit = useCallback(() => {
    if (triggerImmediately && !isLoading && initialPrompt) {
      // Use complete to pass the prompt directly, else input might be empty
      void complete(initialPrompt);
    }
    // oxlint-disable-next-line react-hooks/exhaustive-deps
  }, [triggerImmediately]);

  // Focus the input
  useEffect(() => {
    if (enabled) {
      retryWithTimeout(
        () => {
          const input = inputRef.current;
          if (input?.view) {
            input.view.focus();
            initialSubmit();
            return true;
          }
          return false;
        },
        { retries: 3, delay: 100, initialDelay: 100 },
      ); // Wait for animation to complete

      selectAllText(inputRef.current?.view);
    }
  }, [enabled, initialSubmit]);

  // Reset the input when the prompt changes
  useEffect(() => {
    if (enabled) {
      setInput(initialPrompt || "");
    }
  }, [enabled, initialPrompt, setInput]);

  const { theme } = useTheme();

  const handleAcceptCompletion = () => {
    acceptChange(completion);
    setCompletion("");
  };

  // Reject discards the suggestion but keeps the prompt open for refinement.
  const handleDeclineCompletion = () => {
    stop();
    setShowInputPrompt(true);
    inputRef.current?.view?.focus();
  };

  const showCompletionBanner =
    enabled && triggerImmediately && (completion || isLoading);
  // Set default output area to below if not specified
  outputArea = outputArea ?? "below";

  const showInput = enabled && (!triggerImmediately || showInputPrompt);

  const completionBanner = (
    <div
      className={cn(
        "w-full bg-(--cm-background) flex justify-center transition-all duration-300 ease-in-out overflow-hidden",
        showCompletionBanner
          ? "max-h-20 opacity-100 translate-y-0"
          : "max-h-0 opacity-0 -translate-y-2",
      )}
    >
      <CompletionBanner
        status={isLoading ? "loading" : "generated"}
        onAccept={handleAcceptCompletion}
        onReject={handleDeclineCompletion}
        showInputPrompt={showInputPrompt}
        setShowInputPrompt={setShowInputPrompt}
        runCell={runCell}
        className="mt-4 mb-3 w-lg"
      />
    </div>
  );

  const renderMergeEditor = (originalCode: string, modifiedCode: string) => {
    return (
      <CodeMirrorMerge className="cm" theme={theme}>
        <Original
          onChange={onChange}
          value={originalCode}
          extensions={baseExtensions}
        />
        <Modified
          value={modifiedCode}
          editable={false}
          readOnly={true}
          extensions={baseExtensions}
        />
      </CodeMirrorMerge>
    );
  };

  const renderCompletionEditor = () => {
    if (completion && enabled) {
      return renderMergeEditor(currentCode, completion);
    }
    // If there is no completion and there is previous cell code, it means there is an AI change to the cell.
    // And we want to render the previous cell code as the original
    if (!completion && previousCellCode) {
      return renderMergeEditor(previousCellCode, currentCode);
    }
  };

  const completionButtons = (
    <>
      <AcceptCompletionButton
        isLoading={isLoading}
        onAccept={handleAcceptCompletion}
        size="xs"
        multipleCompletions={false}
        label={AI_EDIT_ACTION_LABELS.accept}
        acceptShortcut="Mod-↵"
        runCell={runCell}
      />
      <RejectCompletionButton
        onDecline={handleDeclineCompletion}
        size="xs"
        multipleCompletions={false}
        label={AI_EDIT_ACTION_LABELS.decline}
        declineShortcut="Shift-Mod-Delete"
      />
    </>
  );

  return (
    <div
      data-ai-input-open={showInput}
      className={cn("flex flex-col w-full rounded-[inherit]", className)}
    >
      <div
        className={cn(
          "flex items-center gap-2 px-3 transition-all rounded-[inherit] rounded-b-none duration-300",
          showInput && "max-h-[400px] border-b min-h-11 visible",
          !showInput && "max-h-0 min-h-0 invisible",
        )}
      >
        {enabled && (
          <>
            <SparklesIcon className="text-(--blue-10) shrink-0" size={16} />
            <PromptInput
              inputRef={inputRef}
              className="h-full my-0 py-2 flex items-center"
              onClose={() => {
                stop();
                declineChange();
              }}
              value={input}
              onChange={setInput}
              onSubmit={() => {
                if (!isLoading) {
                  if (inputRef.current?.view) {
                    storePrompt(inputRef.current.view);
                  }
                  handleSubmit();
                }
              }}
              onKeyDown={createAiCompletionOnKeydown({
                handleAcceptCompletion,
                handleDeclineCompletion,
                isLoading,
                hasCompletion: completion.trim().length > 0,
              })}
            />

            <div className="-mr-1.5 py-1.5">
              <div className="flex flex-row items-center justify-end gap-0.5">
                <SendButton
                  isLoading={isLoading}
                  onStop={stop}
                  onSendClick={handleSubmit}
                  isEmpty={!input.trim()}
                  showStopLabel={true}
                />
                <AddContextButton
                  handleAddContext={() => addContextCompletion(inputRef)}
                  isLoading={isLoading}
                />
                <AIModelDropdown
                  triggerClassName="h-7 text-xs"
                  iconSize="small"
                  forRole="edit"
                  displayIconOnly={true}
                  placeholder="Edit model"
                />
              </div>
              {completion && (
                <div className="mt-1 flex items-center gap-1">
                  {completionButtons}
                </div>
              )}
            </div>

            <div className="h-full w-px bg-border mx-2" />
            <Tooltip content="Include code from other cells">
              <div className="flex flex-row items-start gap-1 overflow-hidden">
                <Checkbox
                  data-testid="include-other-cells-checkbox"
                  id={includeOtherCellsCheckboxId}
                  checked={includeOtherCells}
                  onCheckedChange={(checked) =>
                    setIncludeOtherCells(Boolean(checked))
                  }
                />
                <Label
                  htmlFor={includeOtherCellsCheckboxId}
                  className="text-muted-foreground text-xs whitespace-nowrap ellipsis"
                >
                  All code
                </Label>
              </div>
            </Tooltip>
            <Button
              data-testid="decline-completion-button"
              variant="text"
              size="icon"
              disabled={isLoading}
              onClick={() => {
                stop();
                declineChange();
              }}
            >
              <XIcon className="text-(--red-10)" size={16} />
            </Button>
          </>
        )}
      </div>
      {outputArea === "above" && completionBanner}
      {renderCompletionEditor()}
      {(!completion || !enabled) && !previousCellCode && children}
      {/* By default, show the completion banner below the code */}
      {outputArea === "below" && completionBanner}
    </div>
  );
};

interface CompletionBannerProps {
  status: "loading" | "generated";
  onAccept: () => void;
  onReject: () => void;
  showInputPrompt: boolean;
  setShowInputPrompt: (show: boolean) => void;
  runCell: () => void;
  className?: string;
}

const CompletionBanner: React.FC<CompletionBannerProps> = ({
  status,
  onAccept,
  onReject,
  className,
  showInputPrompt,
  setShowInputPrompt,
  runCell,
}) => {
  const isLoading = status === "loading";

  return (
    <div
      className={cn(
        "flex flex-row items-center gap-6 rounded-md py-2 px-2.5 text-sm border border-border bg-(--gray-2) dark:bg-(--gray-3)",
        className,
      )}
    >
      <div className="flex flex-row items-center gap-2">
        {isLoading ? (
          <Loader2Icon
            className="animate-spin text-blue-600 mb-px"
            size={15}
            strokeWidth={2}
            aria-label="Generating fix"
          />
        ) : (
          <CircleCheckIcon
            className="text-green-600 mb-px"
            size={15}
            strokeWidth={2}
            aria-label="Fix generated"
          />
        )}

        <p className="transition-opacity duration-200 text-muted-foreground">
          {isLoading ? "Generating fix..." : "Showing fix"}
        </p>
      </div>

      <div className="flex flex-row items-center gap-1">
        <Label
          htmlFor="show-input-prompt"
          className="text-muted-foreground text-xs whitespace-nowrap ellipsis"
        >
          Show prompt
        </Label>
        <Switch
          checked={showInputPrompt}
          onCheckedChange={setShowInputPrompt}
          size="xs"
        />
      </div>

      <div className="flex flex-row items-center gap-2 ml-auto">
        <AcceptCompletionButton
          isLoading={isLoading}
          onAccept={onAccept}
          size="xs"
          label={AI_EDIT_ACTION_LABELS.accept}
          runCell={runCell}
        />
        <RejectCompletionButton
          onDecline={onReject}
          size="xs"
          label={AI_EDIT_ACTION_LABELS.decline}
        />
      </div>
    </div>
  );
};