ALTER TABLE IF EXISTS gateway_tasks
  ADD COLUMN IF NOT EXISTS api_key_name text,
  ADD COLUMN IF NOT EXISTS api_key_prefix text,
  ADD COLUMN IF NOT EXISTS requested_model text,
  ADD COLUMN IF NOT EXISTS resolved_model text,
  ADD COLUMN IF NOT EXISTS request_id text,
  ADD COLUMN IF NOT EXISTS usage jsonb NOT NULL DEFAULT '{}'::jsonb,
  ADD COLUMN IF NOT EXISTS metrics jsonb NOT NULL DEFAULT '{}'::jsonb,
  ADD COLUMN IF NOT EXISTS billing_summary jsonb NOT NULL DEFAULT '{}'::jsonb,
  ADD COLUMN IF NOT EXISTS final_charge_amount numeric NOT NULL DEFAULT 0,
  ADD COLUMN IF NOT EXISTS response_started_at timestamptz,
  ADD COLUMN IF NOT EXISTS response_finished_at timestamptz,
  ADD COLUMN IF NOT EXISTS response_duration_ms bigint;

ALTER TABLE IF EXISTS gateway_task_attempts
  ADD COLUMN IF NOT EXISTS request_id text,
  ADD COLUMN IF NOT EXISTS usage jsonb NOT NULL DEFAULT '{}'::jsonb,
  ADD COLUMN IF NOT EXISTS metrics jsonb NOT NULL DEFAULT '{}'::jsonb,
  ADD COLUMN IF NOT EXISTS response_started_at timestamptz,
  ADD COLUMN IF NOT EXISTS response_finished_at timestamptz,
  ADD COLUMN IF NOT EXISTS response_duration_ms bigint;

CREATE INDEX IF NOT EXISTS idx_gateway_tasks_request_id
  ON gateway_tasks(request_id)
  WHERE request_id IS NOT NULL;

CREATE INDEX IF NOT EXISTS idx_gateway_tasks_api_key_created
  ON gateway_tasks(api_key_id, created_at DESC)
  WHERE api_key_id IS NOT NULL;

UPDATE gateway_tasks
SET requested_model = model
WHERE requested_model IS NULL;

UPDATE gateway_tasks AS t
SET api_key_name = COALESCE(t.api_key_name, k.name),
    api_key_prefix = COALESCE(t.api_key_prefix, k.key_prefix)
FROM gateway_api_keys AS k
WHERE t.api_key_id = k.id::text;

WITH billing_totals AS (
  SELECT
    t.id,
    COUNT(line.value) AS line_count,
    COALESCE(SUM(
      CASE
        WHEN jsonb_typeof(line.value) = 'object'
          AND line.value ? 'amount'
          AND line.value->>'amount' ~ '^-?[0-9]+(\.[0-9]+)?$'
        THEN (line.value->>'amount')::numeric
        ELSE 0
      END
    ), 0) AS total_amount
  FROM gateway_tasks AS t
  LEFT JOIN LATERAL jsonb_array_elements(COALESCE(t.billings, '[]'::jsonb)) AS line(value) ON true
  GROUP BY t.id
)
UPDATE gateway_tasks AS t
SET final_charge_amount = billing_totals.total_amount,
    billing_summary = jsonb_build_object(
      'lineCount', billing_totals.line_count,
      'totalAmount', billing_totals.total_amount,
      'currency', 'resource',
      'finalCharge', jsonb_build_object(
        'amount', billing_totals.total_amount,
        'currency', 'resource',
        'simulated', COALESCE(t.run_mode = 'simulation', false)
      )
    )
FROM billing_totals
WHERE t.id = billing_totals.id
  AND COALESCE(t.billing_summary, '{}'::jsonb) = '{}'::jsonb;
