/**
 * pi-tokens v1.1 — Token cost intelligence for Pi
 *
 * Not just totals — WHERE the money goes. Per-tool cost breakdown
 * with $/call estimates. Find the expensive tool calls.
 *
 * /tokens              — session summary with cost estimate
 * /tokens tools        — per-tool cost breakdown
 * /tokens expensive    — top 10 most expensive individual calls
 * /tokens messages     — per-message breakdown
 * /tokens budget <$N>  — set session budget, warn when approaching
 * /tokens reset        — reset counters
 */

import type { ExtensionAPI } from '@anthropic-ai/claude-code'

// Cost per million tokens (rough averages)
const MODEL_COSTS: Record<string, { input: number; output: number }> = {
  'default':    { input: 3, output: 15 },
  'haiku':      { input: 0.25, output: 1.25 },
  'sonnet':     { input: 3, output: 15 },
  'opus':       { input: 15, output: 75 },
  'gpt-4o':     { input: 2.5, output: 10 },
  'gpt-5':      { input: 5, output: 20 },
  'gemini':     { input: 0.15, output: 0.60 },
}

interface ToolCall {
  toolName: string
  inputTokens: number
  outputTokens: number
  cost: number
  timestamp: number
}

interface TokenState {
  totalInput: number
  totalOutput: number
  messageCount: number
  sessionStart: number
  toolAgg: Map<string, { calls: number; tokens: number; cost: number }>
  expensiveCalls: ToolCall[]  // top 20 most expensive individual calls
  budget: number | null       // session budget in dollars
  budgetWarned: boolean
  model: string
}

const state: TokenState = {
  totalInput: 0, totalOutput: 0, messageCount: 0, sessionStart: Date.now(),
  toolAgg: new Map(), expensiveCalls: [], budget: null, budgetWarned: false, model: 'default',
}

function getCostRate(): { input: number; output: number } {
  for (const [key, rate] of Object.entries(MODEL_COSTS)) {
    if (state.model.toLowerCase().includes(key)) return rate
  }
  return MODEL_COSTS.default
}

function estimateCost(inputTok: number, outputTok: number): number {
  const rate = getCostRate()
  return (inputTok * rate.input + outputTok * rate.output) / 1_000_000
}

function totalCost(): number { return estimateCost(state.totalInput, state.totalOutput) }

function fmtCost(c: number): string {
  if (c < 0.001) return '<$0.001'
  if (c < 0.01) return `$${c.toFixed(3)}`
  return `$${c.toFixed(2)}`
}

function fmtNum(n: number): string {
  if (n >= 1_000_000) return `${(n / 1_000_000).toFixed(1)}M`
  if (n >= 1_000) return `${(n / 1_000).toFixed(1)}K`
  return String(n)
}

function estimateToolTokens(text: string): number {
  return text ? Math.ceil(text.length / 4) : 0
}

function formatSummary(): string {
  const elapsed = Math.max(1, Math.round((Date.now() - state.sessionStart) / 60000))
  const total = state.totalInput + state.totalOutput
  const cost = totalCost()
  const rate = Math.round(total / elapsed)
  const costRate = cost / elapsed

  const lines = [
    `## Token Cost Intelligence`, '',
    `| Metric | Value |`, `|--------|-------|`,
    `| Input | ${fmtNum(state.totalInput)} |`,
    `| Output | ${fmtNum(state.totalOutput)} |`,
    `| **Total** | **${fmtNum(total)}** |`,
    `| **Cost** | **${fmtCost(cost)}** |`,
    `| Cost/min | ${fmtCost(costRate)} |`,
    `| Tok/min | ~${fmtNum(rate)} |`,
    `| Messages | ${state.messageCount} |`,
    `| Model | ${state.model} |`,
    `| Uptime | ${elapsed}m |`,
  ]

  if (state.budget !== null) {
    const pct = Math.round(cost / state.budget * 100)
    const remaining = state.budget - cost
    lines.push(`| Budget | ${fmtCost(state.budget)} (${pct}% used, ${fmtCost(remaining)} left) |`)
  }

  // Top 3 expensive tools
  if (state.toolAgg.size > 0) {
    lines.push('', '### Costliest tools')
    const sorted = Array.from(state.toolAgg.entries()).sort((a, b) => b[1].cost - a[1].cost).slice(0, 5)
    for (const [tool, agg] of sorted) {
      const avgCost = agg.calls > 0 ? agg.cost / agg.calls : 0
      lines.push(`- \`${tool}\`: ${fmtCost(agg.cost)} total (${agg.calls} calls, ${fmtCost(avgCost)}/call)`)
    }
  }

  return lines.join('\n')
}

