precision highp float;
precision highp int;

#define PI 3.1415926535897932384626433832795

in vec2 vUv;

layout(location = 0) out vec4 FragColor;
layout(location = 1) out vec4 FragColorExt;

uniform float uTime;
uniform sampler2D tLastState;
uniform sampler2D tLastExtState;
uniform sampler2D tWorld;
uniform sampler2D tPresence;
uniform float uForagerRatio;
uniform sampler2D uMaterialProps;
uniform vec4 uAntSpawn;

const float ANT_CARRY_STRENGTH = 1.0;
const float sampleDistance = 20.;
const float cellSize = 1. / WORLD_SIZE;

float rand(vec2 co) {
    float a = 12.9898;
    float b = 78.233;
    float c = 43758.5453;
    float dt = dot(co.xy ,vec2(a,b));
    float sn = mod(dt, 3.14);
    return fract(sin(sn) * c);
}

vec2 roundUvToCellCenter(vec2 uv) {
    return floor(uv * WORLD_SIZE) / WORLD_SIZE + cellSize * 0.5;
}

bool tryDropFood(vec2 pos) {
    float currentMatId = texture(tWorld, roundUvToCellCenter(pos)).x;
    float belowMatId = texture(tWorld, roundUvToCellCenter(pos - vec2(0., cellSize))).x;
    return int(currentMatId) == MAT_HOME || int(belowMatId) == MAT_HOME;
}

bool isObstacle(vec2 pos) {
    float materialId = texture(tWorld, roundUvToCellCenter(pos)).x;
    int matInt = int(materialId);
    // ants can only move through air and home cells
    return matInt != MAT_AIR && matInt != MAT_HOME;
}

float smell(vec2 pos, float isCarrying) {
    vec2 value = texture(tWorld, pos).yz;

    if (isCarrying > 0.5) {
        return value.y;
    }

    return value.x;
}

vec2 applyOffsetToPos(vec2 pos, vec2 offsetDirectaion) {
    vec2 newPos = clamp(pos + offsetDirectaion * cellSize, 0., 1.);

    if (!isObstacle(pos) && isObstacle(newPos)) {
        return pos;
    }

    return newPos;
}

float getMaxScentStorage(vec2 antDataUv) {
    float factor = 0.8 + rand(antDataUv * 100.) * 0.2;

    return SCENT_MAX_STORAGE * factor;
}

