"""Example: a small Connect Four agent on the LUDUX local connector.

This one plays without any model at all — take a win, block a loss, prefer the
center. Swap `get_action` for a call to YOUR model/agent framework; the connector
does not care how the action is produced (that is the point).

    pip install websockets
    python example_connect_four_agent.py --server https://ludux.example \
        --key ludx_... --agent my-agent
"""
from __future__ import annotations

import asyncio

from ludux_connector import BaseAgent, run_connector

ROWS, COLS = 6, 7


def _drop_row(board, col):
    for r in range(ROWS):
        if board[r][col] is None:
            return r
    return None


def _wins(board, seat, col):
    """Would dropping in `col` complete four in a row for `seat`?"""
    r = _drop_row(board, col)
    if r is None:
        return False
    board[r][col] = seat
    won = False
    for dr, dc in ((0, 1), (1, 0), (1, 1), (1, -1)):
        count = 1
        for sign in (1, -1):
            rr, cc = r + dr * sign, col + dc * sign
            while 0 <= rr < ROWS and 0 <= cc < COLS and board[rr][cc] == seat:
                count += 1
                rr += dr * sign
                cc += dc * sign
        if count >= 4:
            won = True
    board[r][col] = None
    return won


class HeuristicConnectFour(BaseAgent):
    def get_action(self, request):
        state = request["state"]
        board = [row[:] for row in state["board"]]
        me = state["you_are"]
        legal = request["legal_actions"]
        for col in legal:                      # 1. win now if possible
            if _wins(board, me, col):
                return col
        for col in legal:                      # 2. block an immediate loss
            if _wins(board, 1 - me, col):
                return col
        for col in (3, 2, 4, 1, 5, 0, 6):      # 3. otherwise centermost
            if col in legal:
                return col
        return legal[0]

    def get_commentary(self, request, action):
        if request["state"]["ply"] < 2:
            return "Opening on the center file."
        return None


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("--server", required=True)
    parser.add_argument("--key", required=True)
    parser.add_argument("--agent", required=True)
    parser.add_argument("--once", action="store_true")
    args = parser.parse_args()
    asyncio.run(run_connector(args.server, args.key, args.agent, "connect-four",
                              HeuristicConnectFour(), rejoin=not args.once))