import { join } from "../collector/index.js";
import { arrayLike, sequence } from "../sync/index.js";
import { emptyIterableIterator } from "../utils.js";
import { BitArray } from "./types.js";

const BYTE_SIZE = Uint8Array.BYTES_PER_ELEMENT * 8;
const FULL_BYTE = getMask(BYTE_SIZE);

function getMask(bitCount: number) {
  return ~(~0b0 << bitCount);
}

function getByteIndex(index: number) {
  return Math.floor(index / BYTE_SIZE);
}

function getBitIndex(index: number) {
  return index % BYTE_SIZE;
}

function setBit(byte: number, bitIndex: number, value: boolean) {
  const mask = 0b1 << bitIndex;
  return value ? (byte | mask) : (byte & ~mask);
}

function* yieldBits(byte: number, bitCount: number) {
  for (let i = 0; i < bitCount; i++) {
    yield Boolean((byte >>> i) & 0b1);
  }
}

function* getBytes(bits: Iterable<boolean>) {
  const bitIterator = bits[Symbol.iterator]();

  while (true) {
    let next = bitIterator.next();

    if (next.done) {
      break;
    }

    let byte = 0;

    for (let i = 0; i < BYTE_SIZE; i++) {
      byte |= Number(next.value) << i;

      next = bitIterator.next();

      if (next.done) {
        break;
      }
    }

    yield byte;
  }
}

export class BitArrayImpl implements BitArray {
  readonly #length: number;
  readonly #bits: Uint8Array;
  readonly #wholeBytes: number;
  readonly #remainingBits: number;
  readonly #remainingBitsMask: number;

  public constructor(length: number) {
    const exactBytes = length / BYTE_SIZE;
    const remainingBits = length % BYTE_SIZE;

    this.#length = length;
    this.#bits = new Uint8Array(Math.ceil(exactBytes));
    this.#wholeBytes = Math.floor(exactBytes);
    this.#remainingBits = remainingBits;
    this.#remainingBitsMask = getMask(remainingBits);
  }

  *[Symbol.iterator]() {
    for (let i = 0; i < this.#wholeBytes; i++) {
      yield* yieldBits(this.#bits[i], BYTE_SIZE);
    }

    if (this.#remainingBits > 0) {
      yield* yieldBits(this.#bits[this.#wholeBytes], this.#remainingBits);
    }
  }

  #ensureValidIndex(index: number) {
    if (index < 0 || index >= this.#length) {
      throw new RangeError("The index is outside the BitArray range.");
    }
  }

  #ensureSameSize(other: BitArray) {
    if (this.#length !== other.length) {
      throw new TypeError("The BitArrays do not have the same length.");
    }
  }

  public get length() {
    return this.#length;
  }

