import { createHash, randomUUID } from "node:crypto";
import { readFile, realpath, stat } from "node:fs/promises";
import { isAbsolute, relative, resolve, extname } from "node:path";
import type { Chat, FileHandle, LMStudioClient } from "@lmstudio/sdk";

export const IMAGE_LIMIT = 8 * 1024 * 1024;
export const hashBytes = (bytes: Buffer): string => createHash("sha256").update(bytes).digest("hex");
const cache = new Map<string, string>(); // Digests only; never chat candidates or image buffers.
export interface Budget { remaining: number; }

export async function fileDigest(file: FileHandle, budget: Budget): Promise<string | undefined> {
  if (!file.isImage()) return undefined;
  try {
    if (file.identifier.startsWith("base64:")) {
      const encoded = file.identifier.slice(7);
      if (encoded.length > Math.ceil(IMAGE_LIMIT / 3) * 4) return undefined;
      const bytes = Buffer.from(encoded, "base64");
      if (bytes.length > IMAGE_LIMIT || bytes.length > budget.remaining) return undefined;
      budget.remaining -= bytes.length;
      return hashBytes(bytes);
    }
    const path = await file.getFilePath();
    const info = await stat(path);
    const key = `${path}:${info.size}:${info.mtimeMs}:${info.ctimeMs}`;
    const saved = cache.get(key);
    if (saved) return saved;
    if (!info.isFile() || info.size > IMAGE_LIMIT || info.size > budget.remaining) return undefined;
    budget.remaining -= info.size;
    const bytes = await readFile(path);
    if (bytes.length > IMAGE_LIMIT) return undefined;
    const digest = hashBytes(bytes);
    if (cache.size >= 256) cache.delete(cache.keys().next().value!);
    cache.set(key, digest);
    return digest;
  } catch { return undefined; }
}

export interface ImageCandidate {
  version: 1;
  id: string;
  relativePath: string;
  sha256: string;
  rootHash: string;
  createdAt: number;
}

export async function candidateBytes(root: string, name: string): Promise<{ bytes: Buffer; extension: string }> {
  if (isAbsolute(name) || !name.trim()) throw new Error("Invalid relative image path.");
  const directory = await realpath(root);
  const target = await realpath(resolve(directory, name));
  const rel = relative(directory, target);
  if (!rel || isAbsolute(rel) || rel === ".." || rel.startsWith("../") || rel.startsWith("..\\")) {
    throw new Error("Image path escapes the working directory.");
  }
  const info = await stat(target);
  if (!info.isFile() || info.size > IMAGE_LIMIT) throw new Error("Candidate exceeds image size limit.");
  const bytes = await readFile(target);
  if (bytes.length > IMAGE_LIMIT) throw new Error("Candidate exceeds image size limit.");
  return { bytes, extension: extname(target).toLowerCase() };
}

export async function makeCandidate(root: string, name: string): Promise<ImageCandidate> {
  return (await readCandidate(root, name)).candidate;
}

export async function readCandidate(root: string, name: string) {
  const input = await candidateBytes(root, name);
  const candidate: ImageCandidate = { version: 1, id: randomUUID(), relativePath: name, sha256: hashBytes(input.bytes),
    rootHash: hashBytes(Buffer.from(await realpath(root))), createdAt: Date.now() };
  return { ...input, candidate };
}

// Only receipts paired with an Analyze tool call in THIS history qualify. No directory-global queue.
export function recentCandidates(history: Chat): ImageCandidate[] {
  const messages = history.getMessagesArray().slice(-64);
  const calls = new Set<string>();
  const groups: ImageCandidate[][] = [];
  for (const message of messages) {
    if (message.isAssistantMessage()) {
      for (const call of message.getToolCallRequests()) {
        const name = call.name.replace(/[^a-z]/gi, "").toLowerCase();
        if (call.id && (name.endsWith("analyzelocalimageadvanced") || name.endsWith("multianalyzelocalimagesadvanced"))) calls.add(call.id);
      }
    }
    if (message.getRole() !== "tool") continue;
    for (const entry of message.getToolCallResults()) {
      if (!entry.toolCallId || !calls.has(entry.toolCallId)) continue;
      try {
        const parsed = JSON.parse(entry.content);
        const values = Array.isArray(parsed._aiaImageCandidates) ? parsed._aiaImageCandidates.slice(0, 4) : [parsed._aiaImageCandidate];
        const group: ImageCandidate[] = [];
        for (const value of values) {
        if (value?.version !== 1 || typeof value.id !== "string" || typeof value.relativePath !== "string"
          || !/^[a-f0-9]{64}$/.test(value.sha256) || !/^[a-f0-9]{64}$/.test(value.rootHash)
          || !Number.isFinite(value.createdAt) || value.createdAt > Date.now() + 60_000
          || Date.now() - value.createdAt > 24 * 60 * 60 * 1000) continue;
          group.push(value);
        }
        if (group.length) groups.push(group);
      } catch { /* Other tool outputs are not attachment candidates. */ }
    }
  }
  return groups.reverse().flat().filter((item, i, all) => all.findIndex(other => other.sha256 === item.sha256) === i).slice(0, 8);
}

export async function historyDigests(history: Chat, client: LMStudioClient, budget: Budget): Promise<Set<string>> {
  const digests = new Set<string>();
  for (const message of history.getMessagesArray().slice(-64).reverse()) {
    if (!message.isUserMessage()) continue;
    for (const file of message.getFiles(client)) {
      if (!file.isImage()) continue;
      const stamped = /^aia-([a-f0-9]{64})\./.exec(file.name)?.[1];
      const digest = stamped ?? await fileDigest(file, budget);
      if (digest) digests.add(digest);
    }
  }
  return digests;
}

