import { Option, none, some } from 'fp-ts/lib/Option.js';
import Prando from 'prando';

// a shrinkable generator needs to maintain some state,
// namely the sequence of generated values in order to be
// able to backtrack and shrink.

export interface Generators {

  genInt(): Gen<number>;

  genIntWithin(min: number, max: number): Gen<number>;

  genBytes(): GenBytes;

  genString(): Gen<string>;

  genList<A>(generator: Gen<A>): Gen<A[]>;

  replay(): void;
}

export interface Gen<A> {
  generate(): A;

  generateN(n: number): A[];
}

interface GenBytes {
  reject(): void;
  generate(n: number): number[];
}

class BytesGen implements GenBytes {
  prng: Prando;
  bytes: number[][]; // the internal sequence of generated values

  constructor(prng: Prando) {
    this.prng = prng;
    this.bytes = [];
  }

  generate(n: number): number[] {
    if (n < 0) {
      throw new Error('n must be greater than 0');
    }
    return this.generateBytes(n);
  }

  generateBytes(n: number): number[] {
    const gens = [];
    for (let i = 0; i < n; i++) {
      gens.push(this.prng.nextInt(0, 255));
    }
    this.bytes.push([...gens]);
    return gens;
  }

  reject() {
    if (this.bytes.length === 0) {
      throw new Error('No more bytes to shrink');
    }
    this.bytes.pop();
  }

  replay(): GenBytes {
    return new ReplayBytes(this.bytes);
  }
}

class ReplayBytes implements GenBytes {
  bytes: number[][]; // the internal sequence of generated values
  index: number = 0; // the current index in the sequence of generated values

  constructor(bytes: number[][]) {
    this.bytes = bytes;
    this.index = 0;
  }

  generate(n: number): number[] {
    if (this.index >= this.bytes.length) {
      throw new Error('No more bytes to shrink');
    }

    const chunk = [...this.bytes[this.index++]];

    if (chunk.length < n) {
      const zeros = new Array(n - chunk.length).fill(0);
      chunk.unshift(...zeros);
    }

    return chunk;
  }

  reject() {
    if (this.index >= this.bytes.length) {
      throw new Error('No more bytes to shrink');
    }
    this.index++;
  }
}

class GenInt implements Gen<number> {
  constructor(private bytes: GenBytes) {
    this.bytes = bytes;
  }

  generate(): number {
    const bytes = this.bytes.generate(1);
    return bytes[0];
  }

  generateN(n: number): number[] {
    const result = [];
    for (let i = 0; i < n; i++) {
      result.push(this.bytes.generate(1)[0]);
    }

    return result;
  }
}

class GenIntWithin implements Gen<number> {
  constructor(private bytes: GenBytes, private min: number, private max: number) {
    this.bytes = bytes;
    this.min = min;
    this.max = max;
  }

  generate(): number {
    const range = this.max - this.min;
    const numBytes = Math.ceil(Math.log2(range) / 8);
    const maxBytes = Math.pow(2, numBytes * 8);
    const upperBound = maxBytes - (maxBytes % range);
    let bytes = this.bytes.generate(numBytes);
    let rejection = 0;
    while (true) {
      let value = 0;
      for (let i = 0; i < bytes.length; i++) {
        value = (value << 8) | bytes[i];
      }

      if (value >= upperBound) {
        this.bytes.reject();
        bytes = this.bytes.generate(numBytes);
        continue;
      }

      if (rejection > 0) {
        console.log(`Rejection sampling: rejected ${rejection} values before accepting ${value}`);
      }
      return this.min + (value % range);
    }
  }

  generateN(n: number): number[] {
    const result = [];
    for (let i = 0; i < n; i++) {
      result.push(this.bytes.generate(1)[0]);
    }

    return result;
  }

}

class GenASCIIString implements Gen<string> {
  constructor(private bytes: GenBytes) {
    this.bytes = bytes
  }

  generate(): string {
    const length = this.bytes.generate(1)[0]; // limit to 20 characters
    const bytes = this.bytes.generate(length);
    return String.fromCharCode(...bytes.map(b => b % 128)); // limit to ASCII
  }

  generateN(n: number): string[] {
    const result = [];
    for (let i = 0; i < n; i++) {
      result.push(this.generate());
    }
    return result;
  }
}

class Seeded implements Generators {

  constructor(private bytes: BytesGen) {
    this.bytes = bytes;
  }

  replay(): Generators {
    return new Replayed(this.bytes);
  }

  genInt(): Gen<number> {
    return new GenInt(this.bytes);
  }
  genIntWithin(min: number, max: number): Gen<number> {
    return new GenIntWithin(this.bytes, min, max);
  }
  genBytes(): GenBytes {
    return this.bytes;
  }
  genString(): Gen<string> {
    return new GenASCIIString(this.bytes);
  }
  genList<A>(generator: Gen<A>): Gen<A[]> {
    throw new Error('Method not implemented.');
  }

}

class Replayed implements Generators {
  bytes : GenBytes;

  constructor (bytes: BytesGen) {
    this.bytes = bytes.replay();
  }

  replay(): Generators {
    throw new Error('Cannot replay a replayed generator');
  }

  genInt(): Gen<number> {
    return new GenInt(this.bytes);
  }
  genIntWithin(min: number, max: number): Gen<number> {
    return new GenIntWithin(this.bytes, min, max);
  }
  genBytes(): GenBytes {
    return this.bytes;
  }
  genString(): Gen<string> {
    return new GenASCIIString(this.bytes);
  }
  genList<A>(generator: Gen<A>): Gen<A[]> {
    throw new Error('Method not implemented.');
  }

}


const generators = (seed: number | undefined = undefined) => {
  const prng = new Prando(seed ? seed : Math.random() * 10000000);
  const bytes = new BytesGen(prng);

  return new Seeded(bytes);
};

export default generators;