import { Database, and, eq, inArray, isNotNull, sql } from "../src/drizzle/index.js"
import { BillingTable, BlackPlans, SubscriptionTable, UsageTable } from "../src/schema/billing.sql.js"

if (process.argv.length < 3) {
  console.error("Usage: bun black-stats.ts <plan>")
  process.exit(1)
}
const plan = process.argv[2] as (typeof BlackPlans)[number]
if (!BlackPlans.includes(plan)) {
  console.error("Usage: bun black-stats.ts <plan>")
  process.exit(1)
}
const cutoff = new Date(Date.UTC(2026, 1, 0, 23, 59, 59, 999))

// get workspaces
const workspaces = await Database.use((tx) =>
  tx
    .select({ workspaceID: BillingTable.workspaceID })
    .from(BillingTable)
    .where(
      and(isNotNull(BillingTable.subscriptionID), sql`JSON_UNQUOTE(JSON_EXTRACT(subscription, '$.plan')) = ${plan}`),
    ),
)
if (workspaces.length === 0) throw new Error(`No active Black ${plan} subscriptions found`)

const week = sql<number>`YEARWEEK(${UsageTable.timeCreated}, 3)`
const workspaceIDs = workspaces.map((row) => row.workspaceID)
// Get subscription spend
const spend = await Database.use((tx) =>
  tx
    .select({
      workspaceID: UsageTable.workspaceID,
      week,
      amount: sql<number>`COALESCE(SUM(${UsageTable.cost}), 0)`,
    })
    .from(UsageTable)
    .where(
      and(inArray(UsageTable.workspaceID, workspaceIDs), sql`JSON_UNQUOTE(JSON_EXTRACT(enrichment, '$.plan')) = 'sub'`),
    )
    .groupBy(UsageTable.workspaceID, week),
)

// Get pay per use spend
const ppu = await Database.use((tx) =>
  tx
    .select({
      workspaceID: UsageTable.workspaceID,
      week,
      amount: sql<number>`COALESCE(SUM(${UsageTable.cost}), 0)`,
    })
    .from(UsageTable)
    .where(
      and(
        inArray(UsageTable.workspaceID, workspaceIDs),
        sql`(${UsageTable.enrichment} IS NULL OR JSON_UNQUOTE(JSON_EXTRACT(enrichment, '$.plan')) != 'sub')`,
      ),
    )
    .groupBy(UsageTable.workspaceID, week),
)

const models = await Database.use((tx) =>
  tx
    .select({
      workspaceID: UsageTable.workspaceID,
      model: UsageTable.model,
      amount: sql<number>`COALESCE(SUM(${UsageTable.cost}), 0)`,
    })
    .from(UsageTable)
    .where(
      and(inArray(UsageTable.workspaceID, workspaceIDs), sql`JSON_UNQUOTE(JSON_EXTRACT(enrichment, '$.plan')) = 'sub'`),
    )
    .groupBy(UsageTable.workspaceID, UsageTable.model),
)

const tokens = await Database.use((tx) =>
  tx
    .select({
      workspaceID: UsageTable.workspaceID,
      week,
      input: sql<number>`COALESCE(SUM(${UsageTable.inputTokens}), 0)`,
      cacheRead: sql<number>`COALESCE(SUM(${UsageTable.cacheReadTokens}), 0)`,
      output: sql<number>`COALESCE(SUM(${UsageTable.outputTokens}), 0) + COALESCE(SUM(${UsageTable.reasoningTokens}), 0)`,
    })
    .from(UsageTable)
    .where(
      and(inArray(UsageTable.workspaceID, workspaceIDs), sql`JSON_UNQUOTE(JSON_EXTRACT(enrichment, '$.plan')) = 'sub'`),
    )
    .groupBy(UsageTable.workspaceID, week),
)

const allWeeks = [...spend, ...ppu].map((row) => row.week)
const weeks = [...new Set(allWeeks)].sort((a, b) => a - b)
const spendMap = new Map<string, Map<number, number>>()
const totals = new Map<string, number>()
const ppuMap = new Map<string, Map<number, number>>()
const ppuTotals = new Map<string, number>()
const modelMap = new Map<string, { model: string; amount: number }[]>()
const tokenMap = new Map<string, Map<number, { input: number; cacheRead: number; output: number }>>()

