#include "CoopAIController.h"

#include "CoopGame.h"

#include <algorithm>
#include <array>
#include <cmath>
#include <limits>

namespace {

static bool canPlacePieceForSide(const std::array<CoopGame::Cell, CoopGame::COLS * CoopGame::ROWS>& board,
                                 const CoopGame::Piece& p,
                                 CoopGame::PlayerSide side) {
    for (int cy = 0; cy < 4; ++cy) {
        for (int cx = 0; cx < 4; ++cx) {
            if (!CoopGame::cellFilled(p, cx, cy)) {
                continue;
            }

            const int bx = p.x + cx;
            const int by = p.y + cy;

            // Keep the AI strictly in the correct half.
            if (side == CoopGame::PlayerSide::Right) {
                if (bx < 10 || bx >= CoopGame::COLS) {
                    return false;
                }
            } else {
                if (bx < 0 || bx >= 10) {
                    return false;
                }
            }

            // Above the visible board is allowed.
            if (by < 0) {
                continue;
            }

            if (by >= CoopGame::ROWS) {
                return false;
            }

            if (board[by * CoopGame::COLS + bx].occupied) {
                return false;
            }
        }
    }
    return true;
}

static int dropYFor(const std::array<CoopGame::Cell, CoopGame::COLS * CoopGame::ROWS>& board,
                    CoopGame::Piece p,
                    CoopGame::PlayerSide side) {
    // Assumes p is currently placeable.
    while (true) {
        CoopGame::Piece next = p;
        next.y += 1;
        if (!canPlacePieceForSide(board, next, side)) {
            return p.y;
        }
        p = next;
        if (p.y > CoopGame::ROWS) {
            return p.y;
        }
    }
}

static void applyPiece(std::array<uint8_t, CoopGame::COLS * CoopGame::ROWS>& occ,
                      const CoopGame::Piece& p) {
    for (int cy = 0; cy < 4; ++cy) {
        for (int cx = 0; cx < 4; ++cx) {
            if (!CoopGame::cellFilled(p, cx, cy)) {
                continue;
            }
            const int bx = p.x + cx;
            const int by = p.y + cy;
            if (by < 0 || by >= CoopGame::ROWS || bx < 0 || bx >= CoopGame::COLS) {
                continue;
            }
            occ[by * CoopGame::COLS + bx] = 1;
        }
    }
}

struct Eval {
    double score = -std::numeric_limits<double>::infinity();
    int rot = 0;
    int x = 10;
};

static Eval evaluateBestPlacementForSide(const CoopGame& game, CoopGame::PlayerSide side) {
    const auto& board = game.boardRef();

    std::array<uint8_t, CoopGame::COLS * CoopGame::ROWS> occ{};
    for (int i = 0; i < CoopGame::COLS * CoopGame::ROWS; ++i) {
        occ[i] = board[i].occupied ? 1 : 0;
    }

    const CoopGame::Piece cur = game.current(side);

    Eval best{};

    // Iterate rotations and x positions. IMPORTANT: allow x to go slightly out of bounds
    // because our pieces are represented in a 4x4 mask and many rotations have leading
    // empty columns. For example, placing a vertical I/J/L into column 0 often requires
    // p.x == -1 or p.x == -2 so the filled cells land at bx==0.
    // canPlacePieceForSide() enforces the actual half-board bounds.
    for (int rot = 0; rot < 4; ++rot) {
        int xmin = (side == CoopGame::PlayerSide::Right) ? 6 : -3;
        int xmax = (side == CoopGame::PlayerSide::Right) ? 22 : 13;
        for (int x = xmin; x <= xmax; ++x) {
            CoopGame::Piece p = cur;
            p.rot = rot;
            p.x = x;

            // If this rotation/x is illegal at the current y, try near the top spawn band.
            if (!canPlacePieceForSide(board, p, side)) {
                p.y = -2;
                if (!canPlacePieceForSide(board, p, side)) {
                    continue;
                }
            }

            p.y = dropYFor(board, p, side);

            auto occ2 = occ;
            applyPiece(occ2, p);

            // Count completed full rows (all 20 cols) after placement.
            int fullRows = 0;
            for (int y = 0; y < CoopGame::ROWS; ++y) {
                bool full = true;
                for (int cx = 0; cx < CoopGame::COLS; ++cx) {
                    if (!occ2[y * CoopGame::COLS + cx]) {
                        full = false;
                        break;
                    }
                }
                if (full) {
                    ++fullRows;
                }
            }

            // Right-half column heights + holes + bumpiness.
            std::array<int, 10> heights{};
            int aggregateHeight = 0;
            int holes = 0;

            for (int c = 0; c < 10; ++c) {
                const int bx = (side == CoopGame::PlayerSide::Right) ? (10 + c) : c;
                int h = 0;
                bool found = false;
                for (int y = 0; y < CoopGame::ROWS; ++y) {
                    if (occ2[y * CoopGame::COLS + bx]) {
                        h = CoopGame::ROWS - y;
                        found = true;
                        // Count holes below the first filled cell.
                        for (int yy = y + 1; yy < CoopGame::ROWS; ++yy) {
                            if (!occ2[yy * CoopGame::COLS + bx]) {
                                ++holes;
                            }
                        }
                        break;
                    }
                }
                heights[c] = found ? h : 0;
                aggregateHeight += heights[c];
            }

            int bump = 0;
            for (int i = 0; i < 9; ++i) {
                bump += std::abs(heights[i] - heights[i + 1]);
            }

            // Reward sync potential: rows where the right half is full (10..19).
            int sideHalfFullRows = 0;
            for (int y = 0; y < CoopGame::ROWS; ++y) {
                bool full = true;
                int start = (side == CoopGame::PlayerSide::Right) ? 10 : 0;
                int end = (side == CoopGame::PlayerSide::Right) ? 20 : 10;
                for (int bx = start; bx < end; ++bx) {
                    if (!occ2[y * CoopGame::COLS + bx]) {
                        full = false;
                        break;
                    }
                }
                if (full) {
                    ++sideHalfFullRows;
                }
            }

            // Simple heuristic:
            // - Strongly prefer completed full rows
            // - Prefer making the right half complete (helps cooperative clears)
            // - Penalize holes and excessive height/bumpiness
            double s = 0.0;
            // Strongly prefer full-line clears across the whole board (rare but best).
            s += static_cast<double>(fullRows) * 12000.0;
            // Heavily prefer completing the player's half — make this a primary objective.
            s += static_cast<double>(sideHalfFullRows) * 6000.0;
            // Penalize holes and height less aggressively so completing half-rows is prioritized.
            s -= static_cast<double>(holes) * 180.0;
            s -= static_cast<double>(aggregateHeight) * 4.0;
            s -= static_cast<double>(bump) * 10.0;

            // Reduce center bias so edge placements to complete rows are not punished.
            double centerTarget = (side == CoopGame::PlayerSide::Right) ? 15.0 : 4.5;
            const double centerBias = -std::abs((x + 1.5) - centerTarget) * 1.0;
            s += centerBias;

            if (s > best.score) {
                best.score = s;
                best.rot = rot;
                best.x = x;
            }
        }
    }

    return best;
}

} // namespace

