import React, { useEffect, useRef, useState } from "react";
import {
  ArrowsOutSimple,
  ArrowsInSimple,
  CaretRight,
  Trash,
  ChatCircle,
  Plus,
  SpinnerGap,
  CaretDown,
} from "@phosphor-icons/react";
import { api, streamChat, readPreference, savePreference } from "./api.js";
import { handleTabListKeyDown, IconButton, Modal } from "./ui.jsx";
import { useConfirmation } from "./hooks/useConfirmation.js";
import { ConfirmationDialog } from "./ConfirmationDialog.jsx";
import { AssistantTranscript } from "./AssistantTranscript.jsx";
import { AssistantComposer } from "./AssistantComposer.jsx";
import { AssistantSources } from "./AssistantSources.jsx";
import { ChatSwitcher } from "./ChatSwitcher.jsx";

export function Assistant({
  material,
  providers,
  context,
  onClearContext,
  currentPage,
  screenshotRef,
  onClose,
  expanded,
  onExpand,
  onSettings,
  onPage,
  onRemoveSource,
  onRemoveMaterial,
  notify,
  onSaveNote,
  draftRequest,
  chatRequest,
  comparison,
  onComparisonPage,
}) {
  const [tab, setTab] = useState("chats"),
    [chats, setChats] = useState([]),
    [selected, setSelected] = useState(""),
    [cache, setCache] = useState({}),
    [drafts, setDrafts] = useState({}),
    [busy, setBusy] = useState({}),
    [errors, setErrors] = useState({}),
    [creating, setCreating] = useState(false),
    [listOpen, setListOpen] = useState(false),
    [contextOpen, setContextOpen] = useState(false),
    [loading, setLoading] = useState(false),
    [renaming, setRenaming] = useState(null),
    [renameTitle, setRenameTitle] = useState(""),
    [savingTitle, setSavingTitle] = useState(false),
    [providerId, setProviderId] = useState(readPreference("provider", "")),
    [model, setModel] = useState(readPreference("model", "")),
    [reasoningEffort, setReasoningEffort] = useState(() => {
      const stored = readPreference("reasoningEffort", "auto");
      return ["auto", "low", "medium", "high"].includes(stored)
        ? stored
        : "auto";
    }),
    [attachImage, setAttachImage] = useState(false),
    [below, setBelow] = useState(false);
  const alive = useRef(true),
    draftsRef = useRef({}),
    controllers = useRef(new Map()),
    handledRequest = useRef(null),
    handledChatRequest = useRef(null),
    midRef = useRef(material?.id),
    chatTriggerRef = useRef(null),
    transcript = useRef(),
    inputRef = useRef(),
    nearBottom = useRef(true);
  midRef.current = material?.id;
  draftsRef.current = drafts;
  const { confirm, dialogProps } = useConfirmation();
  const provider = providers.find((p) => p.id === providerId),
    list = cache[selected] || [],
    draft = drafts[selected || "new"] || "";
  useEffect(() => {
    if (!provider?.options.vision) setAttachImage(false);
  }, [provider?.options.vision]);
  useEffect(() => {
    alive.current = true;
    return () => { alive.current = false; for (const controller of controllers.current.values()) controller.abort(); };
  }, []);
  const updateDrafts = (change) => {
    if (!alive.current) return;
    const next = typeof change === "function" ? change(draftsRef.current) : change;
    draftsRef.current = next;
    setDrafts(next);
    if (midRef.current) savePreference(`drafts-${midRef.current}`, next);
  };
  const refreshChats = async (mid = material?.id) => {
    if (!mid) return;
    const next = await api("/materials/" + mid + "/chats");
    if (!alive.current || midRef.current !== mid) return;
    setChats(next);
    const remembered = readPreference("chat-" + mid, "");
    setSelected((v) =>
      next.some((c) => c.id === v)
        ? v
        : next.some((c) => c.id === remembered)
          ? remembered
          : next[0]?.id || "",
    );
  };
  useEffect(() => {
    for (const c of controllers.current.values()) c.abort();
    controllers.current.clear();
    setCache({});
    setErrors({});
    setBusy({});
    setChats([]);
    setSelected("");
    const restoredDrafts = material?.id ? readPreference(`drafts-${material.id}`, {}) : {};
    draftsRef.current = restoredDrafts;
    setDrafts(restoredDrafts);
    if (material?.id)
      refreshChats(material.id).catch((e) => notify(e.message, "error"));
    return () => {
      for (const c of controllers.current.values()) c.abort();
    };
  }, [material?.id]);
  useEffect(() => {
    if (!providers.length) return;
    if (!provider) {
      setProviderId(
        (providers.find((p) => p.hasKey && p.models.length) || providers[0])
          ?.id || "",
      );
      return;
    }
    if (!provider.models.includes(model)) setModel(provider.models[0] || "");
  }, [providers, providerId, model]);
  useEffect(() => {
    if (!selected || cache[selected]) {
      setLoading(false);
      return;
    }
    let live = true;
    setLoading(true);
    api("/chats/" + selected + "/messages")
      .then((rows) => {
        if (live) setCache((v) => ({ ...v, [selected]: rows }));
      })
      .catch((e) => notify(e.message, "error"))
      .finally(() => {
        if (live) setLoading(false);
      });
    return () => {
      live = false;
    };
  }, [selected, material?.id]);
  useEffect(() => {
    setContextOpen(false);
    if (context) {
      setTab("chats");
      setAttachImage(
        material?.status === "image-only" && !!provider?.options.vision,
      );
      setTimeout(() => inputRef.current?.focus(), 50);
    }
  }, [context]);
  useEffect(() => {
    if (!transcript.current) return;
    if (nearBottom.current) {
      transcript.current.scrollTop = transcript.current.scrollHeight;
      setBelow(false);
    } else setBelow(true);
  }, [list, selected]);
  const chooseChat = (cid) => {
    savePreference("chat-" + material.id, cid);
    setSelected(cid);
    setLoading(false);
    nearBottom.current = true;
    setBelow(false);
    const c = chats.find((v) => v.id === cid);
    const p = providers.find((p) => p.id === c?.provider_id);
    if (p) {
      setProviderId(p.id);
      if (p.models.includes(c.model)) setModel(c.model);
    }
  };
  const newChat = async () => {
    if (!material) return;
    setTab("chats");
    setCreating(true);
    const targetMaterialId = material.id;
    try {
      const c = await api("/chats", {
        method: "POST",
        body: { materialId: targetMaterialId },
      });
      if (!alive.current || midRef.current !== targetMaterialId) return;
      setChats((v) => [c, ...v]);
      setCache((v) => ({ ...v, [c.id]: [] }));
      setSelected(c.id);
      savePreference("chat-" + material.id, c.id);
      nearBottom.current = true;
      setTimeout(() => inputRef.current?.focus(), 50);
      return c.id;
    } catch (e) {
      notify(e.message, "error");
    } finally {
      if (alive.current) setCreating(false);
    }
  };
  const updateRows = (cid, fn) =>
    setCache((v) => ({ ...v, [cid]: fn(v[cid] || []) }));
  const send = async (retryUser) => {
    if (!material || creating || busy[selected]) return;
    if (!provider?.hasKey || !model) {
      onSettings();
      notify(
        "Add an API key and at least one model ID to start a conversation.",
      );
      return;
    }
    const question = retryUser ? retryUser.content : draft.trim();
    if (!question) return;
    const mid = material.id;
    setCreating(true);
    const cid = selected || (await newChat());
    setCreating(false);
    if (!cid) return;
    if (!alive.current || midRef.current !== mid) return;
    const original = cache[cid] || [];
    let existing = original;
    if (retryUser)
      existing = existing.slice(
        0,
        existing.findIndex((m) => m.id === retryUser.id) + 1,
      );
    const tempUser = "u" + Date.now(),
      tempAssistant = "a" + Date.now();
    const source = retryUser?.context || {
      selection: context?.quote || "",
      page:
        context?.page || (material.kind === "book" ? currentPage : undefined),
    };
    const image =
      retryUser?.context?.image ||
      (attachImage && material.kind === "book"
        ? screenshotRef.current?.(source.page)
        : undefined);
    if (attachImage && !image && !retryUser) {
      notify(
        "Wait for the PDF page to finish rendering before attaching it.",
        "error",
      );
      return;
    }
    const payload = {
      content: question,
      providerId: provider.id,
      model,
      ...(reasoningEffort !== "auto" && (provider.options?.protocol !== "anthropic" || provider.options?.reasoning === true) ? { reasoningEffort } : {}),
      selection: source.selection || "",
      ...(source.page ? { page: source.page } : {}),
      ...(image ? { image } : {}),
      ...(retryUser ? { retryMessageId: retryUser.id } : {}),
      ...(retryUser?.context?.comparison || comparison
        ? { comparison: retryUser?.context?.comparison || comparison }
        : {}),
    };
    const user = retryUser || {
      id: tempUser,
      role: "user",
      content: question,
      status: "complete",
      context: {
        selection: payload.selection,
        page: payload.page,
        image,
        comparison: payload.comparison,
      },
      created_at: new Date().toISOString(),
    };
    setCache((v) => ({
      ...v,
      [cid]: [
        ...existing,
        ...(retryUser ? [] : [user]),
        {
          id: tempAssistant,
          role: "assistant",
          content: "",
          status: "streaming",
          model,
          citations: [],
        },
      ],
    }));
    updateDrafts((v) => ({ ...v, [cid]: "", new: "" }));
    setBusy((v) => ({ ...v, [cid]: true }));
    setErrors((v) => ({ ...v, [cid]: "" }));
    nearBottom.current = true;
    const controller = new AbortController();
    controllers.current.set(cid, controller);
    let received = false,
      finished = false,
      aid = tempAssistant,
      buffer = "",
      frame;
    const flush = () => {
      if (!buffer || !alive.current || midRef.current !== mid) return;
      const text = buffer;
      buffer = "";
      updateRows(cid, (rows) =>
        rows.map((m) =>
          m.id === aid ? { ...m, content: m.content + text } : m,
        ),
      );
    };
    try {
      await streamChat(cid, payload, controller.signal, (event) => {
        if (!alive.current || midRef.current !== mid) return;
        if (event.type === "context") {
          received = true;
          aid = event.assistantId;
          updateRows(cid, (rows) =>
            rows.map((m) =>
              m.id === tempUser
                ? { ...m, id: event.userId }
                : m.id === tempAssistant
                  ? {
                      ...m,
                      id: aid,
                      citations: event.citations,
                      context: {
                        estimatedTokens: event.estimatedTokens,
                        mode: event.mode,
                      },
                    }
                  : m,
            ),
          );
        }
        if (event.type === "delta") {
          buffer += event.text;
          if (!frame)
            frame = requestAnimationFrame(() => {
              frame = null;
              flush();
            });
        }
        if (event.type === "error")
          setErrors((v) => ({ ...v, [cid]: event.error }));
        if (event.type === "done") {
          finished = true;
          if (frame) cancelAnimationFrame(frame);
          frame = null;
          flush();
          updateRows(cid, (rows) =>
            rows.map((m) =>
              m.id === aid ? { ...m, status: event.status } : m,
            ),
          );
        }
      });
      if (!finished && alive.current && midRef.current === mid) {
        if (frame) cancelAnimationFrame(frame);
        frame = null;
        flush();
        setErrors((current) => ({ ...current, [cid]: "The response ended before completion. Retry this question to continue." }));
        updateRows(cid, (rows) => rows.map((message) => message.id === aid ? { ...message, status: "interrupted" } : message));
      }
    } catch (e) {
      if (!alive.current || midRef.current !== mid) return;
      flush();
      const cancelled = e.name === "AbortError";
      setErrors((v) => ({
        ...v,
        [cid]: cancelled ? "Response stopped." : e.message,
      }));
      if (!received) {
        setCache((v) => ({ ...v, [cid]: original }));
        if (!retryUser) updateDrafts((v) => ({ ...v, [cid]: question }));
      } else
        updateRows(cid, (rows) =>
          rows.map((m) =>
            m.id === aid
              ? { ...m, status: cancelled ? "cancelled" : "interrupted" }
              : m,
          ),
        );
    } finally {
      if (frame) cancelAnimationFrame(frame);
      controllers.current.delete(cid);
      if (alive.current && midRef.current === mid) {
        setBusy((v) => ({ ...v, [cid]: false }));
        refreshChats(mid).catch((e) => notify(e.message, "error"));
      }
    }
  };
  const removeChat = (c) =>
    confirm({
      title: "Delete conversation?",
      description: c.title,
      detail:
        "All messages in this conversation will be deleted. Other conversations and the material are kept.",
      label: "Delete conversation",
      action: async () => {
        controllers.current.get(c.id)?.abort();
        await api("/chats/" + c.id, { method: "DELETE" });
        setCache((v) => {
          const n = { ...v };
          delete n[c.id];
          return n;
        });
        await refreshChats();
        updateDrafts((drafts) => {
          const next = { ...drafts };
          delete next[c.id];
          return next;
        });
        notify("Conversation deleted.");
      },
    });
  const removeMessage = (message) =>
    confirm({
      title: "Delete message?",
      description: message.content.slice(0, 160),
      detail:
        "This message will be removed from the conversation and future chat history. Other messages are kept.",
      label: "Delete message",
      action: async () => {
        await api("/messages/" + message.id, { method: "DELETE" });
        setCache((cache) => ({
          ...cache,
          [selected]: (cache[selected] || []).filter(
            (item) => item.id !== message.id,
          ),
        }));
        await refreshChats();
        notify("Message deleted.");
      },
    });
  const renameChat = (chat) => {
    setRenaming(chat);
    setRenameTitle(chat.title);
  };
  const saveChatTitle = async (event) => {
    event.preventDefault();
    if (!renameTitle.trim() || savingTitle) return;
    setSavingTitle(true);
    try {
      await api("/chats/" + renaming.id, {
        method: "PATCH",
        body: { title: renameTitle.trim() },
      });
      await refreshChats();
      setRenaming(null);
    } catch (error) {
      notify(error.message, "error");
    } finally {
      setSavingTitle(false);
    }
  };
  useEffect(() => {
    if (
      !draftRequest ||
      draftRequest.materialId !== material?.id ||
      handledRequest.current === draftRequest.token
    )
      return;
    handledRequest.current = draftRequest.token;
    setTab("chats");
    (async () => {
      const cid = await newChat();
      if (cid && midRef.current === draftRequest.materialId) {
        updateDrafts((current) => ({ ...current, [cid]: draftRequest.text }));
        setTimeout(() => inputRef.current?.focus(), 100);
      }
    })();
  }, [draftRequest, material?.id]);
  useEffect(() => {
    if (
      !chatRequest ||
      handledChatRequest.current === chatRequest.token ||
      chatRequest.materialId !== material?.id ||
      !chats.some((c) => c.id === chatRequest.chatId)
    )
      return;
    handledChatRequest.current = chatRequest.token;
    setTab("chats");
    chooseChat(chatRequest.chatId);
  }, [chatRequest, chats.length]);
  useEffect(() => {
    if (chatRequest?.chatId !== selected || !cache[selected]) return;
    const frame = requestAnimationFrame(() => {
      transcript.current
        ?.querySelector(
          `[data-message-id="${CSS.escape(chatRequest.messageId || "")}"]`,
        )
        ?.scrollIntoView({ block: "center" });
    });
    return () => cancelAnimationFrame(frame);
  }, [chatRequest, selected, cache[selected]?.length]);
  return (
    <>
      <aside className="assistant-panel panel" aria-label="Assistant">
        <div className="panel-title">
          <span>Assistant</span>
          <div className="assistant-panel-actions">
            <IconButton
              ref={chatTriggerRef}
              label="Choose conversation"
              aria-haspopup="dialog"
              onClick={() => setListOpen(true)}
              disabled={!material}
            >
              <ChatCircle size={18} />
            </IconButton>
            <IconButton
              label="New chat"
              disabled={!material || creating}
              aria-busy={creating}
              onClick={newChat}
            >
              {creating ? (
                <SpinnerGap className="spin" size={18} />
              ) : (
                <Plus size={18} />
              )}
            </IconButton>
            <IconButton
              label={expanded ? "Restore assistant size" : "Expand assistant"}
              aria-pressed={expanded}
              onClick={onExpand}
            >
              {expanded ? (
                <ArrowsInSimple size={18} />
              ) : (
                <ArrowsOutSimple size={18} />
              )}
            </IconButton>
            <IconButton label="Hide assistant panel" onClick={onClose}>
              <CaretRight size={18} />
            </IconButton>
          </div>
        </div>
        <div
          className="segmented panel-tabs"
          role="tablist"
          aria-label="Assistant views"
          data-selected={tab === "sources" ? "second" : "first"}
          onKeyDown={handleTabListKeyDown}
        >
          <button
            role="tab"
            aria-selected={tab === "chats"}
            tabIndex={tab === "chats" ? 0 : -1}
            type="button"
            className={tab === "chats" ? "active" : ""}
            onClick={() => setTab("chats")}
          >
            Chats <span className="count">{chats.length}</span>
          </button>
          <button
            role="tab"
            aria-selected={tab === "sources"}
            tabIndex={tab === "sources" ? 0 : -1}
            type="button"
            className={tab === "sources" ? "active" : ""}
            onClick={() => setTab("sources")}
          >
            Sources{" "}
            <span className="count">
              {material?.kind === "book" ? 1 : material?.sources?.length || 0}
            </span>
          </button>
        </div>
        {tab === "sources" ? (
          <AssistantSources
            material={material}
            onRemoveSource={onRemoveSource}
            onRemoveMaterial={onRemoveMaterial}
          />
        ) : (
          <>
            <AssistantTranscript
              loading={loading}
              messages={list}
              material={material}
              busy={busy}
              selectedChat={selected}
              errors={errors}
              transcriptRef={transcript}
              nearBottom={nearBottom}
              below={below}
              setBelow={setBelow}
              onRemoveMessage={removeMessage}
              onRetry={send}
              onPage={onPage}
              notify={notify}
              onSaveNote={onSaveNote}
              onComparisonPage={onComparisonPage}
            />
            {context && (context.quote || context.page) && (
              <div className="selected-context">
                <div className="context-heading">
                  <button
                    type="button"
                    className="text-button context-toggle"
                    aria-expanded={contextOpen}
                    onClick={() => setContextOpen((open) => !open)}
                  >
                    <CaretDown
                      size={14}
                      className={contextOpen ? "" : "rotate"}
                    />
                    <span>
                      {context.quote
                        ? `Selected passage${context.page ? ` · Page ${context.page}` : ""}`
                        : `Book page ${context.page}`}
                    </span>
                  </button>
                  <IconButton
                    label="Remove selected passage"
                    onClick={onClearContext}
                  >
                    <Trash size={14} />
                  </IconButton>
                </div>
                {contextOpen && (
                  <div className="context-details">
                    {context.quote && (
                      <blockquote>
                        {context.quote.slice(0, 500)}
                        {context.quote.length > 500 ? "…" : ""}
                      </blockquote>
                    )}
                    {context.page && (
                      <button
                        className="text-button"
                        onClick={() => onPage(context.page)}
                      >
                        Page {context.page} · open in reader
                      </button>
                    )}
                  </div>
                )}
              </div>
            )}
            <AssistantComposer
              material={material}
              currentPage={currentPage}
              context={context}
              providers={providers}
              provider={provider}
              providerId={providerId}
              model={model}
              reasoningEffort={reasoningEffort}
              attachImage={attachImage}
              draft={draft}
              selectedChat={selected}
              busy={busy}
              creating={creating}
              inputRef={inputRef}
              onModelChange={({
                providerId: nextProviderId,
                model: nextModel,
              }) => {
                setProviderId(nextProviderId);
                savePreference("provider", nextProviderId);
                setModel(nextModel);
                savePreference("model", nextModel);
              }}
              onReasoningChange={(value) => {
                setReasoningEffort(value);
                savePreference("reasoningEffort", value);
              }}
              onAttachImageChange={setAttachImage}
              onDraftChange={(value) =>
                updateDrafts((current) => ({
                  ...current,
                  [selected || "new"]: value,
                }))
              }
              onSettings={onSettings}
              onStop={() => controllers.current.get(selected)?.abort()}
              onSend={send}
            />
          </>
        )}
      </aside>
      <ChatSwitcher
        triggerRef={chatTriggerRef}
        chats={chats}
        busy={busy}
        selectedChat={selected}
        listOpen={listOpen}
        onOpenChange={setListOpen}
        onChooseChat={(cid) => {
          setTab("chats");
          chooseChat(cid);
        }}
        onRenameChat={renameChat}
        onRemoveChat={removeChat}
      />
      <Modal
        open={!!renaming}
        onOpenChange={(open) => {
          if (!open && !savingTitle) setRenaming(null);
        }}
        title="Rename conversation"
      >
        <form onSubmit={saveChatTitle}>
          <label>
            Title
            <input
              value={renameTitle}
              onChange={(event) => setRenameTitle(event.target.value)}
              maxLength={160}
              required
            />
          </label>
          <div className="form-actions">
            <button
              type="button"
              disabled={savingTitle}
              onClick={() => setRenaming(null)}
            >
              Cancel
            </button>
            <button
              className="primary"
              disabled={!renameTitle.trim() || savingTitle}
            >
              {savingTitle ? "Saving…" : "Save title"}
            </button>
          </div>
        </form>
      </Modal>
      <ConfirmationDialog {...dialogProps} />
    </>
  );
}