#include "world/wetness.h"
#include "world/terrain.h"
#include "world/site.h"
#include "render/texture.h"

#include <stdlib.h>
#include <string.h>

#define WET_RES     2048                      /* cells per side of the whole ground */
#define WET_CHUNK   32                        /* cells per side of a chunk */
#define WET_CHUNKS  (WET_RES / WET_CHUNK)
#define CHUNK_CELLS (WET_CHUNK * WET_CHUNK)
#define HALF        (TERRAIN_SIZE * 0.5f)
#define CELL        (TERRAIN_SIZE / (float)WET_RES)  /* meters */
#define CELL_AREA   (CELL * CELL)                    /* m^2: 1 liter on it is 1 / CELL_AREA mm */
#define SOIL_BLOCK  4                                /* soil and grass looked up per 4 x 4 cells */

#define FLOW_STEP   (1.0f / 30.0f)  /* real seconds per step of running water */
#define FLOW_MAX    0.45f           /* most of a cell's water that may leave in one step */
#define SLOW_STEP   0.25f           /* real seconds between soaking / drying steps */
#define EVAPORATE   (1.5f / 3600.0f) /* mm of standing water per game second (paved: twice) */
#define DRAIN_TIME  1200.0f         /* game seconds: soaked soil drains down to field capacity */
#define FIELD_CAP   0.7f            /* ... which is this much of what it can hold */
#define PEE_FADE    (2.0f * 3600.0f) /* game seconds: a pee stain / smell fades */
#define MUD_FORMS   60.0f           /* game seconds for soaked bare ground to turn to mud */
#define MUD_DRIES   (4.0f * 3600.0f) /* ... and to dry up again */
#define WATER_SHOWN 25.0f           /* mm of standing water at full texture value */
#define PEE_SHOWN   2.0f            /* mm of pee at a full stain */

/* holds: mm of water the top soil takes; soak: mm per game second into dry soil;
   dry: game hours to dry out; mud: how readily bare ground turns to mud;
   speed: how fast water runs over it (m/s at a slope of 1); keep: mm that stays put. */
static const struct { float holds, soak, dry, mud, speed, keep; } SOIL[SOIL_COUNT] = {
    [SOIL_MEADOW] = { 14.0f, 0.10f, 5.0f, 0.8f, 0.8f, 2.0f },
    [SOIL_YARD]   = {  8.0f, 0.03f, 4.0f, 1.0f, 1.6f, 1.5f },
    [SOIL_TRACK]  = {  8.0f, 0.02f, 4.0f, 1.0f, 1.6f, 1.5f },
    [SOIL_FIELD]  = { 20.0f, 0.15f, 6.0f, 1.0f, 1.2f, 2.0f },
    [SOIL_FOREST] = { 22.0f, 0.25f, 8.0f, 0.3f, 0.6f, 2.5f },
    [SOIL_GRAVEL] = {  4.0f, 1.50f, 1.5f, 0.0f, 1.2f, 0.5f },
    [SOIL_PAVED]  = {  0.0f, 0.00f, 1.0f, 0.0f, 3.0f, 0.8f },
};

typedef struct WetChunk {
    int   cx, cz;
    bool  running;  /* standing water above what stays put: runs every flow step */
    bool  dirty;    /* texture out of date */
    bool  dead;     /* all dry: freed once the texture is cleared */
    float ground[CHUNK_CELLS]; /* mm, ground height incl. trampled paths */
    float water[CHUNK_CELLS];  /* mm standing on top */
    float soil[CHUNK_CELLS];   /* mm held in the soil */
    float pee[CHUNK_CELLS];    /* mm of pee in it all */
    float mud[CHUNK_CELLS];    /* 0..1 */
    float grass[CHUNK_CELLS];  /* how much grass grows here (0..1) */
    unsigned char kind[CHUNK_CELLS];
    float dw[CHUNK_CELLS], dp[CHUNK_CELLS]; /* flow in / out this step */
    unsigned char texels[CHUNK_CELLS * 4];
} WetChunk;

/* ---------- Cells and chunks ---------- */

