// ✅ Context Window Scheduler — deterministic token budget allocation
// file: src/gateway/ContextScheduler.ts

import { TokenAwareRequestBuilder } from './TokenAwareRequestBuilder'
import type { ChatMessage } from '../types'

const MODEL_CONTEXT_LIMITS: Record<string, number> = {
  'claude-3-5-sonnet-20241022': 200_000,
  'gpt-4o': 128_000,
  'gpt-4o-mini': 128_000,
}

export interface ContextBudget {
  systemTokens: number
  ragTokens: number
  historyTokens: number
  outputTokens: number
  totalAvailable: number
}

export interface PackedContext {
  messages: ChatMessage[]
  budget: ContextBudget
  pruned: number  // Number of history turns pruned
  ragChunksIncluded: number
}

export class ContextScheduler {
  private tokenizer = new TokenAwareRequestBuilder()

  pack(
    systemPrompt: string,
    ragChunks: { content: string; tokenCount: number }[],
    history: ChatMessage[],
    model: string,
    maxOutputTokens: number,
  ): PackedContext {
    const limit = MODEL_CONTEXT_LIMITS[model] ?? 128_000
    const safetyMargin = 200  // Buffer against tokeniser approximation errors

    // ── Budget allocation ───────────────────────────────────────────────────
    const systemTokens = this.tokenizer.countTokens(systemPrompt) + 4
    const outputReserve = maxOutputTokens
    const availableForContent = limit - systemTokens - outputReserve - safetyMargin

    // Allocate 55% to RAG, 45% to history.
    // 55/45 is calibrated for knowledge-retrieval tasks where context quality
    // outweighs history continuity. For conversational applications with minimal
    // RAG, adjust to 40/60. These values should be A/B tested against your eval
    // dataset (Part 22) using Context Recall as the guide metric.
    const ragBudget = Math.floor(availableForContent * 0.55)
    const historyBudget = availableForContent - ragBudget

    // ── RAG packing — greedy, respects budget ───────────────────────────────
    const includedChunks: string[] = []
    let ragTokensConsumed = 0

    for (const chunk of ragChunks) {
      if (ragTokensConsumed + chunk.tokenCount > ragBudget) break
      includedChunks.push(chunk.content)
      ragTokensConsumed += chunk.tokenCount
    }

    // ── History pruning — newest turns first ────────────────────────────────
    const prunedHistory: ChatMessage[] = []
    let historyTokensConsumed = 0
    let pruned = 0

    for (let i = history.length - 1; i >= 0; i--) {
      const cost = this.tokenizer.countTokens(history[i].content) + 4
      if (historyTokensConsumed + cost > historyBudget) {
        pruned++
        continue
      }
      prunedHistory.unshift(history[i])
      historyTokensConsumed += cost
    }

    // ── Attention-optimal prompt layout ─────────────────────────────────────
    // Critical: place RAG evidence LAST (before user turn) — highest attention position
    // Place system instructions FIRST — second-highest attention position
    const ragSection = includedChunks.length > 0
      ? `<knowledge_vault>\n${includedChunks.join('\n\n---\n\n')}\n</knowledge_vault>`
      : ''

    const messages: ChatMessage[] = [
      { role: 'system', content: systemPrompt },
      ...prunedHistory,
      // RAG evidence as the final assistant/system turn before user message
      ...(ragSection ? [{ role: 'user' as const, content: ragSection }] : []),
    ]

    return {
      messages,
      budget: {
        systemTokens,
        ragTokens: ragTokensConsumed,
        historyTokens: historyTokensConsumed,
        outputTokens: outputReserve,
        totalAvailable: limit,
      },
      pruned,
      ragChunksIncluded: includedChunks.length,
    }
  }
}

// ✅ Conversation summariser — replaces oldest turns with summary
// file: src/gateway/HistorySummariser.ts

import { generateText } from 'ai'
import { createAnthropic } from '@ai-sdk/anthropic'

const anthropic = createAnthropic()

export async function summariseHistory(
  turns: ChatMessage[],
  maxSummaryTokens = 400,
): Promise<string> {
  const formattedHistory = turns
    .map(m => `${m.role.toUpperCase()}: ${m.content}`)
    .join('\n\n')

  const { text } = await generateText({
    model: anthropic('claude-3-haiku-20240307'),
    maxTokens: maxSummaryTokens,
    temperature: 0,
    prompt: `Summarise the following conversation turns into a compact factual summary.
Preserve: decisions made, facts established, user goals stated, technical constraints mentioned.
Omit: pleasantries, repeated information, intermediate reasoning.

CONVERSATION:
${formattedHistory}

SUMMARY (max ${maxSummaryTokens} tokens):`,
  })

  return `[Earlier conversation summary]: ${text}`
}

// ✅ Attention-optimal prompt assembler used in ContextScheduler.pack()
// Critical layout rule: [system] [history] [RAG] [user_message]
// NOT:                  [system] [RAG] [history] [user_message]

// The RAG section immediately before the user turn maximises retrieval attention.
// This is the single layout change with the largest measurable impact on RAG quality.