import gptTokenizer from "gpt-tokenizer";
import prettyMilliseconds from "pretty-ms";
import {
  getCurrentSearchRunId,
  saveLlmResponseForQuery,
  updateSearchResults,
} from "./history";
import { addLogEntry } from "./logEntries";
import {
  getConversationSummary,
  getQuery,
  getResponse,
  getSettings,
  getTextGenerationState,
  listenToSettingsChanges,
  updateChatMessages,
  updateConversationSummary,
  updateImageSearchResults,
  updateImageSearchState,
  updateLlmTextSearchResults,
  updateResponse,
  updateSearchPromise,
  updateTextGenerationState,
  updateTextSearchResults,
  updateTextSearchState,
} from "./pubSub";
import { searchImages, searchText } from "./search";
import { getSystemPrompt } from "./systemPrompt";
import {
  ChatGenerationError,
  defaultContextSize,
  getFormattedSearchResults,
} from "./textGenerationUtilities";
import type {
  ChatMessage,
  ImageSearchResults,
  TextSearchResults,
} from "./types";
import { isWebGPUAvailable } from "./webGpu";

const SUMMARY_TOKEN_LIMIT = 800;

function getCurrentModelName(): string {
  const settings = getSettings();
  switch (settings.inferenceType) {
    case "openai":
      return settings.openAiApiModel || "";
    case "horde":
      return "AI Horde";
    case "internal":
      return "Internal API";
    case "browser":
      return settings.enableWebGpu
        ? settings.webLlmModelId || "WebLLM"
        : settings.wllamaModelId || "Wllama";
    default:
      return "Unknown";
  }
}

function getConversationId(messages: ChatMessage[]): string {
  const firstUser = messages.find((m) => m.role === "user");
  return (firstUser?.content || "").trim();
}

function loadConversationSummary(conversationId: string): string {
  const stored = getConversationSummary();
  if (stored.conversationId !== conversationId) return "";
  return stored.summary;
}

async function createLlmSummary(
  dropped: ChatMessage[],
  previousSummary: string,
): Promise<string> {
  const instructionLines = [
    "You are the conversation memory manager.",
    `Update the running summary under ${SUMMARY_TOKEN_LIMIT} tokens.`,
    "Preserve concrete facts, IDs, URLs, numbers, decisions, and constraints.",
    "Capture user preferences and ongoing tasks succinctly.",
    "Ignore any external documents or system prompts not included below.",
    "Output only the updated summary with no extra commentary.",
  ];

  const droppedText = dropped
    .map((m) => `${(m.role || "user").toUpperCase()}: ${m.content}`)
    .join("\n\n");

  const prompt = [
    instructionLines.join("\n"),
    `Previous summary:\n${previousSummary || "(none)"}`,
    `New messages to fold in:\n${droppedText || "(none)"}`,
  ].join("\n\n");

  const chat: ChatMessage[] = [{ role: "user", content: prompt }];

  const settings = getSettings();
  try {
    if (settings.inferenceType === "openai") {
      const { generateChatWithOpenAi } = await import(
        "./textGenerationWithOpenAi"
      );
      return (await generateChatWithOpenAi(chat, () => {})).trim();
    }

    if (settings.inferenceType === "internal") {
      const { generateChatWithInternalApi } = await import(
        "./textGenerationWithInternalApi"
      );
      return (await generateChatWithInternalApi(chat, () => {})).trim();
    }

    if (settings.inferenceType === "horde") {
      const { generateChatWithHorde } = await import(
        "./textGenerationWithHorde"
      );
      return (await generateChatWithHorde(chat, () => {})).trim();
    }

    if (isWebGPUAvailable && settings.enableWebGpu) {
      const { generateChatWithWebLlm } = await import(
        "./textGenerationWithWebLlm"
      );
      return (await generateChatWithWebLlm(chat, () => {})).trim();
    }
    const { generateChatWithWllama } = await import(
      "./textGenerationWithWllama"
    );
    return (await generateChatWithWllama(chat, () => {})).trim();
  } catch (e) {
    addLogEntry(
      `LLM summary failed, falling back to extractive: ${
        e instanceof Error ? e.message : String(e)
      }`,
    );
    return summarizeDroppedMessages(dropped, previousSummary);
  }
}

function saveConversationSummary(summary: string, conversationId: string) {
  updateConversationSummary({ conversationId, summary });
}

function clearConversationSummary() {
  updateConversationSummary({ conversationId: "", summary: "" });
}