static float cellCenter(int g) { return -HALF + ((float)g + 0.5f) * CELL; }

static WetChunk *chunkAt(const Wetness *w, int cx, int cz) {
    if (cx < 0 || cz < 0 || cx >= WET_CHUNKS || cz >= WET_CHUNKS) return NULL;
    return w->chunks[cz * WET_CHUNKS + cx];
}

static WetChunk *newChunk(Terrain *t, int cx, int cz) {
    Wetness *w = &t->wet;
    WetChunk *c = calloc(1, sizeof(WetChunk));
    if (!c) return NULL;
    c->cx = cx;
    c->cz = cz;
    for (int bz = 0; bz < WET_CHUNK; bz += SOIL_BLOCK)
        for (int bx = 0; bx < WET_CHUNK; bx += SOIL_BLOCK) {
            float x = cellCenter(cx * WET_CHUNK + bx) + CELL * (SOIL_BLOCK - 1) * 0.5f;
            float z = cellCenter(cz * WET_CHUNK + bz) + CELL * (SOIL_BLOCK - 1) * 0.5f;
            Soil kind = Site_soil(x, z);
            float grass = Site_grass(x, z);
            for (int k = 0; k < SOIL_BLOCK * SOIL_BLOCK; k++) {
                int i = (bz + k / SOIL_BLOCK) * WET_CHUNK + bx + k % SOIL_BLOCK;
                c->kind[i] = (unsigned char)kind;
                c->grass[i] = grass;
            }
        }
    for (int i = 0; i < CHUNK_CELLS; i++)
        c->ground[i] = Terrain_height(t, cellCenter(cx * WET_CHUNK + i % WET_CHUNK), cellCenter(cz * WET_CHUNK + i / WET_CHUNK)) * 1000.0f;
    w->chunks[cz * WET_CHUNKS + cx] = c;
    w->live[w->live_count++] = cz * WET_CHUNKS + cx;
    return c;
}

#define WET_SAVE_TAG 0x57455431u /* "WET1" */

bool Wetness_write(const Wetness *w, FILE *f) {
    unsigned tag = WET_SAVE_TAG, n = 0, res = WET_RES, chunk = WET_CHUNK;
    for (int i = 0; i < w->live_count; i++) { const WetChunk *c = w->chunks[w->live[i]]; if (c && !c->dead) n++; }
    if (fwrite(&tag, 4, 1, f) != 1 || fwrite(&res, 4, 1, f) != 1 || fwrite(&chunk, 4, 1, f) != 1 || fwrite(&n, 4, 1, f) != 1) return false;
    for (int i = 0; i < w->live_count; i++) {
        const WetChunk *c = w->chunks[w->live[i]];
        if (!c || c->dead) continue;
        int head[2] = { c->cx, c->cz };
        if (fwrite(head, sizeof head, 1, f) != 1 || fwrite(c->water, sizeof c->water, 1, f) != 1 || fwrite(c->soil, sizeof c->soil, 1, f) != 1 ||
            fwrite(c->pee, sizeof c->pee, 1, f) != 1 || fwrite(c->mud, sizeof c->mud, 1, f) != 1) return false;
    }
    return true;
}

bool Wetness_read(Terrain *t, FILE *f) {
    unsigned tag, n, res, chunk;
    if (fread(&tag, 4, 1, f) != 1 || tag != WET_SAVE_TAG || fread(&res, 4, 1, f) != 1 || fread(&chunk, 4, 1, f) != 1 || fread(&n, 4, 1, f) != 1) return false;
    if (res != WET_RES || chunk != WET_CHUNK) return false; /* (another grid: the ground starts dry) */
    for (unsigned k = 0; k < n; k++) {
        int head[2];
        WetChunk tmp;
        if (fread(head, sizeof head, 1, f) != 1 || fread(tmp.water, sizeof tmp.water, 1, f) != 1 || fread(tmp.soil, sizeof tmp.soil, 1, f) != 1 ||
            fread(tmp.pee, sizeof tmp.pee, 1, f) != 1 || fread(tmp.mud, sizeof tmp.mud, 1, f) != 1) return false;
        if (head[0] < 0 || head[1] < 0 || head[0] >= WET_CHUNKS || head[1] >= WET_CHUNKS) continue;
        WetChunk *c = chunkAt(&t->wet, head[0], head[1]);
        if (!c) c = newChunk(t, head[0], head[1]);
        if (!c) return false;
        memcpy(c->water, tmp.water, sizeof c->water);
        memcpy(c->soil, tmp.soil, sizeof c->soil);
        memcpy(c->pee, tmp.pee, sizeof c->pee);
        memcpy(c->mud, tmp.mud, sizeof c->mud);
        c->running = true;
        c->dirty = true;
        c->dead = false;
    }
    return true;
}

