#Requires -Version 5.1
<#
.SYNOPSIS
  Cosmos DB — Cost / RU-consumption report.

.DESCRIPTION
  Identifies which database (product) and which container drives the most
  Azure Cosmos DB cost. For SERVERLESS accounts cost is dominated by Request
  Units (RU) consumed; for PROVISIONED accounts it is the provisioned
  throughput. The billing mode is auto-detected.

  Reports:
    1. Account billing mode (serverless vs provisioned) + region
    2. Actual billed cost (Cost Management, last 30d) by service
    3. RU consumption by database, ranked, with est. monthly $
    4. RU consumption by container for the top databases
    5. Storage (DataUsage) by database

.PARAMETER Account
  Cosmos account name. Falls back to $env:COSMOS_ACCOUNT_NAME.

.PARAMETER ResourceGroup
  Resource group. Falls back to $env:COSMOS_RESOURCE_GROUP.

.PARAMETER Days
  Lookback window (days) for RU metrics. Default 7 (or $env:DAYS).

.PARAMETER Top
  Rows per table. Default 15 (or $env:TOP).

.PARAMETER Drill
  Number of top databases to drill into by container. Default 3.

.PARAMETER Rate
  USD per 1M RU (serverless). Default 0.25 (or $env:SERVERLESS_RU_RATE).

.EXAMPLE
  $env:COSMOS_ACCOUNT_NAME='cosmos-mywisprai'; $env:COSMOS_RESOURCE_GROUP='rg-mywisprai'
  ./scripts/cosmos-cost-report.ps1

.EXAMPLE
  ./scripts/cosmos-cost-report.ps1 -Account cosmos-mywisprai -ResourceGroup rg-mywisprai -Days 7

.NOTES
  Prerequisites: Azure CLI installed and authenticated (az login).
#>
[CmdletBinding()]
param(
  [string]$Account       = $env:COSMOS_ACCOUNT_NAME,
  [string]$ResourceGroup = $env:COSMOS_RESOURCE_GROUP,
  [int]   $Days          = $(if ($env:DAYS) { [int]$env:DAYS } else { 7 }),
  [int]   $Top           = $(if ($env:TOP)  { [int]$env:TOP }  else { 15 }),
  [int]   $Drill         = 3,
  [double]$Rate          = $(if ($env:SERVERLESS_RU_RATE) { [double]$env:SERVERLESS_RU_RATE } else { 0.25 })
)

$ErrorActionPreference = 'Stop'

if (-not $Account)       { Write-Error 'Set COSMOS_ACCOUNT_NAME or pass -Account'; exit 2 }
if (-not $ResourceGroup) { Write-Error 'Set COSMOS_RESOURCE_GROUP or pass -ResourceGroup'; exit 2 }
if (-not (Get-Command az -ErrorAction SilentlyContinue)) { Write-Error 'az CLI not found'; exit 2 }

$start = (Get-Date).ToUniversalTime().AddDays(-$Days).ToString('yyyy-MM-ddTHH:mm:ssZ')
$end   = (Get-Date).ToUniversalTime().ToString('yyyy-MM-ddTHH:mm:ssZ')

Write-Host '════════════════════════════════════════════════════════════════'
Write-Host ' Cosmos DB Cost Report'
Write-Host "   Account: $Account   Resource group: $ResourceGroup"
Write-Host "   RU window: last ${Days}d   ($start -> $end)"
Write-Host '════════════════════════════════════════════════════════════════'

$rid    = az cosmosdb show -n $Account -g $ResourceGroup --query id -o tsv
$caps   = az cosmosdb show -n $Account -g $ResourceGroup --query "capabilities[].name" -o tsv
$region = az cosmosdb show -n $Account -g $ResourceGroup --query "locations[0].locationName" -o tsv
$mode   = if ($caps -match 'EnableServerless') { 'SERVERLESS' } else { 'PROVISIONED' }

Write-Host ''
Write-Host "▶ Billing mode : $mode"
Write-Host "▶ Region       : $region"
Write-Host ("▶ Est. RU rate : `$$Rate per 1M RU (serverless)")

