// Practical tokenisation — using Tiktoken
// GPT-4o uses o200k_base (200k vocab); GPT-4 / GPT-3.5-turbo use cl100k_base (100k vocab)
import { encoding_for_model } from 'tiktoken'

// Note: encoding_for_model('gpt-4o') resolves to o200k_base
const enc = encoding_for_model('gpt-4o')

function tokenise(text: string): { ids: number[]; count: number } {
  const ids = enc.encode(text)
  return { ids: Array.from(ids), count: ids.length }
}

// Cost and latency examples:
console.log(tokenise('Hello, world!').count)         // → 4
console.log(tokenise('Héllo, wörld!').count)         // → 8  (non-ASCII ≈ 2× tokens)
console.log(tokenise('    '.repeat(10)).count)       // → 10 (each leading space = token)
console.log(tokenise('userId').count)                // → 1  (common camelCase merge)
console.log(tokenise('x_forwarded_for').count)       // → 5  (underscored tokens fragment)

// ✅ Token-aware request builder — baseline for the Workspace gateway
// file: src/gateway/TokenAwareRequestBuilder.ts

// Use o200k_base for GPT-4o; cl100k_base for GPT-4/GPT-3.5. Using gpt-4o as a
// close approximation for all modern models (all share the same vocabulary size range).
import { encoding_for_model } from 'tiktoken'
import type { ChatMessage } from '../types'

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

export class TokenAwareRequestBuilder {
  private enc = encoding_for_model('gpt-4o') // o200k_base (200k vocab) — close approximation for all modern models

  countTokens(text: string): number {
    return this.enc.encode(text).length
  }

  countMessagesTokens(messages: ChatMessage[]): number {
    // Each message has 4 tokens overhead: role, name, content separator, delimiters
    return messages.reduce(
      (total, msg) => total + this.countTokens(msg.content) + 4,
      3, // Priming tokens for every request
    )
  }

  buildWithBudget(
    systemPrompt: string,
    history: ChatMessage[],
    model: string,
    maxOutputTokens: number,
  ): ChatMessage[] {
    const limit = MODEL_TOKEN_LIMITS[model] ?? 128_000
    const systemTokens = this.countTokens(systemPrompt) + 4
    const availableForHistory = limit - systemTokens - maxOutputTokens - 100 // safety margin

    // Prune from oldest turn, preserving the last user message
    const pruned: ChatMessage[] = []
    let consumed = 0

    for (let i = history.length - 1; i >= 0; i--) {
      const cost = this.countTokens(history[i].content) + 4
      if (consumed + cost > availableForHistory) break
      pruned.unshift(history[i])
      consumed += cost
    }

    return [{ role: 'system', content: systemPrompt }, ...pruned]
  }
}

// ✅ Correct Pattern — Streaming metrics instrumented at the gateway
// file: src/gateway/MeteredStream.ts

export interface StreamMetrics {
  ttft: number      // ms from request sent to first token received
  itl: number[]     // ms between each successive token
  tps: number       // tokens per second (average decode phase)
  totalTokens: number
  totalMs: number
}

export async function* meteredStream(
  stream: AsyncIterable<string>,
  requestSentAt: number,
): AsyncGenerator<string, StreamMetrics> {
  let firstToken = true
  let ttft = 0
  const itl: number[] = []
  let prevTokenAt = 0
  let totalTokens = 0

  for await (const token of stream) {
    const now = performance.now()

    if (firstToken) {
      ttft = now - requestSentAt
      prevTokenAt = now
      firstToken = false
    } else {
      itl.push(now - prevTokenAt)
      prevTokenAt = now
    }

    totalTokens++
    yield token
  }

  const totalMs = performance.now() - requestSentAt
  const decodeMs = totalMs - ttft
  const tps = totalTokens > 1 ? ((totalTokens - 1) / decodeMs) * 1000 : 0

  return { ttft, itl, tps, totalTokens, totalMs }
}

// ✅ Prefix-cache-stable system prompt construction
// file: src/gateway/SystemPromptBuilder.ts

export class SystemPromptBuilder {
  // STABLE SECTION — Must never change between requests for the same org.
  // Placed first to maximise prefix cache coverage.
  private buildStablePrefix(orgName: string, tools: string): string {
    return [
      `You are the Enterprise Collaborative AI Workspace assistant for ${orgName}.`,
      'You help engineers draft technical documents, create architecture diagrams, and retrieve knowledge.',
      '',
      '## Capabilities',
      tools,  // Tool schemas — must be deterministically serialised
      '',
      '## Response Rules',
      '- Always respond in structured JSON matching the WorkspaceAction schema.',
      '- Never fabricate document content. Ground all claims in retrieved knowledge.',
    ].join('\n')
  }

  // DYNAMIC SECTION — Placed last; not cached.
  private buildDynamicSuffix(userContext: { name: string; timezone: string }): string {
    return `\n## Current Session\nUser: ${userContext.name}\nTimezone: ${userContext.timezone}`
  }

  build(orgName: string, tools: string, userContext: { name: string; timezone: string }): string {
    // Stable prefix first → maximises KV cache hit rate
    // Dynamic suffix last → only un-cached portion re-computed each call
    return this.buildStablePrefix(orgName, tools) + this.buildDynamicSuffix(userContext)
  }
}

VRAM for KV cache ≈ 2 × num_layers × num_heads × head_dim × context_length × bytes_per_element

For a 70B parameter model (typical cloud hosted):
  2 × 80 × 8 × 128 × 100,000 × 2 bytes (fp16) ≈ 32 GB for 100k token context

// ✅ Type-safe model config for different task profiles
// file: src/gateway/ModelConfigs.ts

export type TaskProfile = 'structured_output' | 'document_draft' | 'agent_plan' | 'embedding'

export const MODEL_CONFIGS: Record<TaskProfile, ModelConfig> = {
  structured_output: {
    model: 'claude-3-5-sonnet-20241022',
    temperature: 0,
    max_tokens: 2048,
    // Use JSON Schema constrained decoding when provider supports it
  },
  document_draft: {
    model: 'claude-3-5-sonnet-20241022',
    temperature: 0.4,
    max_tokens: 8192,
  },
  agent_plan: {
    model: 'claude-3-5-sonnet-20241022',
    temperature: 0.1,      // Near-deterministic planning
    max_tokens: 4096,
  },
  embedding: {
    model: 'text-embedding-3-large',
    dimensions: 1536,      // Matryoshka: reducible to 256 for fast ANN
  },
} as const