for (const row of spend) {
  const workspace = spendMap.get(row.workspaceID) ?? new Map<number, number>()
  const total = totals.get(row.workspaceID) ?? 0
  const amount = toNumber(row.amount)
  workspace.set(row.week, amount)
  totals.set(row.workspaceID, total + amount)
  spendMap.set(row.workspaceID, workspace)
}

for (const row of ppu) {
  const workspace = ppuMap.get(row.workspaceID) ?? new Map<number, number>()
  const total = ppuTotals.get(row.workspaceID) ?? 0
  const amount = toNumber(row.amount)
  workspace.set(row.week, amount)
  ppuTotals.set(row.workspaceID, total + amount)
  ppuMap.set(row.workspaceID, workspace)
}

for (const row of models) {
  const current = modelMap.get(row.workspaceID) ?? []
  current.push({ model: row.model, amount: toNumber(row.amount) })
  modelMap.set(row.workspaceID, current)
}

for (const row of tokens) {
  const workspace = tokenMap.get(row.workspaceID) ?? new Map()
  workspace.set(row.week, {
    input: toNumber(row.input),
    cacheRead: toNumber(row.cacheRead),
    output: toNumber(row.output),
  })
  tokenMap.set(row.workspaceID, workspace)
}

const users = await Database.use((tx) =>
  tx
    .select({
      workspaceID: SubscriptionTable.workspaceID,
      subscribed: SubscriptionTable.timeCreated,
      subscription: BillingTable.subscription,
    })
    .from(SubscriptionTable)
    .innerJoin(BillingTable, eq(SubscriptionTable.workspaceID, BillingTable.workspaceID))
    .where(
      and(inArray(SubscriptionTable.workspaceID, workspaceIDs), sql`${SubscriptionTable.timeCreated} <= ${cutoff}`),
    ),
)

const counts = new Map<string, number>()
for (const user of users) {
  const current = counts.get(user.workspaceID) ?? 0
  counts.set(user.workspaceID, current + 1)
}

const rows = users
  .map((user) => {
    const workspace = spendMap.get(user.workspaceID) ?? new Map<number, number>()
    const ppuWorkspace = ppuMap.get(user.workspaceID) ?? new Map<number, number>()
    const count = counts.get(user.workspaceID) ?? 1
    const amount = (totals.get(user.workspaceID) ?? 0) / count
    const ppuAmount = (ppuTotals.get(user.workspaceID) ?? 0) / count
    const monthStart = user.subscribed ? startOfMonth(user.subscribed) : null
    const modelRows = (modelMap.get(user.workspaceID) ?? []).sort((a, b) => b.amount - a.amount).slice(0, 3)
    const modelTotal = totals.get(user.workspaceID) ?? 0
    const modelCells = modelRows.map((row) => ({
      model: row.model,
      percent: modelTotal > 0 ? `${((row.amount / modelTotal) * 100).toFixed(1)}%` : "0.0%",
    }))
    const modelData = [0, 1, 2].map((index) => modelCells[index] ?? { model: "-", percent: "-" })
    const weekly = Object.fromEntries(
      weeks.map((item) => {
        const value = (workspace.get(item) ?? 0) / count
        const beforeMonth = monthStart ? isoWeekStart(item) < monthStart : false
        return [formatWeek(item), beforeMonth ? "-" : formatMicroCents(value)]
      }),
    )
    const ppuWeekly = Object.fromEntries(
      weeks.map((item) => {
        const value = (ppuWorkspace.get(item) ?? 0) / count
        const beforeMonth = monthStart ? isoWeekStart(item) < monthStart : false
        return [formatWeek(item), beforeMonth ? "-" : formatMicroCents(value)]
      }),
    )
    const tokenWorkspace = tokenMap.get(user.workspaceID) ?? new Map()
    const weeklyTokens = Object.fromEntries(
      weeks.map((item) => {
        const t = tokenWorkspace.get(item) ?? { input: 0, cacheRead: 0, output: 0 }
        const beforeMonth = monthStart ? isoWeekStart(item) < monthStart : false
        return [
          formatWeek(item),
          beforeMonth
            ? { input: "-", cacheRead: "-", output: "-" }
            : {
                input: Math.round(t.input / count),
                cacheRead: Math.round(t.cacheRead / count),
                output: Math.round(t.output / count),
              },
        ]
      }),
    )
    return {
      workspaceID: user.workspaceID,
      useBalance: user.subscription?.useBalance ?? false,
      subscribed: formatDate(user.subscribed),
      subscribedAt: user.subscribed?.getTime() ?? 0,
      amount,
      ppuAmount,
      models: modelData,
      weekly,
      ppuWeekly,
      weeklyTokens,
    }
  })
  .sort((a, b) => a.subscribedAt - b.subscribedAt)