# ── 1. Actual billed cost (Cost Management, last 30d) ─────────────
Write-Host ''
Write-Host '── Actual billed cost — last 30d by service (resource group) ──'
$sub    = az account show --query id -o tsv
$cmFrom = (Get-Date).ToUniversalTime().AddDays(-30).ToString('yyyy-MM-ddT00:00:00Z')
$cmTo   = (Get-Date).ToUniversalTime().ToString('yyyy-MM-ddT00:00:00Z')
$cmBody = @{
  type       = 'ActualCost'
  timeframe  = 'Custom'
  timePeriod = @{ from = $cmFrom; to = $cmTo }
  dataset    = @{
    granularity = 'None'
    aggregation = @{ totalCost = @{ name = 'Cost'; function = 'Sum' } }
    grouping    = @(@{ type = 'Dimension'; name = 'ServiceName' })
  }
} | ConvertTo-Json -Depth 10 -Compress
$tmp = New-TemporaryFile
Set-Content -Path $tmp -Value $cmBody -NoNewline
try {
  $rows = az rest --method post `
    --url "https://management.azure.com/subscriptions/$sub/resourceGroups/$ResourceGroup/providers/Microsoft.CostManagement/query?api-version=2023-11-01" `
    --body "@$tmp" --headers "Content-Type=application/json" `
    --query "properties.rows" -o json 2>$null | ConvertFrom-Json
  if ($rows) {
    $rows | Sort-Object { [double]$_[0] } -Descending | ForEach-Object {
      '   {0,-28} ${1,8:N2} {2}' -f ([string]$_[1]).Substring(0,[Math]::Min(28,([string]$_[1]).Length)), [double]$_[0], $_[2]
    }
  } else {
    Write-Host '   (no cost data — needs Cost Management reader on the subscription)'
  }
} catch {
  Write-Host '   (cost query unavailable — continuing with RU metrics)'
} finally {
  Remove-Item $tmp -ErrorAction SilentlyContinue
}

# Helper: query TotalRequestUnits split by a dimension, return [pscustomobject]@{ k; ru }
function Get-RuByDimension {
  param([string]$Filter)
  $json = az monitor metrics list --resource $rid --metric TotalRequestUnits `
    --aggregation Total --interval P1D --start-time $start --end-time $end `
    --filter $Filter --top 500 `
    --query "value[0].timeseries[].{k: metadatavalues[0].value, ru: sum(data[].total)}" `
    -o json 2>$null
  if (-not $json) { return @() }
  $parsed = $json | ConvertFrom-Json
  if (-not $parsed) { return @() }
  ,@($parsed | ForEach-Object {
    [pscustomobject]@{ k = $(if ($_.k) { $_.k } else { '<none>' }); ru = [double]($_.ru) }
  })
}

function Show-RuTable {
  param([object[]]$Rows)
  $ranked = @($Rows | Sort-Object ru -Descending)
  $total  = ($ranked | Measure-Object ru -Sum).Sum
  if (-not $total -or $total -eq 0) { $total = 1.0 }
  '   {0,-28} {1,16} {2,8}  {3,10}' -f 'name', "RU (${Days}d)", 'share', 'est $/mo' | Write-Host
  '   ' + ('-' * 68) | Write-Host
  foreach ($r in ($ranked | Select-Object -First $Top)) {
    $est   = $r.ru / $Days * 30.0 / 1000000.0 * $Rate
    $share = 100.0 * $r.ru / $total
    $name  = ([string]$r.k).Substring(0,[Math]::Min(28,([string]$r.k).Length))
    '   {0,-28} {1,16} {2,7:N1}%  {3,9:N2}' -f $name, ('{0:N0}' -f $r.ru), $share, $est | Write-Host
  }
  $proj = $total / $Days * 30.0 / 1000000.0 * $Rate
  '   ' + ('-' * 68) | Write-Host
  '   {0,-28} {1,16} {2,8}  {3,9:N2}' -f 'TOTAL', ('{0:N0}' -f $total), '', $proj | Write-Host
}

# ── 2. RU by database ─────────────────────────────────────────────
Write-Host ''
Write-Host "── RU consumption by database (product) — last ${Days}d ──"
$dbRows = Get-RuByDimension -Filter "DatabaseName eq '*'"
Show-RuTable -Rows $dbRows

# ── 3. Drill into the top databases by container ──────────────────
$topDbs = @($dbRows | Where-Object { $_.ru -gt 0 -and $_.k -ne '<none>' } |
            Sort-Object ru -Descending | Select-Object -First $Drill -ExpandProperty k)
foreach ($db in $topDbs) {
  Write-Host ''
  Write-Host "── RU by container in '$db' — last ${Days}d ──"
  $coll = Get-RuByDimension -Filter "DatabaseName eq '$db' and CollectionName eq '*'"
  Show-RuTable -Rows $coll
}

# ── 4. Storage by database ────────────────────────────────────────
Write-Host ''
Write-Host '── Storage (DataUsage) by database — latest snapshot ──'
$storeJson = az monitor metrics list --resource $rid --metric DataUsage `
  --aggregation Total --interval PT1H --start-time $start --end-time $end `
  --filter "DatabaseName eq '*'" --top 200 `
  --query "value[0].timeseries[].{k: metadatavalues[0].value, b: max(data[].total)}" `
  -o json 2>$null
$store = if ($storeJson) { @($storeJson | ConvertFrom-Json) } else { @() }
if ($store.Count -eq 0) {
  Write-Host '   (no storage data)'
} else {
  foreach ($s in ($store | Sort-Object { [double]$_.b } -Descending)) {
    $name = ([string]$s.k).Substring(0,[Math]::Min(28,([string]$s.k).Length))
    '   {0,-28} {1,10:N2} MB' -f $name, ([double]$s.b / 1MB) | Write-Host
  }
}

Write-Host ''
Write-Host '════════════════════════════════════════════════════════════════'
Write-Host ' Notes:'
Write-Host "  - Serverless cost ~= RU consumed x `$$Rate/1M + storage(`$~0.25/GB-mo)."
Write-Host "  - 'est `$/mo' linearly projects the ${Days}d window to 30 days."
Write-Host '  - High RU + low request count => expensive per-op (cross-partition'
Write-Host '    queries / large docs) — prime rightsizing target.'
Write-Host '  - A *_locks container burning RU is usually lock-polling overhead.'
Write-Host '════════════════════════════════════════════════════════════════'