/* The chunk and index of cell (gx, gz), made if missing and make is set. */
static WetChunk *cellAt(Terrain *t, int gx, int gz, bool make, int *index) {
    if (gx < 0 || gz < 0 || gx >= WET_RES || gz >= WET_RES) return NULL;
    WetChunk *c = chunkAt(&t->wet, gx / WET_CHUNK, gz / WET_CHUNK);
    if (!c && make) c = newChunk(t, gx / WET_CHUNK, gz / WET_CHUNK);
    if (c) *index = (gz % WET_CHUNK) * WET_CHUNK + gx % WET_CHUNK;
    return c;
}

static const WetChunk *cellAtConst(const Wetness *w, float x, float z, int *index) {
    int gx = (int)floorf((x + HALF) / CELL), gz = (int)floorf((z + HALF) / CELL);
    if (gx < 0 || gz < 0 || gx >= WET_RES || gz >= WET_RES) return NULL;
    const WetChunk *c = chunkAt(w, gx / WET_CHUNK, gz / WET_CHUNK);
    if (c) *index = (gz % WET_CHUNK) * WET_CHUNK + gx % WET_CHUNK;
    return c;
}

static float saturation(const WetChunk *c, int i) {
    float holds = SOIL[c->kind[i]].holds;
    return holds > 0.0f ? clampf(c->soil[i] / holds, 0.0f, 1.0f) : 0.0f;
}

/* ---------- Public: setup ---------- */

bool Wetness_init(Wetness *w) {
    memset(w, 0, sizeof(*w));
    w->chunks = calloc((size_t)WET_CHUNKS * WET_CHUNKS, sizeof(WetChunk *));
    w->live = calloc((size_t)WET_CHUNKS * WET_CHUNKS, sizeof(int));
    w->texture = Texture_createRGBA8(WET_RES, WET_RES);
    if (!w->chunks || !w->live || !w->texture) {
        Wetness_destroy(w);
        return false;
    }
    return true;
}

void Wetness_destroy(Wetness *w) {
    if (w->chunks)
        for (int i = 0; i < WET_CHUNKS * WET_CHUNKS; i++) free(w->chunks[i]);
    free(w->chunks);
    free(w->live);
    Texture_destroy(&w->texture);
    memset(w, 0, sizeof(*w));
}

/* ---------- Public: water coming down ---------- */

void Wetness_pour(Terrain *t, float x, float z, float liters, float radius, float pee) {
    if (liters <= 0.0f) return;
    radius = fmaxf(radius, CELL * 0.5f);
    int g0x = (int)floorf((x - radius + HALF) / CELL), g1x = (int)floorf((x + radius + HALF) / CELL);
    int g0z = (int)floorf((z - radius + HALF) / CELL), g1z = (int)floorf((z + radius + HALF) / CELL);
    /* weights fall off from the middle; at least the cell it hits */
    float total = 0.0f;
    for (int pass = 0; pass < 2; pass++) {
        for (int gz = g0z; gz <= g1z; gz++)
            for (int gx = g0x; gx <= g1x; gx++) {
                float dx = cellCenter(gx) - x, dz = cellCenter(gz) - z;
                float weight = fmaxf(0.0f, 1.0f - sqrtf(dx * dx + dz * dz) / radius);
                if (g0x == g1x && g0z == g1z) weight = 1.0f;
                if (weight <= 0.0f) continue;
                if (pass == 0) { total += weight; continue; }
                int i;
                WetChunk *c = cellAt(t, gx, gz, true, &i);
                if (!c) continue;
                float mm = liters / CELL_AREA * weight / total;
                c->water[i] += mm;
                c->pee[i] += mm * pee;
                c->running = c->dirty = true;
                c->dead = false; /* wet again before it was cleared */
            }
        if (total <= 0.0f) { /* smaller than a cell and between centers: the cell it hits */
            int i;
            WetChunk *c = cellAt(t, (int)floorf((x + HALF) / CELL), (int)floorf((z + HALF) / CELL), true, &i);
            if (!c) return;
            c->water[i] += liters / CELL_AREA;
            c->pee[i] += liters / CELL_AREA * pee;
            c->running = c->dirty = true;
            c->dead = false;
            return;
        }
    }
}

