import type { Message } from '../models/types.js';
import type { MemoryStore, PromptMemorySection } from './store.js';

export interface AdaptiveMemoryConfig {
  maxTokens: number;
  maxSnippetsPerSection?: number;
}

interface ScoredSnippet {
  text: string;
  score: number;
  overlapScore: number;
}

const STOPWORDS = new Set([
  'a', 'an', 'and', 'are', 'as', 'at', 'be', 'by', 'for', 'from', 'has', 'he', 'in', 'is', 'it', 'its',
  'of', 'on', 'or', 'that', 'the', 'to', 'was', 'were', 'will', 'with', 'you', 'your', 'we', 'this',
  'they', 'their', 'our', 'but', 'not', 'can', 'just', 'into', 'about', 'after', 'before', 'than', 'then',
]);

export function buildAdaptiveMemoryContext(opts: {
  store: MemoryStore;
  userMessage: string;
  recentMessages: Message[];
  config: AdaptiveMemoryConfig;
}): string {
  const sections = opts.store.getPromptSections();
  if (sections.length === 0) {
    return '';
  }

  const maxChars = Math.max(1, opts.config.maxTokens) * 4;
  const sectionBlocks = buildSectionBlocks({
    sections,
    userMessage: opts.userMessage,
    recentMessages: opts.recentMessages,
    maxSnippetsPerSection: opts.config.maxSnippetsPerSection ?? 6,
  });

  if (sectionBlocks.length === 0) {
    return '';
  }

  const blocksByPriority = sectionBlocks.sort((a, b) => b.score - a.score || a.index - b.index);
  const selected: typeof blocksByPriority = [];
  let usedChars = 0;

  for (const block of blocksByPriority) {
    const withSeparator = selected.length > 0 ? 2 : 0;
    const projectedChars = usedChars + withSeparator + block.content.length;
    if (projectedChars > maxChars) {
      continue;
    }
    selected.push(block);
    usedChars = projectedChars;
  }

  if (selected.length === 0) {
    const best = blocksByPriority[0];
    return best.content.slice(0, maxChars);
  }

  selected.sort((a, b) => a.index - b.index);
  return selected.map((block) => block.content).join('\n\n');
}

export function buildRecentMemoryContext(store: MemoryStore, maxTokens: number): string {
  const full = store.getContextForPrompt();
  if (!full) {
    return '';
  }

  const maxChars = Math.max(1, maxTokens) * 4;
  if (full.length <= maxChars) {
    return full;
  }
  return full.slice(-maxChars);
}

function buildSectionBlocks(opts: {
  sections: PromptMemorySection[];
  userMessage: string;
  recentMessages: Message[];
  maxSnippetsPerSection: number;
}): Array<{ score: number; content: string; index: number }> {
  const weightedContextTokens = buildWeightedContextTokens(opts.userMessage, opts.recentMessages);
  const blocks: Array<{ score: number; content: string; index: number }> = [];

  for (let sectionIndex = 0; sectionIndex < opts.sections.length; sectionIndex++) {
    const section = opts.sections[sectionIndex];
    const snippets = splitIntoSnippets(section.content).slice(-opts.maxSnippetsPerSection);
    if (snippets.length === 0) {
      continue;
    }

    const scored = scoreSnippets(snippets, weightedContextTokens);
    const bestSnippets = scored
      .filter((snippet) => snippet.overlapScore > 0)
      .sort((a, b) => b.score - a.score)
      .slice(0, Math.min(3, scored.length))
      .map((snippet) => snippet.text.trim())
      .filter((snippet) => snippet.length > 0);

    if (bestSnippets.length === 0) {
      continue;
    }

    const block = `## ${section.title}\n\n${bestSnippets.join('\n\n')}`;
    const blockScore = scored.reduce((sum, item) => sum + item.score, 0) / scored.length;
    blocks.push({ score: blockScore, content: block, index: sectionIndex });
  }

  return blocks;
}

function splitIntoSnippets(content: string): string[] {
  return content
    .split(/\n\s*\n/g)
    .map(part => part.trim())
    .filter(part => part.length > 0);
}

function scoreSnippets(snippets: string[], weightedContextTokens: Map<string, number>): ScoredSnippet[] {
  const totalContextWeight = Math.max(
    1,
    Array.from(weightedContextTokens.values()).reduce((sum, value) => sum + value, 0),
  );

  return snippets.map((snippet, index) => {
    const snippetTokens = new Set(tokenize(snippet));
    let overlapWeight = 0;

    for (const token of snippetTokens) {
      overlapWeight += weightedContextTokens.get(token) ?? 0;
    }

    const overlapScore = overlapWeight / totalContextWeight;
    const recencyScore = (index + 1) / snippets.length;
    const score = overlapScore * 0.8 + recencyScore * 0.2;

    return { text: snippet, score, overlapScore };
  });
}

function buildWeightedContextTokens(userMessage: string, recentMessages: Message[]): Map<string, number> {
  const weighted = new Map<string, number>();
  const recent = recentMessages.slice(-6);

  addTokens(weighted, userMessage, 1.0);

  for (let i = 0; i < recent.length; i++) {
    const msg = recent[i];
    const recencyWeight = 0.25 + ((i + 1) / recent.length) * 0.55;
    addTokens(weighted, typeof msg.content === 'string' ? msg.content : '', recencyWeight);
  }

  return weighted;
}

function addTokens(target: Map<string, number>, text: string, weight: number): void {
  for (const token of tokenize(text)) {
    target.set(token, (target.get(token) ?? 0) + weight);
  }
}

function tokenize(text: string): string[] {
  return text
    .toLowerCase()
    .split(/[^a-z0-9]+/)
    .filter(token => token.length >= 3 && !STOPWORDS.has(token));
}