void CoopAIController::reset() {
    m_lastPieceSeq = 0;
    m_hasPlan = false;
    m_targetRot = 0;
    m_targetX = 10;
    m_moveTimerMs = 0.0;
    m_moveDir = 0;
    m_rotateTimerMs = 0.0;
}

void CoopAIController::computePlan(const CoopGame& game, CoopGame::PlayerSide side) {
    const Eval best = evaluateBestPlacementForSide(game, side);
    m_targetRot = best.rot;
    m_targetX = best.x;
    m_hasPlan = true;
    m_moveTimerMs = 0.0;
    m_moveDir = 0;
    m_rotateTimerMs = 0.0;
}

void CoopAIController::update(CoopGame& game, CoopGame::PlayerSide side, double frameMs) {
    const uint64_t seq = game.currentPieceSequence(side);
    if (seq != m_lastPieceSeq) {
        m_lastPieceSeq = seq;
        m_hasPlan = false;
        m_moveTimerMs = 0.0;
        m_moveDir = 0;
        m_rotateTimerMs = 0.0;
    }

    if (!m_hasPlan) {
        computePlan(game, side);
    }

    const CoopGame::Piece cur = game.current(side);

    // Clamp negative deltas (defensive; callers should pass >= 0).
    const double dt = std::max(0.0, frameMs);

    // Update timers.
    if (m_moveTimerMs > 0.0) {
        m_moveTimerMs -= dt;
        if (m_moveTimerMs < 0.0) m_moveTimerMs = 0.0;
    }
    if (m_rotateTimerMs > 0.0) {
        m_rotateTimerMs -= dt;
        if (m_rotateTimerMs < 0.0) m_rotateTimerMs = 0.0;
    }

    // Rotate toward target first.
    const int curRot = ((cur.rot % 4) + 4) % 4;
    const int tgtRot = ((m_targetRot % 4) + 4) % 4;
    int diff = (tgtRot - curRot + 4) % 4;
    if (diff != 0) {
        // Human-ish rotation rate limiting.
        if (m_rotateTimerMs <= 0.0) {
            const int dir = (diff == 3) ? -1 : 1;
            game.rotate(side, dir);
            m_rotateTimerMs = m_rotateIntervalMs;
        }
        // While rotating, do not also slide horizontally in the same frame.
        m_moveDir = 0;
        m_moveTimerMs = 0.0;
        return;
    }

    // Move horizontally toward target.
    int desiredDir = 0;
    if (cur.x < m_targetX) desiredDir = +1;
    else if (cur.x > m_targetX) desiredDir = -1;

    if (desiredDir == 0) {
        // Aligned: do nothing. Gravity controls fall speed (no AI hard drops).
        m_moveDir = 0;
        m_moveTimerMs = 0.0;
        return;
    }

    // DAS/ARR-style horizontal movement pacing.
    if (m_moveDir != desiredDir) {
        // New direction / initial press: move immediately, then wait DAS.
        game.move(side, desiredDir);
        m_moveDir = desiredDir;
        m_moveTimerMs = m_dasMs;
        return;
    }

    // Holding direction: repeat every ARR once DAS has elapsed.
    if (m_moveTimerMs <= 0.0) {
        game.move(side, desiredDir);
        m_moveTimerMs = m_arrMs;
    }
}