// Copyright 2024 Google LLC

// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at

//     https://www.apache.org/licenses/LICENSE-2.0

// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "common.h"
#include "math.h"
#include "../third_party/z80_trace.h"

#ifdef WASM
FILE * const stderr=NULL;
int fprintf(FILE *stream, const char *format, ...) {return 0;}

#define PAGE_SIZE   0x10000
extern unsigned char __heap_base;
static int _heap_end = (int)&__heap_base;

WASM_EXPORT("alloc")
void * alloc(int size) {
    int ptr = (_heap_end+0xf) & ~0xf; // round up to 16 byte alignment
    _heap_end = ptr + size;
    int pages_needed = (_heap_end+PAGE_SIZE-1) / PAGE_SIZE;
    int pages_n = __builtin_wasm_memory_size(0);
    if (pages_n<pages_needed) {
        __builtin_wasm_memory_grow(0, pages_needed-pages_n);
    }
    return (void *)ptr;
}
#endif

// soup state
BUFFER(rng_state, uint64_t, 1)
BUFFER(soup_dims, int, 2)
DYNAMIC_BUFFER(soup, uint8_t);
DYNAMIC_BUFFER(idx_in_batch, int);

// stats
BUFFER(counts, int, 256)
//BUFFER(counts2, int, 256*256)

// batch scheduling
BUFFER(batch_pair_n, int, 1)
DYNAMIC_BUFFER(batch_idx, int)
DYNAMIC_BUFFER(batch, uint8_t)
DYNAMIC_BUFFER(batch_z80, uint8_t)
DYNAMIC_BUFFER(batch_pc, uint8_t)
DYNAMIC_BUFFER(batch_read_mask, uint32_t)
DYNAMIC_BUFFER(batch_write_mask, uint32_t)

WASM_EXPORT("get_tape_len") int get_tape_len() {return TAPE_LENGTH;}

WASM_EXPORT("alloc_soup")
void alloc_soup(int h, int w) {
    int size = w*h*TAPE_LENGTH;
    soup_alloc(size);
    idx_in_batch_alloc(w*h);
    soup_dims[0] = h;
    soup_dims[1] = w;
}

WASM_EXPORT("alloc_batch")
void alloc_batch(int pair_n) {
    batch_idx_alloc(pair_n*2);
    batch_alloc(pair_n*PAIR_LENGTH);
    batch_z80_alloc(pair_n*sizeof(z80));
    batch_pc_alloc(pair_n);
    batch_read_mask_alloc(pair_n);
    batch_write_mask_alloc(pair_n);
}

inline uint64_t rand64(/*uint64_t seed*/) {
  uint64_t z = (rng_state[0] += 0x9e3779b97f4a7c15);
  z = (z ^ (z >> 30)) * 0xbf58476d1ce4e5b9;
  z = (z ^ (z >> 27)) * 0x94d049bb133111eb;
  return z ^ (z >> 31);
}

WASM_EXPORT("init")
void init(int seed) {
    rng_state[0] = seed;
    for (int i=0; i<soup_len; ++i) {
        soup[i] = rand64();
    }
}

WASM_EXPORT("mutate")
void mutate(int n) {
    for (int i=0; i<n; ++i) {
        uint64_t rnd = rand64();
        uint8_t v = rnd&0xff; rnd >>= 8;
        soup[rnd % soup_len] = v;
    }
}

WASM_EXPORT("prepare_batch") int prepare_batch() {
    const int h=soup_dims[0], w=soup_dims[1], tape_n=h*w;
    int pair_n = 0, collision_count=0, pos=0;
    for (int i=0; i<tape_n; ++i) {idx_in_batch[i] = -1;}
    int max_pair_n = batch_idx_len/2;
    while (pair_n<max_pair_n && collision_count<16) {
        uint64_t rnd = rand64();
        int dir = (rnd&1)*2-1; rnd>>=1;
        int horizontal = rnd&1; rnd>>=1;
        int i = rnd % tape_n, j=i;
        if (idx_in_batch[i]>=0) { ++collision_count; continue;}
        if (horizontal) {
            if (i % w == 0)     { dir =  1; }
            if ((i+1) % w == 0) { dir = -1; }
            j += dir;
        } else {
            if (i < w)          { dir =  1; }
            if (tape_n-i-1 < w) { dir = -1; }
            j += dir*w;
        }
        if (idx_in_batch[j]>=0) { ++collision_count; continue;}
        idx_in_batch[i] = idx_in_batch[j] = pair_n;
        batch_idx[2*pair_n] = i;
        batch_idx[2*pair_n+1] = j;
        for (int k=0; k<TAPE_LENGTH; ++k) {
            batch[pos+k] = soup[i*TAPE_LENGTH+k];
            batch[pos+k+TAPE_LENGTH] = soup[j*TAPE_LENGTH+k];
        }
        pos += TAPE_LENGTH*2;
        pair_n++;
        collision_count = 0;
    }
    batch_pair_n[0] = pair_n;
    return pair_n;
}