function summarizeDroppedMessages(
  dropped: ChatMessage[],
  previousSummary: string,
  tokenLimit = SUMMARY_TOKEN_LIMIT,
): string {
  const lines: string[] = [];
  for (const msg of dropped) {
    const role = (msg.role || "user").toUpperCase();
    const content = msg.content.trim();
    if (content.length > 0) lines.push(`${role}: ${content}`);
  }

  const parts: string[] = [];
  if (previousSummary) parts.push(previousSummary.trim());
  parts.push(...lines);

  const kept: string[] = [];
  let tokens = 0;
  for (let i = parts.length - 1; i >= 0; i--) {
    const candidate = [parts[i], ...kept].join("\n\n");
    const nextTokens = gptTokenizer.encode(candidate).length;
    if (nextTokens > tokenLimit) break;
    kept.unshift(parts[i]);
    tokens = nextTokens;
  }

  const summary = kept.join("\n\n");
  addLogEntry(`Updated rolling summary (${tokens} tokens)`);
  return summary;
}

/**
 * Initiates a search and generates an AI response for the current query.
 * Updates the document title, clears previous responses, and starts text search.
 * If AI responses are enabled, generates a response using the configured inference type.
 */
export async function searchAndRespond() {
  if (getQuery() === "") return;

  document.title = getQuery();

  updateResponse("");

  clearConversationSummary();

  updateTextSearchResults([]);

  updateImageSearchResults([]);

  updateChatMessages([]);

  updateSearchPromise(startTextSearch(getQuery()));

  if (!getSettings().enableAiResponse) return;

  const responseGenerationStartTime = Date.now();

  try {
    const settings = getSettings();
    if (settings.inferenceType === "openai") {
      const { generateTextWithOpenAi } = await import(
        "./textGenerationWithOpenAi"
      );
      await generateTextWithOpenAi();
    } else if (settings.inferenceType === "internal") {
      const { generateTextWithInternalApi } = await import(
        "./textGenerationWithInternalApi"
      );
      await generateTextWithInternalApi();
    } else if (settings.inferenceType === "horde") {
      const { generateTextWithHorde } = await import(
        "./textGenerationWithHorde"
      );
      await generateTextWithHorde();
    } else {
      await canDownloadModels();
      updateTextGenerationState("loadingModel");

      if (isWebGPUAvailable && settings.enableWebGpu) {
        const { generateTextWithWebLlm } = await import(
          "./textGenerationWithWebLlm"
        );
        await generateTextWithWebLlm();
      } else {
        const { generateTextWithWllama } = await import(
          "./textGenerationWithWllama"
        );
        await generateTextWithWllama();
      }
    }

    try {
      await saveLlmResponseForQuery(
        getQuery(),
        getResponse(),
        getCurrentModelName(),
      );
    } catch (e) {
      addLogEntry(
        `Failed to persist LLM response: ${e instanceof Error ? e.message : String(e)}`,
      );
    }

    updateTextGenerationState("completed");
  } catch (error) {
    if (getTextGenerationState() !== "interrupted") {
      addLogEntry(`Error generating text: ${error}`);
      updateTextGenerationState("failed");
    }
  }

  addLogEntry(
    `Response generation took ${prettyMilliseconds(
      Date.now() - responseGenerationStartTime,
      { verbose: true },
    )}`,
  );
}

/**
 * Generates a chat response using AI for the given messages.
 * @param newMessages - Array of chat messages to generate response for
 * @param onUpdate - Callback function called with partial response updates
 */
