/**
 * Time-bucketed cost/usage rollups (#124).
 *
 * Aggregates events into hourly buckets keyed on (tenant, verified agent, model)
 * so coarse/historical analytics read precomputed sums instead of scanning + JSON-
 * parsing raw rows, and so per-agent cost survives a raw-event purge (the
 * contributing pricing_versions are retained for reconciliation).
 *
 * `applyRollupBatch` runs inside the ingest transaction but is best-effort —
 * a rollup failure must never break ingest; retention/backfill reconciles.
 */
import { sql } from 'drizzle-orm';
import type { AgentLensEvent } from '@agentkitai/agentlens-core';

/** Cost-bearing event types (carry costUsd + a pricing_version). */
export const COST_EVENT_TYPES = new Set(['llm_response', 'cost_tracked']);

/** Truncate an ISO timestamp to its UTC hour, matching the analytics strftime bucket. */
export function bucketStartHour(timestamp: string): string {
  const d = new Date(timestamp);
  if (Number.isNaN(d.getTime())) return timestamp;
  const p = (n: number): string => String(n).padStart(2, '0');
  return `${d.getUTCFullYear()}-${p(d.getUTCMonth() + 1)}-${p(d.getUTCDate())}T${p(d.getUTCHours())}:00:00Z`;
}

export interface RollupBucket {
  tenantId: string;
  verifiedAgentId: string;
  model: string;
  bucketStart: string;
  eventCount: number;
  toolCallCount: number;
  errorCount: number;
  llmCallCount: number;
  inputTokens: number;
  outputTokens: number;
  cacheReadTokens: number;
  cacheWriteTokens: number;
  costUsd: number;
  latencySumMs: number;
  latencyCount: number;
  pricingVersions: Set<string>;
}

function payloadOf(e: AgentLensEvent): Record<string, unknown> {
  return (e.payload ?? {}) as Record<string, unknown>;
}
function num(v: unknown): number {
  return typeof v === 'number' && Number.isFinite(v) ? v : 0;
}
function verifiedAgentIdOf(e: AgentLensEvent): string {
  const v = e.metadata?.verifiedAgentId;
  return typeof v === 'string' && v ? v : '';
}
function modelOf(e: AgentLensEvent): string {
  const m = payloadOf(e).model;
  return typeof m === 'string' ? m : '';
}
function isErrorEvent(e: AgentLensEvent): boolean {
  return e.severity === 'error' || e.severity === 'critical' || e.eventType === 'tool_error';
}
function latencyOf(e: AgentLensEvent): number | null {
  const p = payloadOf(e);
  if (e.eventType === 'llm_response' && typeof p.latencyMs === 'number') return p.latencyMs;
  if (e.eventType === 'tool_response' && typeof p.durationMs === 'number') return p.durationMs;
  return null;
}

/**
 * Aggregate a batch of events into per-(agent, model, hour) buckets. Pure —
 * `pricingVersion` is the active pricing fingerprint stamped on cost events.
 */
export function aggregateBatch(
  events: AgentLensEvent[],
  pricingVersion: string | null,
): Map<string, RollupBucket> {
  const buckets = new Map<string, RollupBucket>();
  for (const e of events) {
    const tenantId = e.tenantId ?? 'default';
    const verifiedAgentId = verifiedAgentIdOf(e);
    const model = modelOf(e);
    const bucketStart = bucketStartHour(e.timestamp);
    const key = `${tenantId} ${verifiedAgentId} ${model} ${bucketStart}`;
    let b = buckets.get(key);
    if (!b) {
      b = {
        tenantId,
        verifiedAgentId,
        model,
        bucketStart,
        eventCount: 0,
        toolCallCount: 0,
        errorCount: 0,
        llmCallCount: 0,
        inputTokens: 0,
        outputTokens: 0,
        cacheReadTokens: 0,
        cacheWriteTokens: 0,
        costUsd: 0,
        latencySumMs: 0,
        latencyCount: 0,
        pricingVersions: new Set<string>(),
      };
      buckets.set(key, b);
    }
    const p = payloadOf(e);
    const usage = (p.usage ?? p) as Record<string, unknown>;
    b.eventCount += 1;
    if (e.eventType === 'tool_call') b.toolCallCount += 1;
    if (isErrorEvent(e)) b.errorCount += 1;
    if (e.eventType === 'llm_call') b.llmCallCount += 1;
    b.inputTokens += num(usage.inputTokens);
    b.outputTokens += num(usage.outputTokens);
    b.cacheReadTokens += num(usage.cacheReadTokens);
    b.cacheWriteTokens += num(usage.cacheWriteTokens);
    b.costUsd += num(p.costUsd);
    const lat = latencyOf(e);
    if (lat != null) {
      b.latencySumMs += lat;
      b.latencyCount += 1;
    }
    if (pricingVersion && COST_EVENT_TYPES.has(e.eventType)) b.pricingVersions.add(pricingVersion);
  }
  return buckets;
}

