import type { SessionEntry } from "@earendil-works/pi-coding-agent";
import type { SearchInputMessage } from "./types.ts";

const ASSISTANT_CONTEXT_TOKEN_LIMIT = 1_000;
const APPROX_CHARS_PER_TOKEN = 4;

function messageText(message: Record<string, unknown>): string {
  const content = message.content;
  if (typeof content === "string") return content;
  if (!Array.isArray(content)) return "";
  return content
    .filter((item): item is { type: "text"; text: string } => {
      return typeof item === "object" && item !== null && (item as { type?: unknown }).type === "text" && typeof (item as { text?: unknown }).text === "string";
    })
    .map((item) => item.text)
    .join("\n");
}

function toSearchMessage(entry: SessionEntry): SearchInputMessage | undefined {
  if (entry.type !== "message") return undefined;
  const message = entry.message as unknown as Record<string, unknown>;
  const role = message.role;
  if (role !== "user" && role !== "assistant") return undefined;
  const text = messageText(message);
  if (!text) return undefined;
  return {
    type: "message",
    role,
    content: [role === "user" ? { type: "input_text", text } : { type: "output_text", text }],
  };
}

function retainThroughLatestTwoUserMessages(messages: SearchInputMessage[]): SearchInputMessage[] {
  let latestUser = -1;
  for (let index = messages.length - 1; index >= 0; index -= 1) {
    if (messages[index]?.role === "user") {
      latestUser = index;
      break;
    }
  }
  if (latestUser < 0) return [];
  const throughLatestUser = messages.slice(0, latestUser + 1);
  const userIndices = throughLatestUser
    .map((message, index) => (message.role === "user" ? index : -1))
    .filter((index) => index >= 0);
  const start = userIndices.at(-2) ?? userIndices.at(-1) ?? 0;
  return throughLatestUser.slice(start);
}

function truncateAssistantContext(messages: SearchInputMessage[]): SearchInputMessage[] {
  let remainingChars = ASSISTANT_CONTEXT_TOKEN_LIMIT * APPROX_CHARS_PER_TOKEN;
  const result: SearchInputMessage[] = [];
  for (const message of messages) {
    if (message.role !== "assistant") {
      result.push(message);
      continue;
    }
    if (remainingChars <= 0) continue;
    const text = message.content[0]?.text ?? "";
    const truncated = text.length <= remainingChars ? text : `${text.slice(0, Math.max(0, remainingChars - 1))}…`;
    remainingChars -= Math.min(text.length, remainingChars);
    if (truncated) {
      result.push({ type: "message", role: "assistant", content: [{ type: "output_text", text: truncated }] });
    }
  }
  return result;
}

/** Sends the previous user turn, up to ~1k assistant tokens, and the current user turn to raw search. */
export function buildRecentSearchInput(entries: SessionEntry[]): SearchInputMessage[] | undefined {
  const messages = entries.map(toSearchMessage).filter((message): message is SearchInputMessage => message !== undefined);
  const recent = truncateAssistantContext(retainThroughLatestTwoUserMessages(messages));
  return recent.length > 0 ? recent : undefined;
}