  public get(index: number) {
    this.#ensureValidIndex(index);

    return Boolean((this.#bits[getByteIndex(index)] >>> getBitIndex(index)) & 0b1);
  }

  public set(index: number, value: boolean) {
    this.#ensureValidIndex(index);

    const byteIndex = getByteIndex(index);
    this.#bits[byteIndex] = setBit(this.#bits[byteIndex], getBitIndex(index), value);
  }

  public isFull() {
    return this.#bits.subarray(0, this.#wholeBytes).every(byte => byte === FULL_BYTE) && (this.#remainingBits === 0 || this.#bits[this.#wholeBytes] === this.#remainingBitsMask);
  }

  public isEmpty() {
    return this.#bits.every(byte => byte === 0);
  }

  public fill(value: boolean) {
    const b = value ? FULL_BYTE : 0;
    for (let i = 0; i < this.#wholeBytes; i++) {
      this.#bits[i] = b;
    }

    if (this.#remainingBits > 0) {
      this.#bits[this.#wholeBytes] = value ? this.#remainingBitsMask : 0;
    }
  }

  public and(other: BitArray) {
    this.#ensureSameSize(other);

    if (other instanceof BitArrayImpl) {
      for (let i = 0; i < this.#length; i++) {
        this.#bits[i] &= other.#bits[i];
      }
    } else {
      let i = 0;
      for (const byte of getBytes(other)) {
        this.#bits[i++] &= byte;
      }
    }

    return this;
  }

  public or(other: BitArray) {
    this.#ensureSameSize(other);

    if (other instanceof BitArrayImpl) {
      for (let i = 0; i < this.#length; i++) {
        this.#bits[i] |= other.#bits[i];
      }
    } else {
      let i = 0;
      for (const byte of getBytes(other)) {
        this.#bits[i++] |= byte;
      }
    }

    return this;
  }

  public xor(other: BitArray) {
    this.#ensureSameSize(other);

    if (other instanceof BitArrayImpl) {
      for (let i = 0; i < this.#length; i++) {
        this.#bits[i] ^= other.#bits[i];
      }
    } else {
      let i = 0;
      for (const byte of getBytes(other)) {
        this.#bits[i++] ^= byte;
      }
    }

    return this;
  }

  public not() {
    for (let i = 0; i < this.#wholeBytes; i++) {
      this.#bits[i] = ~this.#bits[i] & FULL_BYTE;
    }

    if (this.#remainingBits > 0) {
      this.#bits[this.#wholeBytes] = ~this.#bits[this.#wholeBytes] & this.#remainingBitsMask;
    }

    return this;
  }

  public contains(other: BitArray) {
    this.#ensureSameSize(other);

    return other instanceof BitArrayImpl ?
      arrayLike(this.#bits).zip(arrayLike(other.#bits)).all(([a, b]) => (a & b) === b) :
      sequence(this).zip(sequence(other)).where(([, b]) => b).all(([a]) => a);
  }

  public intersects(other: BitArray) {
    this.#ensureSameSize(other);

    return other instanceof BitArrayImpl ?
      arrayLike(this.#bits).zip(arrayLike(other.#bits)).any(([a, b]) => (a & b) !== 0) :
      sequence(this).zip(sequence(other)).any(([a, b]) => a && b);
  }

  public slice(offset: number, length: number) {
    return new BitArraySlice(this, offset, length);
  }

  public copy() {
    const copy = new BitArrayImpl(this.#length);
    copy.#bits.set(this.#bits);
    return copy;
  }

  public toArray() {
    return sequence(this).toArray();
  }

  public equals(other: BitArray) {
    return other === this || other && (other instanceof BitArrayImpl ? arrayLike(this.#bits).sequenceEquals(arrayLike(other.#bits)) : sequence(this).sequenceEquals(sequence(other)));
  }

  public toString() {
    return sequence(this).select(bit => bit ? '1' : '0').collect(join());
  }
}

export class EmptyBitArray implements BitArray {
  [Symbol.iterator](): Iterator<boolean, any, undefined> {
    return emptyIterableIterator<boolean>();
  }

  get length() {
    return 0;
  }

  isFull() {
    return false;
  }

  isEmpty() {
    return false;
  }

  get(index: number): boolean {
    throw new Error("BitArray has zero length");
  }

  set(index: number, value: boolean) {
    throw new Error("BitArray has zero length");
  }

  fill(value: boolean) { }

  and(other: BitArray) {
    if (other.length !== 0) {
      throw new Error("The other BitArray does not have zero length");
    }

    return this;
  }

  or(other: BitArray) {
    if (other.length !== 0) {
      throw new Error("The other BitArray does not have zero length");
    }

    return this;
  }

  xor(other: BitArray) {
    if (other.length !== 0) {
      throw new Error("The other BitArray does not have zero length");
    }

    return this;
  }

  not() {
    return this;
  }

  contains(other: BitArray) {
    return false;
  }

  intersects(other: BitArray) {
    return false;
  }

  slice(offset: number, length: number) {
    if (offset > 0) {
      throw new Error("offset > 0");
    }

    if (length > 0) {
      throw new Error("length > 0");
    }

    return this;
  }

  copy() {
    return this;
  }

  toArray(): boolean[] {
    return [];
  }

  equals(other: BitArray) {
    return other === this || other && other.length === 0;
  }

  toString() {
    return "";
  }
}

export class BitArraySlice implements BitArray {
  readonly #parent: BitArray;
  readonly #offset: number;
  readonly #length: number;

  public constructor(parent: BitArray, offset: number, length: number) {
    this.#parent = parent;
    this.#offset = offset;
    this.#length = length;
  }

  *[Symbol.iterator]() {
    for (let i = 0; i < this.#length; i++) {
      yield this.#parent.get(i + this.#offset);
    }
  }

  #ensureValidIndex(index: number) {
    if (index < 0 || index >= this.#length) {
      throw new RangeError("The index is outside the BitArray range.");
    }
  }

  #ensureSameSize(other: BitArray) {
    if (this.#length !== other.length) {
      throw new TypeError("The BitArrays do not have the same length.");
    }
  }

  public get length() {
    return this.#length;
  }

  public isFull() {
    for (const bit of this) {
      if (!bit) {
        return false;
      }
    }

    return true;
  }

  public isEmpty() {
    for (const bit of this) {
      if (bit) {
        return false;
      }
    }

    return true;
  }

  public get(index: number) {
    this.#ensureValidIndex(index);

    return this.#parent.get(index + this.#offset);
  }

  public set(index: number, value: boolean) {
    this.#ensureValidIndex(index);

    this.#parent.set(index + this.#offset, value);
  }

  public fill(value: boolean) {
    for (let i = 0; i < this.#length; i++) {
      this.#parent.set(i + this.#offset, value);
    }
  }

  public and(other: BitArray) {
    this.#ensureSameSize(other);

    for (let i = 0; i < this.#length; i++) {
      this.set(i, this.get(i) && other.get(i));
    }

    return this;
  }

  public or(other: BitArray) {
    this.#ensureSameSize(other);

    for (let i = 0; i < this.#length; i++) {
      this.set(i, this.get(i) || other.get(i));
    }

    return this;
  }

  public xor(other: BitArray) {
    this.#ensureSameSize(other);

    for (let i = 0; i < this.#length; i++) {
      this.set(i, this.get(i) !== other.get(i));
    }

    return this;
  }

  public not() {
    for (let i = 0; i < this.#length; i++) {
      this.set(i, !this.get(i));
    }

    return this;
  }

  public contains(other: BitArray) {
    this.#ensureSameSize(other);

    return sequence(this).zip(sequence(other)).where(([, b]) => b).all(([a, b]) => a && b);
  }

  public intersects(other: BitArray) {
    this.#ensureSameSize(other);

    return sequence(this).zip(sequence(other)).any(([a, b]) => a && b);
  }

  public slice(offset: number, length: number) {
    return new BitArraySlice(this.#parent, offset + this.#offset, length);
  }

  public copy() {
    const copy = new BitArrayImpl(this.#length);

    for (let i = 0; i < this.#length; i++) {
      copy.set(i, this.get(i));
    }

    return copy;
  }

  public toArray() {
    return sequence(this).toArray();
  }

  public equals(other: BitArray) {
    return other === this || other && sequence(this).sequenceEquals(sequence(other));
  }

  public toString() {
    return sequence(this).select(bit => bit ? '1' : '0').collect(join());
  }
}