// eslint-disable-next-line @typescript-eslint/no-explicit-any
type SqlRunner = { get: (q: any) => any; run: (q: any) => unknown };

export interface RollupRow {
  event_count: number;
  tool_call_count: number;
  error_count: number;
  llm_call_count: number;
  input_tokens: number;
  output_tokens: number;
  cache_read_tokens: number;
  cache_write_tokens: number;
  cost_usd: number;
  latency_sum_ms: number;
  latency_count: number;
  pricing_versions: string;
}

export interface MergedRollup {
  eventCount: number;
  toolCallCount: number;
  errorCount: number;
  llmCallCount: number;
  inputTokens: number;
  outputTokens: number;
  cacheReadTokens: number;
  cacheWriteTokens: number;
  costUsd: number;
  latencySumMs: number;
  latencyCount: number;
  pricingVersions: string; // JSON array string
}

/**
 * Merge an existing rollup row (or none) with a batch bucket → the new cumulative
 * row. Shared by the SQLite (sync) and Postgres (async) writers so they stay in
 * lockstep. `Number()` coerces existing values (node-postgres returns numeric as
 * strings); the pricing-version set is unioned, not overwritten.
 */
export function mergeRollup(existing: RollupRow | undefined, b: RollupBucket): MergedRollup {
  const versions = new Set<string>(b.pricingVersions);
  if (existing) {
    try {
      for (const v of JSON.parse(existing.pricing_versions) as string[]) versions.add(v);
    } catch {
      /* ignore malformed */
    }
  }
  return {
    eventCount: Number(existing?.event_count ?? 0) + b.eventCount,
    toolCallCount: Number(existing?.tool_call_count ?? 0) + b.toolCallCount,
    errorCount: Number(existing?.error_count ?? 0) + b.errorCount,
    llmCallCount: Number(existing?.llm_call_count ?? 0) + b.llmCallCount,
    inputTokens: Number(existing?.input_tokens ?? 0) + b.inputTokens,
    outputTokens: Number(existing?.output_tokens ?? 0) + b.outputTokens,
    cacheReadTokens: Number(existing?.cache_read_tokens ?? 0) + b.cacheReadTokens,
    cacheWriteTokens: Number(existing?.cache_write_tokens ?? 0) + b.cacheWriteTokens,
    costUsd: Number(existing?.cost_usd ?? 0) + b.costUsd,
    latencySumMs: Number(existing?.latency_sum_ms ?? 0) + b.latencySumMs,
    latencyCount: Number(existing?.latency_count ?? 0) + b.latencyCount,
    pricingVersions: JSON.stringify([...versions]),
  };
}

/**
 * Merge a batch's buckets into `cost_rollups` (read-modify-write per bucket, so
 * count sums and the pricing-version set union are both exact). Runs on the
 * given tx/db runner; idempotent w.r.t. re-running a batch only at the row level
 * (callers pass deduped events).
 */
export function applyRollupBatch(runner: SqlRunner, events: AgentLensEvent[], pricingVersion: string | null): void {
  const buckets = aggregateBatch(events, pricingVersion);
  if (buckets.size === 0) return;
  const now = new Date().toISOString();
  for (const b of buckets.values()) {
    const existing = runner.get(sql`
      SELECT event_count, tool_call_count, error_count, llm_call_count, input_tokens, output_tokens,
             cache_read_tokens, cache_write_tokens, cost_usd, latency_sum_ms, latency_count, pricing_versions
      FROM cost_rollups
      WHERE tenant_id = ${b.tenantId} AND verified_agent_id = ${b.verifiedAgentId}
        AND model = ${b.model} AND bucket_start = ${b.bucketStart} AND granularity = 'hour'
    `) as RollupRow | undefined;

    const row = mergeRollup(existing, b);

    runner.run(sql`
      INSERT OR REPLACE INTO cost_rollups
        (tenant_id, verified_agent_id, model, bucket_start, granularity, event_count, tool_call_count,
         error_count, llm_call_count, input_tokens, output_tokens, cache_read_tokens, cache_write_tokens,
         cost_usd, latency_sum_ms, latency_count, pricing_versions, updated_at)
      VALUES
        (${b.tenantId}, ${b.verifiedAgentId}, ${b.model}, ${b.bucketStart}, 'hour', ${row.eventCount}, ${row.toolCallCount},
         ${row.errorCount}, ${row.llmCallCount}, ${row.inputTokens}, ${row.outputTokens}, ${row.cacheReadTokens}, ${row.cacheWriteTokens},
         ${row.costUsd}, ${row.latencySumMs}, ${row.latencyCount}, ${row.pricingVersions}, ${now})
    `);
  }
}