console.log(`Black ${plan} subscribers: ${rows.length}`)
const header = [
  "workspaceID",
  "subscribed",
  "useCredit",
  "subTotal",
  "ppuTotal",
  "model1",
  "model1%",
  "model2",
  "model2%",
  "model3",
  "model3%",
  ...weeks.flatMap((item) => [
    formatWeek(item) + " sub",
    formatWeek(item) + " ppu",
    formatWeek(item) + " input",
    formatWeek(item) + " cache",
    formatWeek(item) + " output",
  ]),
]
const lines = [header.map(csvCell).join(",")]
for (const row of rows) {
  const model1 = row.models[0]
  const model2 = row.models[1]
  const model3 = row.models[2]
  const cells = [
    row.workspaceID,
    row.subscribed ?? "",
    row.useBalance ? "yes" : "no",
    formatMicroCents(row.amount),
    formatMicroCents(row.ppuAmount),
    model1.model,
    model1.percent,
    model2.model,
    model2.percent,
    model3.model,
    model3.percent,
    ...weeks.flatMap((item) => {
      const t = row.weeklyTokens[formatWeek(item)] ?? { input: "-", cacheRead: "-", output: "-" }
      return [
        row.weekly[formatWeek(item)] ?? "",
        row.ppuWeekly[formatWeek(item)] ?? "",
        String(t.input),
        String(t.cacheRead),
        String(t.output),
      ]
    }),
  ]
  lines.push(cells.map(csvCell).join(","))
}
const output = `${lines.join("\n")}\n`
const file = Bun.file(`black-stats-${plan}.csv`)
await file.write(output)
console.log(`Wrote ${lines.length - 1} rows to ${file.name}`)
const total = rows.reduce((sum, row) => sum + row.amount, 0)
const average = rows.length === 0 ? 0 : total / rows.length
console.log(`Average spending per user: ${formatMicroCents(average)}`)

function formatMicroCents(value: number) {
  return `$${(value / 100000000).toFixed(2)}`
}

function formatDate(value: Date | null | undefined) {
  if (!value) return null
  return value.toISOString().split("T")[0]
}

function formatWeek(value: number) {
  return formatDate(isoWeekStart(value)) ?? ""
}

function startOfMonth(value: Date) {
  return new Date(Date.UTC(value.getUTCFullYear(), value.getUTCMonth(), 1))
}

function isoWeekStart(value: number) {
  const year = Math.floor(value / 100)
  const weekNumber = value % 100
  const jan4 = new Date(Date.UTC(year, 0, 4))
  const day = jan4.getUTCDay() || 7
  const weekStart = new Date(Date.UTC(year, 0, 4 - (day - 1)))
  weekStart.setUTCDate(weekStart.getUTCDate() + (weekNumber - 1) * 7)
  return weekStart
}

function toNumber(value: unknown) {
  if (typeof value === "number") return value
  if (typeof value === "bigint") return Number(value)
  if (typeof value === "string") return Number(value)
  return 0
}

function csvCell(value: string | number) {
  const text = String(value)
  if (!/[",\n]/.test(text)) return text
  return `"${text.replace(/"/g, '""')}"`
}