import { bucketFor, maxBatchForLimits } from './engine/shapes.js';

function adapterText(ctx) {
  const info = ctx?.adapterInfo ?? {};
  return `${info.vendor ?? ''} ${info.architecture ?? ''} ${info.device ?? ''} ${info.description ?? ''}`.toLowerCase();
}

export function deviceProfile(ctx, dtype) {
  const text = adapterText(ctx);
  let kind = 'balanced';
  let maxRows = 64;
  let workBudget = 64 * 32;

  if (/swiftshader|llvmpipe|software|basic render/.test(text)) {
    kind = 'software';
    maxRows = 8;
    workBudget = 8 * 32;
  } else if (/adreno|mali|powervr/.test(text)) {
    kind = 'mobile';
    maxRows = 24;
    workBudget = 24 * 32;
  } else if (/nvidia|geforce|quadro|rtx|radeon|\bamd\b/.test(text)) {
    kind = 'discrete';
    maxRows = 192;
    workBudget = 192 * 32;
  } else if (/intel|apple/.test(text)) {
    kind = 'integrated';
    maxRows = 96;
    workBudget = 96 * 32;
  }

  if (dtype === 'f32') {
    maxRows = Math.max(8, Math.floor(maxRows / 2));
    workBudget = Math.max(8 * 32, Math.floor(workBudget / 2));
  }
  return { kind, maxRows, workBudget };
}

function fits(rows, ctx, dtype, profile) {
  if (!rows.length) return true;
  const S = bucketFor(rows[rows.length - 1].srcIds.length);
  const dtypeBytes = dtype === 'f16' ? 2 : 4;
  const bindingLimit = maxBatchForLimits(ctx, S, dtypeBytes);
  const workLimit = Math.max(1, Math.floor(profile.workBudget / S));
  return rows.length <= Math.min(profile.maxRows, workLimit, bindingLimit);
}

// Cheap, deterministic device fitting: pick from a few conservative device
// classes, respect real WebGPU binding limits, then balance a small tail.
// There is no startup sweep and no persistent autotune state.
export function planDeviceBatches(rows, ctx, dtype) {
  const sorted = [...rows].sort((a, b) => a.srcIds.length - b.srcIds.length || a.li - b.li);
  const profile = deviceProfile(ctx, dtype);
  const batches = [];
  let current = [];

  for (const row of sorted) {
    const next = [...current, row];
    if (current.length && !fits(next, ctx, dtype, profile)) {
      batches.push(current);
      current = [row];
    } else {
      current = next;
    }
  }
  if (current.length) batches.push(current);

  // A greedy cap can leave an inefficient final micro-batch. Move the
  // longest rows from its predecessor until the pair is balanced or the
  // tail reaches its own device limit.
  for (let i = 1; i < batches.length; i += 1) {
    const left = batches[i - 1];
    const right = batches[i];
    while (left.length > right.length + 1) {
      const moved = left[left.length - 1];
      const nextLeft = left.slice(0, -1);
      const nextRight = [moved, ...right];
      if (!fits(nextLeft, ctx, dtype, profile) || !fits(nextRight, ctx, dtype, profile)) break;
      left.pop();
      right.unshift(moved);
    }
  }

  return { batches, profile };
}