function formatToolBreakdown(): string {
  if (state.toolAgg.size === 0) return 'No tool calls tracked yet.'
  const cost = totalCost()
  const rows = ['| Tool | Calls | Tokens | Cost | $/call | % |', '|------|-------|--------|------|--------|---|']
  const sorted = Array.from(state.toolAgg.entries()).sort((a, b) => b[1].cost - a[1].cost)
  for (const [tool, agg] of sorted) {
    const pct = cost > 0 ? Math.round(agg.cost / cost * 100) : 0
    const avgCost = agg.calls > 0 ? agg.cost / agg.calls : 0
    rows.push(`| \`${tool}\` | ${agg.calls} | ${fmtNum(agg.tokens)} | ${fmtCost(agg.cost)} | ${fmtCost(avgCost)} | ${pct}% |`)
  }
  return `## Tool Cost Breakdown\n\n${rows.join('\n')}`
}

function formatExpensive(): string {
  if (state.expensiveCalls.length === 0) return 'No expensive calls tracked yet.'
  const rows = ['| # | Tool | Tokens | Cost | When |', '|---|------|--------|------|------|']
  const sorted = [...state.expensiveCalls].sort((a, b) => b.cost - a.cost).slice(0, 10)
  for (let i = 0; i < sorted.length; i++) {
    const c = sorted[i]
    const ago = Math.round((Date.now() - c.timestamp) / 60000)
    const agoStr = ago < 1 ? 'just now' : `${ago}m ago`
    rows.push(`| ${i + 1} | \`${c.toolName}\` | ${fmtNum(c.inputTokens + c.outputTokens)} | ${fmtCost(c.cost)} | ${agoStr} |`)
  }
  return `## Most Expensive Calls\n\n${rows.join('\n')}`
}

export default function init(pi: ExtensionAPI) {
  pi.on('message_end', (event: any) => {
    const msg = event.message as any
    if (msg?.usage) {
      state.totalInput += msg.usage.input_tokens || 0
      state.totalOutput += msg.usage.output_tokens || 0
    }
    state.messageCount++

    // Budget check
    if (state.budget !== null && !state.budgetWarned) {
      const cost = totalCost()
      if (cost >= state.budget * 0.9) {
        state.budgetWarned = true
        pi.sendMessage({ content: `⚠️ **Budget warning:** ${fmtCost(cost)} of ${fmtCost(state.budget)} (${Math.round(cost/state.budget*100)}%)`, display: true }, { triggerTurn: false })
      }
    }
  })

  pi.on('model_select', (event: any) => {
    state.model = event.model?.id || event.model?.name || 'default'
  })

  pi.on('tool_result', (event: any) => {
    const toolName = event.toolName || 'unknown'
    const content = Array.isArray(event.content) ? event.content.map((c: any) => c.text || '').join('') : ''
    const outputTok = estimateToolTokens(content)
    const inputTok = estimateToolTokens(JSON.stringify(event.input || ''))
    const cost = estimateCost(inputTok, outputTok)

    // Aggregate
    const agg = state.toolAgg.get(toolName) || { calls: 0, tokens: 0, cost: 0 }
    agg.calls++; agg.tokens += inputTok + outputTok; agg.cost += cost
    state.toolAgg.set(toolName, agg)

    // Track expensive calls
    state.expensiveCalls.push({ toolName, inputTokens: inputTok, outputTokens: outputTok, cost, timestamp: Date.now() })
    if (state.expensiveCalls.length > 100) {
      state.expensiveCalls.sort((a, b) => b.cost - a.cost)
      state.expensiveCalls.length = 50
    }
  })

  pi.addCommand({
    name: 'tokens',
    description: 'Token cost intelligence — where the money goes',
    handler: async (args) => {
      const parts = args.trim().split(/\s+/)
      const sub = parts[0]?.toLowerCase()

      if (sub === 'tools') { pi.sendMessage({ content: formatToolBreakdown(), display: true }, { triggerTurn: false }); return }
      if (sub === 'expensive') { pi.sendMessage({ content: formatExpensive(), display: true }, { triggerTurn: false }); return }
      if (sub === 'budget') {
        const amount = parseFloat(parts[1])
        if (isNaN(amount)) { pi.sendMessage({ content: 'Usage: /tokens budget <dollars>', display: true }, { triggerTurn: false }); return }
        state.budget = amount; state.budgetWarned = false
        pi.sendMessage({ content: `Budget set to ${fmtCost(amount)}. Warning at 90%.`, display: true }, { triggerTurn: false }); return
      }
      if (sub === 'reset') {
        state.totalInput = 0; state.totalOutput = 0; state.messageCount = 0
        state.toolAgg.clear(); state.expensiveCalls.length = 0
        state.sessionStart = Date.now(); state.budgetWarned = false
        pi.sendMessage({ content: 'Counters reset.', display: true }, { triggerTurn: false }); return
      }

      pi.sendMessage({ content: formatSummary(), display: true }, { triggerTurn: false })
    },
  })

  pi.addTool({ name: 'tokens_usage', description: 'Token cost intelligence with per-tool breakdown.',
    parameters: { type: 'object', properties: {} }, handler: async () => formatSummary() })

  pi.addTool({ name: 'tokens_expensive', description: 'Show the 10 most expensive individual tool calls.',
    parameters: { type: 'object', properties: {} }, handler: async () => formatExpensive() })
}