import type { ToolSpec } from "../tools/protocol"
import { estimateTokens } from "../utils/tokens"

export interface ContextCategory {
  name: string
  tokens: number
}

export interface ContextAnalysis {
  categories: ContextCategory[]
  totalTokens: number
  maxTokens: number
  percentage: number
  model: string
}

const MODEL_CONTEXT_WINDOWS: Record<string, number> = {
  "gpt-4o": 128_000,
  "gpt-4o-mini": 128_000,
  "gpt-4-turbo": 128_000,
  "glm-5.1": 128_000,
  "glm-5-turbo": 128_000,
  "glm-4.7-flash": 128_000,
  "claude-sonnet": 200_000,
  "claude-opus": 200_000,
  "deepseek-r1": 128_000,
  "qwen3": 128_000,
}

const DEFAULT_CONTEXT_WINDOW = 128_000

export function getContextWindowForModel(model: string): number {
  const lower = model.toLowerCase()
  for (const [name, window] of Object.entries(MODEL_CONTEXT_WINDOWS)) {
    if (lower.includes(name)) return window
  }
  const kMatch = lower.match(/(\d+)k/)
  if (kMatch) return parseInt(kMatch[1], 10) * 1_000
  if (lower.includes("1m") || lower.includes("million")) return 1_000_000
  return DEFAULT_CONTEXT_WINDOW
}

export function analyzeContextUsage(options: {
  systemPrompt: string
  toolSpecs: ToolSpec[]
  messages: Array<{ role: string; content: unknown }>
  skillsContent?: string
  memoryContent?: string
  model: string
}): ContextAnalysis {
  const categories: ContextCategory[] = []

  const systemTokens = estimateTokens(options.systemPrompt)
  if (systemTokens > 0) {
    categories.push({ name: "System prompt", tokens: systemTokens })
  }

  let toolTokens = 0
  for (const spec of options.toolSpecs) {
    toolTokens += estimateTokens(spec.description)
    toolTokens += estimateTokens(JSON.stringify(spec.inputSchema))
  }
  if (toolTokens > 0) {
    categories.push({ name: "Tool definitions", tokens: toolTokens })
  }

  if (options.skillsContent) {
    const skillsTokens = estimateTokens(options.skillsContent)
    if (skillsTokens > 0) {
      categories.push({ name: "Skills", tokens: skillsTokens })
    }
  }

  if (options.memoryContent) {
    const memTokens = estimateTokens(options.memoryContent)
    if (memTokens > 0) {
      categories.push({ name: "Memory", tokens: memTokens })
    }
  }

  let messageTokens = 0
  for (const msg of options.messages) {
    const text = typeof msg.content === "string" ? msg.content : JSON.stringify(msg.content ?? "")
    messageTokens += estimateTokens(text)
  }
  if (messageTokens > 0) {
    categories.push({ name: "Messages", tokens: messageTokens })
  }

  const totalTokens = categories.reduce((sum, c) => sum + c.tokens, 0)
  const maxTokens = getContextWindowForModel(options.model)
  const percentage = maxTokens > 0 ? Math.round((totalTokens / maxTokens) * 1000) / 10 : 0

  const freeTokens = Math.max(0, maxTokens - totalTokens)
  categories.push({ name: "Free space", tokens: freeTokens })

  return { categories, totalTokens, maxTokens, percentage, model: options.model }
}

export function formatContextAnalysis(analysis: ContextAnalysis): string {
  const lines: string[] = [
    "## Context Usage",
    "",
    `**Model:** ${analysis.model}`,
    `**Tokens:** ${analysis.totalTokens.toLocaleString()} / ${analysis.maxTokens.toLocaleString()} (${analysis.percentage}%)`,
    "",
  ]

  const visible = analysis.categories.filter((c) => c.tokens > 0 && c.name !== "Free space")

  if (visible.length > 0) {
    lines.push("| Category | Tokens | Percentage |")
    lines.push("|----------|--------|------------|")
    for (const cat of visible) {
      const pct = analysis.maxTokens > 0 ? ((cat.tokens / analysis.maxTokens) * 100).toFixed(1) : "0.0"
      lines.push(`| ${cat.name} | ${cat.tokens.toLocaleString()} | ${pct}% |`)
    }
    const free = analysis.categories.find((c) => c.name === "Free space")
    if (free && free.tokens > 0) {
      const pct = analysis.maxTokens > 0 ? ((free.tokens / analysis.maxTokens) * 100).toFixed(1) : "0.0"
      lines.push(`| Free space | ${free.tokens.toLocaleString()} | ${pct}% |`)
    }
    lines.push("")
  }

  return lines.join("\n")
}