import { createHash, randomUUID } from "node:crypto";
import { readFile, realpath, stat } from "node:fs/promises";
import { isAbsolute, relative, resolve, extname } from "node:path";
import type { Chat, FileHandle, LMStudioClient } from "@lmstudio/sdk";

export const IMAGE_LIMIT = 8 * 1024 * 1024;
export const hashBytes = (bytes: Buffer): string => createHash("sha256").update(bytes).digest("hex");
const cache = new Map<string, string>(); // Digests only; never chat candidates or image buffers.
export interface Budget { remaining: number; }

export async function fileDigest(file: FileHandle, budget: Budget): Promise<string | undefined> {
  if (!file.isImage()) return undefined;
  try {
    if (file.identifier.startsWith("base64:")) {
      const encoded = file.identifier.slice(7);
      if (encoded.length > Math.ceil(IMAGE_LIMIT / 3) * 4) return undefined;
      const bytes = Buffer.from(encoded, "base64");
      if (bytes.length > IMAGE_LIMIT || bytes.length > budget.remaining) return undefined;
      budget.remaining -= bytes.length;
      return hashBytes(bytes);
    }
    const path = await file.getFilePath();
    const info = await stat(path);
    const key = `${path}:${info.size}:${info.mtimeMs}:${info.ctimeMs}`;
    const saved = cache.get(key);
    if (saved) return saved;
    if (!info.isFile() || info.size > IMAGE_LIMIT || info.size > budget.remaining) return undefined;
    budget.remaining -= info.size;
    const bytes = await readFile(path);
    if (bytes.length > IMAGE_LIMIT) return undefined;
    const digest = hashBytes(bytes);
    if (cache.size >= 256) cache.delete(cache.keys().next().value!);
    cache.set(key, digest);
    return digest;
  } catch { return undefined; }
}

export interface ImageCandidate {
  version: 1;
  id: string;
  relativePath: string;
  sha256: string;
  rootHash: string;
  createdAt: number;
}

export async function candidateBytes(root: string, name: string): Promise<{ bytes: Buffer; extension: string }> {
  if (isAbsolute(name) || !name.trim()) throw new Error("Invalid relative image path.");
  const directory = await realpath(root);
  const target = await realpath(resolve(directory, name));
  const rel = relative(directory, target);
  if (!rel || isAbsolute(rel) || rel === ".." || rel.startsWith("../") || rel.startsWith("..\\")) {
    throw new Error("Image path escapes the working directory.");
  }
  const info = await stat(target);
  if (!info.isFile() || info.size > IMAGE_LIMIT) throw new Error("Candidate exceeds image size limit.");
  const bytes = await readFile(target);
  if (bytes.length > IMAGE_LIMIT) throw new Error("Candidate exceeds image size limit.");
  return { bytes, extension: extname(target).toLowerCase() };
}

export async function makeCandidate(root: string, name: string): Promise<ImageCandidate> {
  return (await readCandidate(root, name)).candidate;
}

export async function readCandidate(root: string, name: string) {
  const input = await candidateBytes(root, name);
  const candidate: ImageCandidate = { version: 1, id: randomUUID(), relativePath: name, sha256: hashBytes(input.bytes),
    rootHash: hashBytes(Buffer.from(await realpath(root))), createdAt: Date.now() };
  return { ...input, candidate };
}

// Only receipts paired with an Analyze tool call in THIS history qualify. No directory-global queue.
export function recentCandidates(history: Chat): ImageCandidate[] {
  const messages = history.getMessagesArray().slice(-64);
  const calls = new Set<string>();
  const groups: ImageCandidate[][] = [];
  for (const message of messages) {
    if (message.isAssistantMessage()) {
      for (const call of message.getToolCallRequests()) {
        const name = call.name.replace(/[^a-z]/gi, "").toLowerCase();
        if (call.id && (name.endsWith("analyzelocalimageadvanced") || name.endsWith("multianalyzelocalimagesadvanced"))) calls.add(call.id);
      }
    }
    if (message.getRole() !== "tool") continue;
    for (const entry of message.getToolCallResults()) {
      if (!entry.toolCallId || !calls.has(entry.toolCallId)) continue;
      try {
        const parsed = JSON.parse(entry.content);
        const values = Array.isArray(parsed._aiaImageCandidates) ? parsed._aiaImageCandidates.slice(0, 4) : [parsed._aiaImageCandidate];
        const group: ImageCandidate[] = [];
        for (const value of values) {
        if (value?.version !== 1 || typeof value.id !== "string" || typeof value.relativePath !== "string"
          || !/^[a-f0-9]{64}$/.test(value.sha256) || !/^[a-f0-9]{64}$/.test(value.rootHash)
          || !Number.isFinite(value.createdAt) || value.createdAt > Date.now() + 60_000
          || Date.now() - value.createdAt > 24 * 60 * 60 * 1000) continue;
          group.push(value);
        }
        if (group.length) groups.push(group);
      } catch { /* Other tool outputs are not attachment candidates. */ }
    }
  }
  return groups.reverse().flat().filter((item, i, all) => all.findIndex(other => other.sha256 === item.sha256) === i).slice(0, 8);
}

export async function historyDigests(history: Chat, client: LMStudioClient, budget: Budget): Promise<Set<string>> {
  const digests = new Set<string>();
  for (const message of history.getMessagesArray().slice(-64).reverse()) {
    if (!message.isUserMessage()) continue;
    for (const file of message.getFiles(client)) {
      if (!file.isImage()) continue;
      const stamped = /^aia-([a-f0-9]{64})\./.exec(file.name)?.[1];
      const digest = stamped ?? await fileDigest(file, budget);
      if (digest) digests.add(digest);
    }
  }
  return digests;
}