// Name: OpenAI Playground with websearch and reasoning + latex rendering
// Description: Query Open AI's API with web search and reasoning parsing, plus an addition to md() to also render math with katex
// Keyword: ai
// Author: benhaotang
// Github: benhaotang

import "@johnlindquist/kit";
import OpenAI from "openai";
import katex from "katex";

/**
 * Enhanced markdown renderer with KaTeX math support
 * Supports: $$ ... $$, \[ ... \], \( ... \)
 */
function mdWithMath(markdown) {
  const mathBlocks = [];
  let processedMarkdown = markdown;

  // Process block math: $$ ... $$
  processedMarkdown = processedMarkdown.replace(
    /\$\$([\s\S]*?)\$\$/g,
    (match, math) => {
      const index = mathBlocks.length;
      mathBlocks.push({
        type: "block",
        math: math.trim(),
        placeholder: `XMATHBLOCKX${index}XMATHBLOCKX`,
      });
      return `XMATHBLOCKX${index}XMATHBLOCKX`;
    },
  );

  // Process block math: \[ ... \]
  processedMarkdown = processedMarkdown.replace(
    /\\\[([\s\S]*?)\\\]/g,
    (match, math) => {
      const index = mathBlocks.length;
      mathBlocks.push({
        type: "block",
        math: math.trim(),
        placeholder: `XMATHBLOCKX${index}XMATHBLOCKX`,
      });
      return `XMATHBLOCKX${index}XMATHBLOCKX`;
    },
  );

  // Process inline math: \( ... \)
  processedMarkdown = processedMarkdown.replace(
    /\\\(([\s\S]*?)\\\)/g,
    (match, math) => {
      const index = mathBlocks.length;
      mathBlocks.push({
        type: "inline",
        math: math.trim(),
        placeholder: `XMATHBLOCKX${index}XMATHBLOCKX`,
      });
      return `XMATHBLOCKX${index}XMATHBLOCKX`;
    },
  );

  // Convert markdown to HTML using Script Kit's built-in md() function
  let html = md(processedMarkdown);

  // Replace placeholders with rendered KaTeX (handle both plain and wrapped in tags)
  mathBlocks.forEach((block) => {
    try {
      const rendered = katex.renderToString(block.math, {
        displayMode: block.type === "block",
        throwOnError: false,
        output: "mathml",
        trust: false,
        strict: "warn",
      });

      const mathHtml =
        block.type === "block"
          ? `<div class="my-4 text-center flex justify-center">${rendered}</div>`
          : `<span class="inline-block">${rendered}</span>`;

      // Replace both plain placeholder and wrapped versions
      const plainPattern = new RegExp(
        block.placeholder.replace(/[.*+?^${}()|[\]\\]/g, "\\$&"),
        "g",
      );
      const strongPattern = new RegExp(
        `<strong>${block.placeholder.replace(/[.*+?^${}()|[\]\\]/g, "\\$&")}</strong>`,
        "g",
      );
      const emPattern = new RegExp(
        `<em>${block.placeholder.replace(/[.*+?^${}()|[\]\\]/g, "\\$&")}</em>`,
        "g",
      );

      html = html.replace(strongPattern, mathHtml);
      html = html.replace(emPattern, mathHtml);
      html = html.replace(plainPattern, mathHtml);
    } catch (error) {
      const errorHtml =
        block.type === "block"
          ? `<div class="my-4 p-2 bg-red-100 border border-red-300 rounded text-red-700">Math Error: ${block.math}</div>`
          : `<span class="bg-red-100 text-red-700 px-1 rounded">Math Error: ${block.math}</span>`;

      // Replace both plain placeholder and wrapped versions
      const plainPattern = new RegExp(
        block.placeholder.replace(/[.*+?^${}()|[\]\\]/g, "\\$&"),
        "g",
      );
      const strongPattern = new RegExp(
        `<strong>${block.placeholder.replace(/[.*+?^${}()|[\]\\]/g, "\\$&")}</strong>`,
        "g",
      );
      const emPattern = new RegExp(
        `<em>${block.placeholder.replace(/[.*+?^${}()|[\]\\]/g, "\\$&")}</em>`,
        "g",
      );

      html = html.replace(strongPattern, errorHtml);
      html = html.replace(emPattern, errorHtml);
      html = html.replace(plainPattern, errorHtml);
    }
  });

  return html;
}