WASM_EXPORT("absorb_batch") void absorb_batch() {
    const uint8_t * src = batch;
    const int pair_n = batch_pair_n[0];
    for (int i=0; i<pair_n*2; ++i) {
        const int tape_idx = batch_idx[i];
        uint8_t * dst = soup + tape_idx*TAPE_LENGTH;
        for (int k=0; k<TAPE_LENGTH; ++k,++src,++dst) {
            *dst = *src;
        }
    }
}

WASM_EXPORT("updateCounts")
void updateCounts() {
    for (int i=0; i<256; ++i) counts[i] = 0;
    for (int i=0; i<soup_len; ++i) {
        counts[soup[i]]++;
    }
    // for (int i=0; i<256*256; ++i) counts2[i] = 0;
    // for (int i=0; i<soup_len-1; ++i) {
    //     counts[soup[i]*256+soup[i+1]]++;
    // }
}


// MARK: batch execution
uint8_t memoryRead(void * arg, uint16_t ofs) {
    const int pair_idx = (int)arg;
    ofs &= (PAIR_LENGTH-1);
    return batch[pair_idx*PAIR_LENGTH + ofs];
}
void memoryWrite(void * arg, uint16_t ofs, uint8_t value) {
    const int pair_idx = (int)arg;
    ofs &= (PAIR_LENGTH-1);
    batch[pair_idx*PAIR_LENGTH + ofs] = value;
}
uint8_t inPort(z80* cpu, uint8_t port) { return 0; }
void outPort(z80* cpu, uint8_t port, uint8_t value) {}

WASM_EXPORT("run_batch") int run(int step_n) {
    z80 cpu;
    int totalOps = 0;
    for (int i=0; i<batch_pair_n[0]; ++i) {
        z80_init(&cpu);
        cpu.read_byte = memoryRead;
        cpu.write_byte = memoryWrite;
        cpu.userdata = (void*)i;
        cpu.port_in = inPort;
        cpu.port_out = outPort;
        int step=0;
        for (; step<step_n && !cpu.halted; ++step) {
            z80_step(&cpu);
        }
        totalOps += step;
    }
    return totalOps;
}


// MARK: stepwise execution
int getOffset(int pair_idx, uint16_t ofs) {
    const int tapeIdx = batch_idx[pair_idx*2 + (ofs & TAPE_LENGTH ? 1:0)];
    return tapeIdx*TAPE_LENGTH + (ofs&(TAPE_LENGTH-1));
}
uint8_t memoryRead_inplace(void * arg, uint16_t ofs) {
    const int pair_idx = (int)arg;
    batch_read_mask[pair_idx] |= 1 << (ofs & (PAIR_LENGTH-1));
    return soup[getOffset(pair_idx, ofs)];
}
void memoryWrite_inplace(void * arg, uint16_t ofs, uint8_t value) {
    const int pair_idx = (int)arg;
    batch_write_mask[pair_idx] |= 1 << (ofs & (PAIR_LENGTH-1));
    soup[getOffset(pair_idx, ofs)] = value;
}

WASM_EXPORT("z80_init") void _z80_init() {
    for (int i=0; i<batch_pair_n[0]; ++i) {
        z80 * cpu = (z80*)batch_z80 + i;    
        z80_init(cpu);
        cpu->port_in = inPort;
        cpu->port_out = outPort;
        cpu->read_byte = memoryRead_inplace;
        cpu->write_byte = memoryWrite_inplace;
        cpu->userdata = (void*)i;
        batch_pc[i] = 0;
        batch_read_mask[i] = 0;
        batch_write_mask[i] = 0;
    }
}

WASM_EXPORT("z80_step") void _z80_step() {
    for (int i=0; i<batch_pair_n[0]; ++i) {
        z80 * cpu = (z80*)batch_z80 + i;
        batch_read_mask[i] = 0;
        batch_write_mask[i] = 0;
        z80_step(cpu);
        batch_pc[i] = cpu->pc;
    }
}

WASM_EXPORT("z80_mark") void z80_mark() {
    z80 * cpu = (z80*)batch_z80;
    z80_init(cpu);
    cpu->pc = 0x1111;
    cpu->sp = 0x2222;
    cpu->a = 0x01;
    cpu->b = 0x02;
    cpu->c = 0x03;
    cpu->d = 0x04;
    cpu->e = 0x05;
    cpu->h = 0x06;
    cpu->l = 0x07;
}