import { Address, Instruction } from '@solana/kit';
import { ReadonlyUint8Array } from '@solana/codecs-core/dist/types/readonly-uint8array';

/** Keys for output buckets */
export type BucketKey = string;

/** A matcher says: “for this program, if the instruction data starts with this prefix, drop it into bucket K” */
export interface IxMatcher<K extends BucketKey = BucketKey> {
  /** program id to match */
  program: Address;
  /** Discriminator/prefix to match at the start of `ix.data` (any length, incl. 1 byte) */
  prefix: Uint8Array;
  /** Target bucket name */
  bucket: K;
}

/** Convenience builder that accepts a single program or an iterable, and a number | number[] | Uint8Array for the prefix */
export function makeIxMatcher<K extends BucketKey>(
  program: Address,
  discriminator: number | readonly number[] | Uint8Array,
  bucket: K,
): IxMatcher<K> {
  const prefix = typeof discriminator === 'number' ? Uint8Array.of(discriminator) : new Uint8Array(discriminator);
  return { program, prefix, bucket };
}

/** Output: buckets + everything else that didn't match any matcher */
export type ExtractedIxs<K extends BucketKey> = {
  buckets: Record<K, Instruction[]>;
  otherIxs: Instruction[];
};

export function extractIxsByMatchers<K extends BucketKey>(
  instructions: Instruction[],
  matchers: readonly IxMatcher<K>[],
): ExtractedIxs<K> {
  // Precreate buckets for stable object shape & perf
  const uniqueBuckets = Array.from(new Set(matchers.map((m) => m.bucket)));

  // Pre-allocate buckets map with correct typing
  const buckets: Record<K, Instruction[]> = {} as Record<K, Instruction[]>;
  for (const k of uniqueBuckets) {
    buckets[k] = [];
  }

  // Index: program -> matchers
  const byProgram = new Map<Address, IxMatcher<K>[]>();
  for (const m of matchers) {
    const arr = byProgram.get(m.program) || [];
    arr.push(m);
    byProgram.set(m.program, arr);
  }

  const otherIxs: Instruction[] = [];

  for (const ix of instructions) {
    const candidates = byProgram.get(ix.programAddress);
    if (!candidates || candidates.length === 0 || !ix.data) {
      otherIxs.push(ix);
      continue;
    }

    let matched = false;
    for (const m of candidates) {
      const p = m.prefix;
      if (ix.data.length >= p.length && prefixEquals(ix.data, p)) {
        buckets[m.bucket].push(ix); // add to every matching bucket
        matched = true;
      }
    }

    if (!matched) {
      otherIxs.push(ix);
    }
  }

  return { buckets, otherIxs };
}

function prefixEquals(buf: ReadonlyUint8Array, prefix: Uint8Array): boolean {
  // early out for equal references
  if (buf === prefix) return true;
  const n = prefix.length;
  for (let i = 0; i < n; i++) {
    if (buf[i] !== prefix[i]) {
      return false;
    }
  }
  return true;
}