void main()	{
    vec4 lastState = texture(tLastState, vUv);
    vec4 lastExtState = texture(tLastExtState, vUv);

    float noise = rand(vUv * 1000. + fract(uTime / 1000.));

    vec2 pos = lastState.xy;
    float angle = lastState.z;
    float isCarrying = float(int(lastState.w) & 1);
    float storage = float(int(lastState.w) >> 1);
    bool wasObstacle = isObstacle(pos);

    float personality = lastExtState.r;
    float cargoMaterialId = lastExtState.g;
    float pathIntDx = lastExtState.b;
    float pathIntDy = lastExtState.a;

    // calculate this ant's index from its UV position in the texture
    ivec2 texSize = textureSize(tLastState, 0);
    int antIndex = int(vUv.y * float(texSize.y)) * texSize.x + int(vUv.x * float(texSize.x));

    if (pos == vec2(0)) { // init new ant
        bool spawnRequested = (uAntSpawn.w > 0.5);

        if (antIndex >= ANTS_START_COUNT && !spawnRequested) {
            // dormant ant, no spawn request — skip
            FragColor = vec4(0);
            FragColorExt = vec4(0);
            return;
        }

        if (antIndex >= ANTS_START_COUNT && spawnRequested) {
            // probabilistic activation: ~brushRadius ants per frame
            float activationChance = uAntSpawn.z / float(ANT_BUDGET);
            float roll = rand(vUv * 50000. + fract(uTime / 1000.));
            if (roll > activationChance) {
                // not selected this frame — stay dormant
                FragColor = vec4(0);
                FragColorExt = vec4(0);
                return;
            }
            // activate at spawn position with scatter
            float scatter = uAntSpawn.z / WORLD_SIZE;
            float rngX = rand(vUv * 10000. + fract(uTime / 1000.));
            float rngY = rand(vUv * 20000. + fract(uTime / 1000.) + 0.5);
            pos = vec2(
                uAntSpawn.x + (rngX - 0.5) * scatter,
                uAntSpawn.y + (rngY - 0.5) * scatter
            );
            pos = clamp(pos, 0., 1.);
            angle = rand(vUv * 42069.) * 2.0 * PI;
            isCarrying = 0.;
            storage = 0.;
            personality = rand(vUv * 42069.);
            cargoMaterialId = 0.;
            pathIntDx = 0.;
            pathIntDy = 0.;
        } else {
            // normal init for starting ants
            #if VIEW_MODE_SIDE
            // spawn on sand surface: random X, scan down from top to find first non-air cell
            float spawnX = rand(vUv * 10000.);
            float pixelSize = 1.0 / float(textureSize(tWorld, 0).x);
            float surfaceY = 0.6; // fallback if scan finds nothing
            for (float scanY = 1.0; scanY > 0.0; scanY -= pixelSize) {
                float matId = texture(tWorld, vec2(spawnX, scanY)).x;
                if (matId > 0.5) { // non-air cell found
                    surfaceY = scanY + pixelSize; // one cell above the surface
                    break;
                }
            }
            pos = vec2(spawnX, surfaceY);
            angle = (rand(vUv * 31337.) > 0.5) ? 0.0 : PI; // face randomly left or right
            #else
            pos = vec2(0.5);
            angle = rand(vUv * 10000.) * 2. * PI;
            #endif
            isCarrying = 0.;
            storage = 0.;
            personality = rand(vUv * 42069.); // 0.0 = pure follower, 1.0 = pure explorer
            cargoMaterialId = 0.;
            pathIntDx = 0.;
            pathIntDy = 0.;
        }
    }

    // --- GRAVITY AND SURFACE CHECK ---
    vec2 cellCenter = roundUvToCellCenter(pos);
    float belowMat = texture(tWorld, cellCenter - vec2(0., cellSize)).x;
    float currentMat = texture(tWorld, cellCenter).x;

    bool belowSolid = (int(belowMat) != MAT_AIR);

    // collision displacement: if current cell is now solid (sand fell on us), push to nearest air
    if (int(currentMat) != MAT_AIR && int(currentMat) != MAT_HOME) {
        vec2 upPos = cellCenter + vec2(0., cellSize);
        vec2 leftPos = cellCenter - vec2(cellSize, 0.);
        vec2 rightPos = cellCenter + vec2(cellSize, 0.);
        if (int(texture(tWorld, upPos).x) == MAT_AIR) {
            pos = upPos;
        } else if (int(texture(tWorld, leftPos).x) == MAT_AIR) {
            pos = leftPos;
        } else if (int(texture(tWorld, rightPos).x) == MAT_AIR) {
            pos = rightPos;
        }
        cellCenter = roundUvToCellCenter(pos);
        belowMat = texture(tWorld, cellCenter - vec2(0., cellSize)).x;
        belowSolid = (int(belowMat) != MAT_AIR);
    }

    bool isFalling = false;

    // GRAVITY: if nothing solid below and not at world bottom, fall multiple cells
    if (!belowSolid && pos.y > cellSize) {
        for (int g = 0; g < 4; g++) {
            if (pos.y <= cellSize) break;
            float matBelow = texture(tWorld, roundUvToCellCenter(pos - vec2(0., cellSize))).x;
            if (int(matBelow) != MAT_AIR) break;
            pos.y -= cellSize;
        }
        isFalling = true;
    }

    if (!isFalling) {
        bool acted = false;

        // ---- PRIORITY 1: DEPOSIT ----
        if (!acted && isCarrying == 1.) {
            float cellMatId = texture(tWorld, roundUvToCellCenter(pos)).x;

            // carrying food and at home -> drop food
            if (int(cargoMaterialId) == MAT_FOOD && tryDropFood(pos)) {
                isCarrying = 0.;
                cargoMaterialId = 0.;
                angle += PI;
                storage = getMaxScentStorage(vUv);
                acted = true;
            }

            // carrying powder and on surface (air cell, solid below) -> deposit
            if (!acted && int(cargoMaterialId) != MAT_FOOD) {
                if (int(cellMatId) == MAT_AIR) {
                    vec2 belowPos = pos - vec2(0., cellSize);
                    float belowMatId = texture(tWorld, roundUvToCellCenter(belowPos)).x;
                    vec2 abovePos = pos + vec2(0., cellSize);
                    float aboveMatId = texture(tWorld, roundUvToCellCenter(abovePos)).x;
                    if ((int(belowMatId) != MAT_AIR || belowPos.y <= 0.)
                        && int(aboveMatId) == MAT_AIR) {
                        isCarrying = 0.;
                        // keep cargoMaterialId set so discretize reads it this frame
                        angle += PI;
                        storage = getMaxScentStorage(vUv);
                        acted = true;
                    }
                }
            }
        }

        // ---- PRIORITY 2: NAVIGATE (carrying) ----
        if (!acted && isCarrying == 1.) {
            if (int(cargoMaterialId) == MAT_FOOD) {
                // follow toHome pheromone
                float sAhead = smell(applyOffsetToPos(pos, vec2(cos(angle), sin(angle)) * sampleDistance), 1.);
                float sLeft = smell(applyOffsetToPos(pos, vec2(cos(angle - ANT_ROTATION_ANGLE), sin(angle - ANT_ROTATION_ANGLE)) * sampleDistance), 1.);
                float sRight = smell(applyOffsetToPos(pos, vec2(cos(angle + ANT_ROTATION_ANGLE), sin(angle + ANT_ROTATION_ANGLE)) * sampleDistance), 1.);

                if (sLeft > sAhead && sLeft > sRight) {
                    angle -= ANT_ROTATION_ANGLE;
                } else if (sRight > sAhead && sRight > sLeft) {
                    angle += ANT_ROTATION_ANGLE;
                }
            } else {
                // carrying powder: bias upward toward surface
                #if VIEW_MODE_SIDE
                float upwardBias = PI * 0.5;
                float angleDiff = upwardBias - angle;
                angleDiff = mod(angleDiff + PI, 2.0 * PI) - PI;
                angle += angleDiff * 0.3;
                #endif
            }
            // add wander noise to carrying ants too
            float noise2 = rand(vUv * 1000. + fract(uTime / 1000.) + 0.2);
            if (noise2 > 0.5) {
                angle += ANT_ROTATION_ANGLE;
            } else {
                angle -= ANT_ROTATION_ANGLE;
            }
            acted = true;
        }

        // ---- PRIORITY 3: FORAGE (not carrying, food scent detected) ----
        if (!acted && isCarrying == 0.) {
            float sAhead = smell(applyOffsetToPos(pos, vec2(cos(angle), sin(angle)) * sampleDistance), 0.);
            float sLeft = smell(applyOffsetToPos(pos, vec2(cos(angle - ANT_ROTATION_ANGLE), sin(angle - ANT_ROTATION_ANGLE)) * sampleDistance), 0.);
            float sRight = smell(applyOffsetToPos(pos, vec2(cos(angle + ANT_ROTATION_ANGLE), sin(angle + ANT_ROTATION_ANGLE)) * sampleDistance), 0.);

            float maxSmell = max(sAhead, max(sLeft, sRight));

            if (maxSmell > SCENT_THRESHOLD) {
                if (sLeft > sAhead && sLeft > sRight) {
                    angle -= ANT_ROTATION_ANGLE;
                } else if (sRight > sAhead && sRight > sLeft) {
                    angle += ANT_ROTATION_ANGLE;
                }
                acted = true;
            }
        }

        // ---- PRIORITY 4: DIG (not carrying, diggable material nearby below surface) ----
        #if VIEW_MODE_SIDE
        if (!acted && isCarrying == 0.) {
            vec2 belowUv = roundUvToCellCenter(pos - vec2(0., cellSize));
            float belowMat2 = texture(tWorld, belowUv).x;
            vec4 belowProps2 = texelFetch(uMaterialProps, ivec2(int(belowMat2), 0), 0);

            // suppress digging if on the surface (air above) — don't dig topsoil into sky
            vec2 aboveUv = roundUvToCellCenter(pos + vec2(0., cellSize));
            float aboveMat2 = texture(tWorld, aboveUv).x;
            bool onSurfaceTop = (int(aboveMat2) == MAT_AIR);

            if (!onSurfaceTop
                && belowProps2.r == BEHAVIOR_POWDER
                && belowProps2.b <= ANT_CARRY_STRENGTH) {
                // bias angle toward ~40 degrees below horizontal (angle of repose)
                float targetAngle = (cos(angle) >= 0.)
                    ? -0.7  // ~-40 degrees (right and down)
                    : -(PI - 0.7);  // ~-(180-40) degrees (left and down)
                float angleDiff2 = targetAngle - angle;
                angleDiff2 = mod(angleDiff2 + PI, 2.0 * PI) - PI;
                angle += angleDiff2 * 0.2;
                acted = true;
            }
        }
        #endif

        // ---- PRIORITY 5: WANDER (fallback) ----
        if (!acted) {
            if (noise < 0.33) {
                angle += ANT_ROTATION_ANGLE;
            } else if (noise < 0.66) {
                angle -= ANT_ROTATION_ANGLE;
            }
            float noise2 = rand(vUv * 1000. + fract(uTime / 1000.) + 0.2);
            if (noise2 > 0.5) {
                angle += ANT_ROTATION_ANGLE * 2.;
            } else {
                angle -= ANT_ROTATION_ANGLE * 2.;
            }
        }

        // ---- MOVEMENT APPLICATION ----
        vec2 offset = vec2(cos(angle), sin(angle));
        pos = applyOffsetToPos(pos, offset);

        if (fract(pos.x) == 0. || fract(pos.y) == 0. || (!wasObstacle && isObstacle(pos + offset * cellSize))) {
            angle += PI * (noise - 0.5);
        }

        // ---- PICKUP (check cell ahead — ants walk in air, can't enter solid cells) ----
        if (isCarrying == 0.) {
            vec2 aheadUv = roundUvToCellCenter(pos + vec2(cos(angle), sin(angle)) * cellSize);
            float aheadMatId = texture(tWorld, aheadUv).x;
            int aheadMatInt = int(aheadMatId);

            // food: pick up from any adjacent cell
            if (aheadMatInt == MAT_FOOD) {
                isCarrying = 1.;
                cargoMaterialId = aheadMatId;
                angle += PI;
                storage = getMaxScentStorage(vUv);
            } else if (aheadMatInt != MAT_AIR && aheadMatInt != MAT_HOME) {
                // powder: don't grab the surface we're walking on — only dig
                // when facing diagonally into material (DIG priority angles us there)
                vec2 belowCell = roundUvToCellCenter(pos - vec2(0., cellSize));
                bool isWalkingSurface = all(equal(aheadUv, belowCell));
                if (!isWalkingSurface) {
                    vec4 props = texelFetch(uMaterialProps, ivec2(aheadMatInt, 0), 0);
                    float behavior = props.r;
                    float hardness = props.b;
                    if (behavior == BEHAVIOR_POWDER && hardness <= ANT_CARRY_STRENGTH) {
                        isCarrying = 1.;
                        cargoMaterialId = aheadMatId;
                        angle += PI;
                        storage = getMaxScentStorage(vUv);
                    }
                }
            }
        }
    } // end !isFalling

    FragColor = vec4(
        pos.x,
        pos.y,
        angle,
        float((uint(max(storage - SCENT_PER_MARKER, 0.)) << 1) + uint(isCarrying))
    );
    FragColorExt = vec4(personality, cargoMaterialId, pathIntDx, pathIntDy);
}