/* ---------- Running water ---------- */

static const int NEIGHBOR[4][2] = { { 1, 0 }, { -1, 0 }, { 0, 1 }, { 0, -1 } };

/* Standing water runs to lower neighbors (ground + water level), the faster
   the steeper and the deeper; pee goes along with it. */
static void flowChunk(Terrain *t, WetChunk *c) {
    bool running = false;
    for (int i = 0; i < CHUNK_CELLS; i++) {
        const float keep = SOIL[c->kind[i]].keep;
        float free_mm = c->water[i] - keep;
        if (free_mm <= 0.0f) continue;
        running = true;
        float level = c->ground[i] + c->water[i];
        int gx = c->cx * WET_CHUNK + i % WET_CHUNK, gz = c->cz * WET_CHUNK + i / WET_CHUNK;
        float depth = clampf(cbrtf((free_mm / 4.0f) * (free_mm / 4.0f)), 0.15f, 2.0f); /* shallow water creeps */

        WetChunk *to[4];
        int at[4];
        float share[4], drop[4], sum = 0.0f;
        for (int k = 0; k < 4; k++) {
            to[k] = NULL;
            share[k] = 0.0f;
            int nx = gx + NEIGHBOR[k][0], nz = gz + NEIGHBOR[k][1], j;
            WetChunk *n = cellAt(t, nx, nz, false, &j);
            float other;
            if (n) other = n->ground[j] + n->water[j];
            else if (nx < 0 || nz < 0 || nx >= WET_RES || nz >= WET_RES) continue;
            else other = Terrain_height(t, cellCenter(nx), cellCenter(nz)) * 1000.0f;
            drop[k] = level - other;
            if (drop[k] <= 0.0f) continue;
            float slope = drop[k] / (CELL * 1000.0f);
            share[k] = SOIL[c->kind[i]].speed * sqrtf(slope) * depth * FLOW_STEP / CELL;
            if (!n) n = cellAt(t, nx, nz, true, &j);
            if (!n) { share[k] = 0.0f; continue; }
            to[k] = n;
            at[k] = j;
            sum += share[k];
        }
        if (sum <= 0.0f) continue;
        float scale = sum > FLOW_MAX ? FLOW_MAX / sum : 1.0f;
        float body = c->water[i] + c->soil[i] + 1e-3f;
        for (int k = 0; k < 4; k++) {
            if (!to[k]) continue;
            float moved = fminf(free_mm * share[k] * scale, drop[k] * 0.3f);
            float pee = c->pee[i] * moved / body;
            c->dw[i] -= moved;
            c->dp[i] -= pee;
            to[k]->dw[at[k]] += moved;
            to[k]->dp[at[k]] += pee;
            to[k]->running = true;
            to[k]->dead = false;
        }
    }
    c->running = running;
}