const openai = new OpenAI({
  apiKey: await env("OPENAI_API_KEY"),
});

let currentPanel: string = ``;
let lastResponse: string = ``;

while (true) {
  const content = await micro(
    {
      input: "",
      placeholder: "Prompt OpenAI",
      strict: false,
      onEscape: () => {
        exit();
      },
      shortcuts: [
        {
          name: "Copy Response",
          key: `${cmd}+g`,
          onPress: async () => {
            if (lastResponse) {
              await copy(lastResponse);
              await clipboard.writeText(lastResponse);
              await toast("Response copied to clipboard!");
            } else {
              await toast("No response to copy yet");
            }
          },
          bar: "right",
          visible: false,
        },
      ],
    },
    currentPanel,
  );

  setLoading(true);

  let model_name: string = "";
  let prompt: string = "";
  let isOSeries: boolean = false;
  let new_content = "";

  if (content.startsWith("ai ")) {
    new_content = content.replace("ai ", "").trim();
  } else {
    new_content = content.trim();
  }

  if (content.startsWith("!o")) {
    model_name = "o4-mini";
    prompt = new_content.replace("!o", "").trim();
    isOSeries = true;
  } else {
    model_name = "gpt-4.1-mini";
    prompt = new_content;
  }

  const requestParams: any = {
    model: model_name,
    tools: [{ type: "web_search_preview" }],
    input: prompt || "Prompt missing?",
    instructions:
      "You are a helpful assistant with web access, today is " + Date,
  };

  // Add reasoning parameters for o-series models
  if (isOSeries) {
    requestParams.reasoning = {
      effort: "medium",
      summary: "auto",
    };
  }

  const response = await openai.responses.create(requestParams);

  let result = "";
  const citations = [];
  let reasoningText = "";

  // Extract text, citations, and reasoning from the response.output array
  if (response?.output && Array.isArray(response.output)) {
    for (const item of response.output) {
      if (item.type === "message" && item.content) {
        for (const contentItem of item.content) {
          if (contentItem.type === "output_text") {
            result = contentItem.text || "";

            // Extract citations from annotations
            if (contentItem.annotations) {
              for (const annotation of contentItem.annotations) {
                if (annotation.type === "url_citation") {
                  citations.push({
                    url: annotation.url,
                    title: annotation.title,
                  });
                }
              }
            }
          }
        }
      } else if (
        item.type === "reasoning" &&
        isOSeries &&
        item.summary &&
        Array.isArray(item.summary)
      ) {
        // Extract reasoning summary for o-series models
        reasoningText = item.summary
          .filter((summaryItem) => summaryItem.type === "summary_text")
          .map((summaryItem) => summaryItem.text)
          .join("\n\n");
      }
    }
  }

  // Prepend reasoning summary if available
  if (reasoningText) {
    const reasoningSummary = reasoningText
      .split("\n")
      .map((line) => `> ${line}`)
      .join("\n");
    result = reasoningSummary + "\n\n" + result;
  }

  // Add citations as numbered list if any exist
  if (citations.length > 0) {
    result += "\n\n**Sources:**\n";
    citations.forEach((citation, index) => {
      result += `${index + 1}. [${citation.title}](${citation.url})\n`;
    });
  }

  lastResponse = "**Prompt:**\n" + prompt + "\n\n**Assistant:**\n" + result;
  currentPanel = mdWithMath(result);
  setLoading(false);
}