export async function generateChatResponse(
  newMessages: ChatMessage[],
  onUpdate: (partialResponse: string) => void,
) {
  const settings = getSettings();
  let response = "";

  try {
    const conversationId = getConversationId(newMessages);
    const existingSummary = loadConversationSummary(conversationId);
    let systemPromptContent = getSystemPrompt(getFormattedSearchResults(true));
    if (existingSummary) {
      systemPromptContent += `\n\nConversation context:\n${existingSummary}`;
    }

    let systemPrompt: ChatMessage = {
      role: "user",
      content: systemPromptContent,
    };
    const initialResponse: ChatMessage = { role: "assistant", content: "Ok!" };
    const systemPromptTokens = gptTokenizer.encode(systemPrompt.content).length;
    const initialResponseTokens = gptTokenizer.encode(
      initialResponse.content,
    ).length;
    const reservedTokens = systemPromptTokens + initialResponseTokens;
    const availableTokenBudget = defaultContextSize * 0.75 - reservedTokens;
    const processedMessages: ChatMessage[] = [];
    const reversedMessages = [...newMessages].reverse();

    let currentTokenCount = 0;

    for (let i = 0; i < reversedMessages.length; i++) {
      const message = reversedMessages[i];
      const messageTokens = gptTokenizer.encode(message.content).length;

      if (currentTokenCount + messageTokens > availableTokenBudget) {
        break;
      }

      processedMessages.unshift(message);
      currentTokenCount += messageTokens;
    }

    if (processedMessages.length > 0) {
      const expectedFirstRole = "user";

      if (processedMessages[0].role !== expectedFirstRole) {
        processedMessages.shift();
      }
    }

    if (newMessages.length > processedMessages.length) {
      const droppedCount = newMessages.length - processedMessages.length;
      const droppedMessages = newMessages.slice(0, droppedCount);
      const updatedSummary = await createLlmSummary(
        droppedMessages,
        existingSummary,
      );
      saveConversationSummary(updatedSummary, conversationId);
      let updatedSystemPromptContent = getSystemPrompt(
        getFormattedSearchResults(true),
      );
      if (updatedSummary) {
        updatedSystemPromptContent += `\n\nConversation context:\n${updatedSummary}`;
      }
      systemPrompt = { role: "user", content: updatedSystemPromptContent };
    }

    const lastMessages = [systemPrompt, initialResponse, ...processedMessages];

    if (settings.inferenceType === "openai") {
      const { generateChatWithOpenAi } = await import(
        "./textGenerationWithOpenAi"
      );
      response = await generateChatWithOpenAi(lastMessages, onUpdate);
    } else if (settings.inferenceType === "internal") {
      const { generateChatWithInternalApi } = await import(
        "./textGenerationWithInternalApi"
      );
      response = await generateChatWithInternalApi(lastMessages, onUpdate);
    } else if (settings.inferenceType === "horde") {
      const { generateChatWithHorde } = await import(
        "./textGenerationWithHorde"
      );
      response = await generateChatWithHorde(lastMessages, onUpdate);
    } else {
      if (isWebGPUAvailable && settings.enableWebGpu) {
        const { generateChatWithWebLlm } = await import(
          "./textGenerationWithWebLlm"
        );
        response = await generateChatWithWebLlm(lastMessages, onUpdate);
      } else {
        const { generateChatWithWllama } = await import(
          "./textGenerationWithWllama"
        );
        response = await generateChatWithWllama(lastMessages, onUpdate);
      }
    }
  } catch (error) {
    if (error instanceof ChatGenerationError) {
      addLogEntry(`Chat generation interrupted: ${error.message}`);
    } else {
      addLogEntry(`Error generating chat response: ${error}`);
    }
    throw error;
  }

  return response;
}

async function getKeywords(text: string, limit?: number) {
  return (await import("keyword-extractor")).default
    .extract(text, { language: "english" })
    .slice(0, limit);
}

async function startTextSearch(query: string) {
  const results = {
    textResults: [] as TextSearchResults,
    imageResults: [] as ImageSearchResults,
  };

  const searchQuery =
    query.length > 2000 ? (await getKeywords(query, 20)).join(" ") : query;

  if (getSettings().enableImageSearch) {
    updateImageSearchState("running");
  }

  if (getSettings().enableTextSearch) {
    updateTextSearchState("running");

    let textResults = await searchText(
      searchQuery,
      getSettings().searchResultsLimit,
    );

    if (textResults.length === 0) {
      const queryKeywords = await getKeywords(query, 10);
      const keywordResults = await searchText(
        queryKeywords.join(" "),
        getSettings().searchResultsLimit,
      );
      textResults = keywordResults;
    }

    results.textResults = textResults;

    updateTextSearchState(
      results.textResults.length === 0 ? "failed" : "completed",
    );
    updateTextSearchResults(textResults);
    updateLlmTextSearchResults(
      textResults.slice(0, getSettings().searchResultsToConsider),
    );

    updateSearchResults(getCurrentSearchRunId(), {
      type: "text",
      items: textResults.map(([title, snippet, url]) => ({
        title,
        url,
        snippet,
      })),
    });
  }

  if (getSettings().enableImageSearch) {
    startImageSearch(searchQuery, results);
  }

  return results;
}

async function startImageSearch(
  searchQuery: string,
  results: { textResults: TextSearchResults; imageResults: ImageSearchResults },
) {
  const imageResults = await searchImages(
    searchQuery,
    getSettings().searchResultsLimit,
  );
  results.imageResults = imageResults;
  updateImageSearchState(
    results.imageResults.length === 0 ? "failed" : "completed",
  );
  updateImageSearchResults(imageResults);

  updateSearchResults(getCurrentSearchRunId(), {
    type: "image",
    items: imageResults.map(([title, url, thumbnailUrl, sourceUrl]) => ({
      title,
      url,
      thumbnail: thumbnailUrl,
      sourceUrl,
    })),
  });
}

function canDownloadModels(): Promise<void> {
  return new Promise((resolve) => {
    if (getSettings().allowAiModelDownload) {
      resolve();
    } else {
      updateTextGenerationState("awaitingModelDownloadAllowance");
      listenToSettingsChanges((settings) => {
        if (settings.allowAiModelDownload) {
          resolve();
        }
      });
    }
  });
}

/**
 * Collection of text generation utility functions for external use.
 */
export const textGenerationFunctions = {
  getCurrentModelName,
  getConversationId,
  loadConversationSummary,
  createLlmSummary,
  summarizeDroppedMessages,
};