static void flowStep(Terrain *t) {
    Wetness *w = &t->wet;
    for (int n = 0; n < w->live_count; n++) {
        WetChunk *c = w->chunks[w->live[n]];
        memset(c->dw, 0, sizeof(c->dw));
        memset(c->dp, 0, sizeof(c->dp));
    }
    int count = w->live_count; /* chunks made while flowing start with nothing to give */
    for (int n = 0; n < count; n++) {
        WetChunk *c = w->chunks[w->live[n]];
        if (c->running) flowChunk(t, c);
    }
    for (int n = 0; n < w->live_count; n++) {
        WetChunk *c = w->chunks[w->live[n]];
        bool changed = false;
        for (int i = 0; i < CHUNK_CELLS; i++) {
            if (c->dw[i] == 0.0f && c->dp[i] == 0.0f) continue;
            c->water[i] = fmaxf(0.0f, c->water[i] + c->dw[i]);
            c->pee[i] = fmaxf(0.0f, c->pee[i] + c->dp[i]);
            changed = true;
        }
        if (changed) c->dirty = true;
    }
}

/* ---------- Soaking in, drying, mud ---------- */

static void slowChunk(Terrain *t, WetChunk *c, float gdt) {
    float most = 0.0f;
    for (int i = 0; i < CHUNK_CELLS; i++) {
        const int kind = c->kind[i];
        const float holds = SOIL[kind].holds;
        float x = cellCenter(c->cx * WET_CHUNK + i % WET_CHUNK), z = cellCenter(c->cz * WET_CHUNK + i / WET_CHUNK);
        c->ground[i] = Terrain_height(t, x, z) * 1000.0f; /* paths get trodden in */

        /* soaking in: dry soil drinks fastest */
        float s = saturation(c, i);
        float soak = fminf(c->water[i], fminf(SOIL[kind].soak * (0.3f + 2.0f * (1.0f - s)) * gdt, holds - c->soil[i]));
        if (soak > 0.0f) {
            c->water[i] -= soak;
            c->soil[i] += soak;
        }
        c->water[i] = fmaxf(0.0f, c->water[i] - EVAPORATE * (kind == SOIL_PAVED ? 2.0f : 1.0f) * gdt);

        /* draining down (fast while soaked) and drying out (slow) */
        float drained = 0.0f;
        if (holds > 0.0f) {
            float field = FIELD_CAP * holds;
            drained = c->soil[i] > field ? (c->soil[i] - field) * fminf(1.0f, gdt / DRAIN_TIME)
                                         : fminf(c->soil[i], holds * (0.3f + 0.7f * s) * gdt / (SOIL[kind].dry * 3600.0f));
            c->soil[i] -= drained;
        }
        /* pee: washed down with the water, and fading */
        float body = c->water[i] + c->soil[i] + drained + 1e-3f;
        c->pee[i] -= c->pee[i] * fminf(1.0f, drained / body + gdt / PEE_FADE);
        if (c->water[i] + c->soil[i] < 0.05f) c->pee[i] *= 1.0f - fminf(1.0f, 3.0f * gdt / PEE_FADE);

        /* mud: soaked bare soft ground */
        s = saturation(c, i);
        float bare = 1.0f - c->grass[i] * (1.0f - Terrain_wear(t, x, z));
        float target = SOIL[kind].mud * bare * smoothstepf(0.55f, 0.95f, s);
        if (target > c->mud[i]) c->mud[i] += (target - c->mud[i]) * fminf(1.0f, gdt / MUD_FORMS);
        else c->mud[i] -= (c->mud[i] - target) * fminf(1.0f, gdt / MUD_DRIES);

        /* what can't be seen any more is gone */
        if (c->water[i] < 0.01f) c->water[i] = 0.0f;
        if (c->water[i] == 0.0f && c->soil[i] < 0.02f) c->soil[i] = 0.0f;
        if (c->water[i] + c->soil[i] == 0.0f && c->pee[i] < 0.02f) c->pee[i] = 0.0f;
        if (c->mud[i] < 0.1f && s < 0.3f) c->mud[i] = 0.0f;

        if (c->water[i] > SOIL[kind].keep) c->running = true;
        most = fmaxf(most, c->water[i] + c->soil[i] + c->pee[i] + c->mud[i]);
    }
    c->dirty = true;
    if (most <= 0.0f) c->dead = true;
}

