3 files changed, 105 insertions, 24 deletions

diff --git a/src/background.ts b/src/background.ts
index 997dd68..48d51a7 100644
--- a/src/background.ts
+++ b/src/background.ts
@@ -1,21 +1,4 @@
-type Identity = {
-  id?: string;
-  email?: string;
-  name?: string;
-};
-
-type ComposeContext = {
-  account: Identity;
-  compose: {
-    subject?: string;
-    to?: string[];
-    cc?: string[];
-    bcc?: string[];
-    bodyPlain?: string;
-    bodyHTML?: string;
-    identityId?: string;
-  };
-};
+import { ComposeContext } from "./types";
 
 async function getComposeContext(): Promise<ComposeContext> {
   const data: ComposeContext = { account: {}, compose: {} };
@@ -34,6 +17,7 @@ async function getComposeContext(): Promise<ComposeContext> {
     data.compose.identityId = details?.identityId ?? undefined;
     data.compose.bodyPlain = details?.plainTextBody ?? "";
     data.compose.bodyHTML = details?.body ?? "";
+    data.compose.isHtml = Boolean(details?.body && details?.body.trim());
 
     const accounts: any[] = (browser as any).accounts?.list
       ? await (browser as any).accounts.list()
diff --git a/src/popup/popup.ts b/src/popup/popup.ts
index 78866e9..dc28bdc 100644
--- a/src/popup/popup.ts
+++ b/src/popup/popup.ts
@@ -1,18 +1,94 @@
-type Payload = {
+import { ComposeContext } from "../types";
+
+interface Payload {
   prompt: string;
-  context: any;
-};
+  context: ComposeContext;
+}
 
 const promptEl: HTMLTextAreaElement = document.getElementById("prompt") as HTMLTextAreaElement;
 const sendBtn: HTMLButtonElement = document.getElementById("send") as HTMLButtonElement;
 const outputEl: HTMLElement = document.getElementById("output") as HTMLElement;
 
-sendBtn.addEventListener("click", async (): Promise<void> => {
-  const ctx: any = await browser.runtime.sendMessage({ type: "getComposeContext" });
+async function getActiveComposeTabId(): Promise<number | undefined> {
+  const tabs: browser.tabs.Tab[] = (await browser.tabs.query({
+    active: true,
+    currentWindow: true,
+  })) as browser.tabs.Tab[];
+  return tabs[0]?.id;
+}
+
+function insertAtReply(body: string, aiText: string, isHtml: boolean): string {
+  if (isHtml) {
+    const parser: DOMParser = new DOMParser();
+    const doc: Document = parser.parseFromString(body, "text/html");
+    const paragraphs: NodeListOf<HTMLParagraphElement> = doc.querySelectorAll("body > p");
+
+    let target: HTMLParagraphElement | null = null;
+    for (const p of Array.from(paragraphs)) {
+      if (p.textContent && p.textContent.trim().length > 0) {
+        target = p;
+        break;
+      }
+    }
+
+    if (target) {
+      target.insertAdjacentHTML("afterend", aiText);
+      return doc.body.innerHTML;
+    }
+
+    const citeDiv: HTMLElement | null = doc.querySelector("div.moz-cite-prefix");
+    if (citeDiv) {
+      citeDiv.insertAdjacentHTML("beforebegin", aiText);
+      return doc.body.innerHTML;
+    }
+
+    doc.body.insertAdjacentHTML("afterbegin", aiText);
+    return doc.body.innerHTML;
+  } else {
+    const replyEnd: number = body.indexOf("\nOn ");
+    if (replyEnd !== -1) {
+      return body.slice(0, replyEnd).trimEnd() + "\n" + aiText + "\n" + body.slice(replyEnd);
+    }
+    return aiText + "\n\n" + body;
+  }
+}
+
+async function handleInsert(): Promise<void> {
+  const ctx: ComposeContext = (await browser.runtime.sendMessage({
+    type: "getComposeContext",
+  })) as ComposeContext;
+
   const payload: Payload = {
     prompt: promptEl.value,
     context: ctx,
   };
+
+  const tabId: number | undefined = await getActiveComposeTabId();
+  if (tabId === undefined) {
+    outputEl.textContent = "No active compose tab.";
+    return;
+  }
+
+  const isHtml: boolean =
+    ctx.compose.isHtml ?? Boolean(ctx.compose.bodyHTML && ctx.compose.bodyHTML.trim());
+  const aiResult: string = isHtml
+    ? "<p>This response will be inserted by AI</p>"
+    : "This response will be inserted by AI";
+
+  const replyPrefix: string | undefined =
+    ctx.compose.bodyPlain?.split("\n\n")[0].trim() || undefined;
+  const oldBody: string = isHtml ? (ctx.compose.bodyHTML ?? "") : (ctx.compose.bodyPlain ?? "");
+
+  const newBody: string = insertAtReply(oldBody, aiResult, isHtml);
+
+  await (browser as any).compose.setComposeDetails(tabId, { body: newBody });
+
   outputEl.textContent = JSON.stringify(payload, null, 2);
-  alert(JSON.stringify(payload, null, 2));
+}
+
+sendBtn.addEventListener("click", (): void => {
+  handleInsert().catch((err: unknown): void => {
+    console.error("Popup insert failed:", err);
+    outputEl.textContent = "Error inserting text. See console.";
+  });
 });
diff --git a/src/types.ts b/src/types.ts
new file mode 100644
index 0000000..2f6c00a
--- /dev/null
+++ b/src/types.ts
@@ -0,0 +1,21 @@
+export interface Identity {
+  id?: string;
+  email?: string;
+  name?: string;
+}
+
+export interface ComposeDetails {
+  subject?: string;
+  to?: string[];
+  cc?: string[];
+  bcc?: string[];
+  bodyPlain?: string;
+  bodyHTML?: string;
+  identityId?: string;
+  isHtml?: boolean;
+}
+
+export interface ComposeContext {
+  account: Identity;
+  compose: ComposeDetails;
+}