import type { ContextEntry, ContextEntryPriority, ContextWindow } from "../core/types.js";

const PRIORITY_WEIGHT: Record<ContextEntryPriority, number> = {
  critical: 100,
  high: 50,
  normal: 10,
  low: 1,
};

/**
 * Step 2: Context Manager
 * Manages the agent's context window with sliding-window summarization,
 * priority-based pruning, and token budget enforcement.
 *
 * "Massive context utilization" — Fable 5's headline capability.
 */
export class ContextManager {
  private window: ContextWindow;

  constructor(budget: number = 128_000) {
    this.window = {
      entries: [],
      totalTokens: 0,
      budget,
    };
  }

  // ── Entry Management ──────────────────────────────────────

  /** Add an entry and enforce budget */
  add(entry: Omit<ContextEntry, "id" | "timestamp">): ContextEntry {
    const full: ContextEntry = {
      id: cryptoUid(),
      timestamp: new Date().toISOString(),
      ...entry,
    };

    this.window.entries.push(full);
    this.window.totalTokens += full.tokenCount;
    this.enforceBudget();
    return full;
  }

  /** Add a batch of entries at once */
  addBatch(entries: Array<Omit<ContextEntry, "id" | "timestamp">>): ContextEntry[] {
    return entries.map((e) => this.add(e));
  }

  // ── Budget Enforcement ────────────────────────────────────

  /**
   * When the budget is exceeded:
   * 1. Summarize the lowest-priority entries (oldest 30% by score)
   * 2. Drop entries that have already been summarized and are low-priority
   */
  private enforceBudget(): void {
    if (this.window.totalTokens <= this.window.budget) return;

    // Sort by priority weight asc, then by timestamp asc (oldest first)
    const sorted = [...this.window.entries].sort((a, b) => {
      const pa = PRIORITY_WEIGHT[a.priority];
      const pb = PRIORITY_WEIGHT[b.priority];
      if (pa !== pb) return pa - pb;
      return new Date(a.timestamp).getTime() - new Date(b.timestamp).getTime();
    });

    let overage = this.window.totalTokens - this.window.budget;
    const toRemove: Set<string> = new Set();

    for (const entry of sorted) {
      if (overage <= 0) break;
      if (entry.priority === "critical") continue; // never drop critical

      if (entry.summary) {
        // Already summarized — safe to drop
        toRemove.add(entry.id);
        this.window.totalTokens -= entry.tokenCount;
        overage -= entry.tokenCount;
      } else {
        // Summarize first, then mark for removal on next pass
        entry.summary = this.summarize(entry.content);
        // Reduce token count to just the summary
        const saved = entry.tokenCount - entry.summary.length;
        this.window.totalTokens -= saved;
        overage -= saved;
        entry.tokenCount = entry.summary.length;
      }
    }

    // Remove entries marked for removal
    if (toRemove.size > 0) {
      this.window.entries = this.window.entries.filter(
        (e) => !toRemove.has(e.id)
      );
    }
  }

  // ── Summarization ─────────────────────────────────────────

  /** Produce a condensed summary of content (character-count based) */
  private summarize(content: string): string {
    const words = content.split(/\s+/);
    if (words.length <= 50) return content;

    // Extract first 2 sentences + key terms
    const sentences = content.match(/[^.!?\n]+[.!?]/g) || [content];
    const keyTerms = this.extractKeyTerms(content, 10);

    let summary = sentences.slice(0, 2).join(" ").trim();
    if (keyTerms.length > 0) {
      summary += ` [Key: ${keyTerms.join(", ")}]`;
    }
    if (summary.length > 500) {
      summary = summary.slice(0, 497) + "...";
    }
    return summary;
  }

  /** Extract significant terms (capitalized words, frequent terms) */
  private extractKeyTerms(content: string, max: number): string[] {
    const words = content.split(/\s+/);
    const freq = new Map<string, number>();

    for (const w of words) {
      const clean = w.replace(/[^a-zA-Z0-9_-]/g, "");
      if (clean.length < 4) continue;
      // Prefer capitalized/technical terms
      if (/^[A-Z]/.test(clean) || /[_-]/.test(clean)) {
        freq.set(clean, (freq.get(clean) || 0) + 1);
      }
    }

    return [...freq.entries()]
      .sort((a, b) => b[1] - a[1])
      .slice(0, max)
      .map(([term]) => term);
  }

  // ── Queries ───────────────────────────────────────────────

  getWindow(): ContextWindow {
    return { ...this.window, entries: [...this.window.entries] };
  }

  getEntries(options?: {
    priority?: ContextEntryPriority;
    source?: string;
    limit?: number;
  }): ContextEntry[] {
    let filtered = this.window.entries;

    if (options?.priority) {
      filtered = filtered.filter((e) => e.priority === options.priority);
    }
    if (options?.source) {
      filtered = filtered.filter((e) => e.source === options.source);
    }
    if (options?.limit) {
      filtered = filtered.slice(0, options.limit);
    }

    return filtered;
  }

  /** Build a single text block from all entries for LLM context injection */
  render(): string {
    return this.window.entries
      .sort((a, b) => {
        const pa = PRIORITY_WEIGHT[a.priority];
        const pb = PRIORITY_WEIGHT[b.priority];
        if (pa !== pb) return pb - pa; // higher priority first
        return (
          new Date(a.timestamp).getTime() -
          new Date(b.timestamp).getTime()
        );
      })
      .map((e) => {
        const content = e.summary || e.content;
        return `[${e.source}] ${content}`;
      })
      .join("\n\n");
  }

  getUsage(): { used: number; budget: number; pct: number } {
    return {
      used: this.window.totalTokens,
      budget: this.window.budget,
      pct: Math.round((this.window.totalTokens / this.window.budget) * 100),
    };
  }

  reset(budget?: number): void {
    this.window = {
      entries: [],
      totalTokens: 0,
      budget: budget ?? this.window.budget,
    };
  }
}

import { randomUUID } from "node:crypto";

function cryptoUid(): string {
  return randomUUID();
}