void Wetness_update(Terrain *t, float dt, float game_dt) {
    Wetness *w = &t->wet;
    if (w->live_count == 0) {
        w->flow_time = w->slow_time = w->slow_game = 0.0f;
        return;
    }
    w->flow_time = fminf(w->flow_time + dt, FLOW_STEP * 4.0f); /* falls behind rather than stalls */
    while (w->flow_time >= FLOW_STEP) {
        w->flow_time -= FLOW_STEP;
        flowStep(t);
    }
    w->slow_time += dt;
    w->slow_game += game_dt;
    if (w->slow_time >= SLOW_STEP) {
        for (int n = 0; n < w->live_count; n++) slowChunk(t, w->chunks[w->live[n]], w->slow_game);
        w->slow_time = w->slow_game = 0.0f;
    }
}

/* ---------- Texture ---------- */

static unsigned char byte(float v) { return (unsigned char)(clampf(v, 0.0f, 1.0f) * 255.0f + 0.5f); }

void Wetness_sync(Wetness *w) {
    for (int n = 0; n < w->live_count; n++) {
        WetChunk *c = w->chunks[w->live[n]];
        if (!c->dirty) continue;
        for (int i = 0; i < CHUNK_CELLS; i++) {
            unsigned char *p = &c->texels[i * 4];
            if (c->dead) { p[0] = p[1] = p[2] = p[3] = 0; continue; }
            p[0] = byte(c->kind[i] == SOIL_PAVED ? fminf(c->water[i] * 0.5f, 1.0f) : saturation(c, i));
            p[1] = byte(sqrtf(c->water[i] / WATER_SHOWN));
            p[2] = byte(c->pee[i] / PEE_SHOWN);
            p[3] = byte(c->mud[i]);
        }
        Texture_updateRGBA8(w->texture, c->cx * WET_CHUNK, c->cz * WET_CHUNK, WET_CHUNK, WET_CHUNK, c->texels);
        c->dirty = false;
        if (c->dead) { /* cleared on the GPU: gone */
            w->chunks[w->live[n]] = NULL;
            free(c);
            w->live[n--] = w->live[--w->live_count];
        }
    }
}

/* ---------- Queries ---------- */

float Wetness_water(const Wetness *w, float x, float z) {
    int i;
    const WetChunk *c = cellAtConst(w, x, z, &i);
    return c ? c->water[i] : 0.0f;
}

float Wetness_mud(const Wetness *w, float x, float z) {
    int i;
    const WetChunk *c = cellAtConst(w, x, z, &i);
    return c ? c->mud[i] : 0.0f;
}

float Wetness_softness(const Wetness *w, float x, float z) {
    int i;
    const WetChunk *c = cellAtConst(w, x, z, &i);
    return c ? SOIL[c->kind[i]].mud * smoothstepf(0.4f, 0.9f, saturation(c, i)) : 0.0f;
}

float Wetness_pee(const Wetness *w, float x, float z, float radius) {
    float most = 0.0f;
    const int n = (int)(radius / CELL);
    for (int kz = -n; kz <= n; kz++)
        for (int kx = -n; kx <= n; kx++) {
            if (kx * kx + kz * kz > n * n) continue;
            int i;
            const WetChunk *c = cellAtConst(w, x + (float)kx * CELL, z + (float)kz * CELL, &i);
            if (c) most = fmaxf(most, c->pee[i]);
        }
    return clampf(most / (PEE_SHOWN * 0.5f), 0.0f, 1.0f);
}

#define STEP_RADIUS 0.12f /* a shoe */

float Wetness_step(Wetness *w, float x, float z) {
    float soft = Wetness_softness(w, x, z), mud = Wetness_mud(w, x, z);
    if (soft <= 0.05f) return mud;
    const int n = (int)(STEP_RADIUS / CELL + 0.5f);
    for (int kz = -n; kz <= n; kz++)
        for (int kx = -n; kx <= n; kx++) {
            int i;
            WetChunk *c = (WetChunk *)cellAtConst(w, x + (float)kx * CELL, z + (float)kz * CELL, &i);
            if (!c) continue;
            c->mud[i] = fminf(1.0f, c->mud[i] + 0.06f * soft);
            c->dirty = true;
            c->dead = false;
        }
    return mud;
}