From 56d31408bf14749a211ea87835f99bbb0ec1caf1 Mon Sep 17 00:00:00 2001
From: nea <nea@nea.moe>
Date: Sun, 12 Mar 2023 00:25:33 +0100
Subject: PGN export

---
 index.html | 13 ++++++++++---
 pgn.py     | 36 ++++++++++++++++++++++++++++++++++++
 server.py  | 62 ++++++++++++++++++++++++++++++++++++++++++++++++--------------
 3 files changed, 94 insertions(+), 17 deletions(-)
 create mode 100644 pgn.py

diff --git a/index.html b/index.html
index 2a68cc2..0ed1c58 100644
--- a/index.html
+++ b/index.html
@@ -149,7 +149,7 @@
 <div id="choice" class="">
 </div>
 <div id="admodal">
-    <div><a href="#" onclick="adinfoClose()">x</a></div>
+    <div><a href="javascript:void(0)" onclick="adinfoClose()">x</a></div>
     <ul>
     <li>Created by Linnea Gräf</li>    
     <li><a href="https://git.nea.moe/nea/chess">Source</a></li>
@@ -162,6 +162,8 @@
 </div>
 <script>
 
+    const admodal = document.getElementById("admodal")
+
     function adinfo() {
         admodal.classList.add('visible');
     }
@@ -175,6 +177,7 @@
         constructor(elem) {
             this.socket = new WebSocket(`${location.protocol.includes('s') ? 'wss' : 'ws'}://${window.location.host}/socket`)
             this.ended = false
+            this.pgn = null
             this.boardState = {}
             this.lastMove = ''
             window.addEventListener('beforeunload', () => {
@@ -198,11 +201,11 @@
             this.legalMoves = []
             this.socket.addEventListener('message', ev => {
                 const message = JSON.parse(ev.data)
-                console.log(message)
                 this.playerColor = message.player_color || this.playerColor
                 this.boardState = parseFEN(message.board)
                 this.lastMove = message.lastmove || ''
                 this.ended ||= message.event === 'game_over'
+                this.pgn = message.pgn || this.pgn
                 this.result = message.result
                 this.legalMoves = message.legalmoves || []
                 this.awaitResync = false
@@ -267,7 +270,7 @@
             let uci = fromField + toField
             this.preTransformationMove = uci
             if (((toField[1] === '8' && this.playerColor === 'white')
-                    || (toField[2] === '1' && this.playerColor === 'black'))
+                    || (toField[1] === '1' && this.playerColor === 'black'))
                 && (this.boardState[fromField].toUpperCase() === 'P')) {
                 this.choiceButton.classList.add('selectingmove')
             } else {
@@ -300,6 +303,10 @@
                     <p><a href="https://youtu.be/yIRT6xRQkf8"><b>${this.result}</b></a></p>
 `
             }
+            if(this.pgn) {
+                this.turnIndicator.innerHTML += `
+                    <p><a href="data:application/vnd.chess-pgn;base64,${btoa(this.pgn)}" download="game.pgn" >Download PGN</a></p>`
+            }
             for (let field in this.fields) {
                 const fieldDOM = this.fields[field]
                 fieldDOM.innerHTML = ""
diff --git a/pgn.py b/pgn.py
new file mode 100644
index 0000000..ee9489b
--- /dev/null
+++ b/pgn.py
@@ -0,0 +1,36 @@
+import datetime
+import json
+
+
+class PgnWriter:
+
+    def __init__(self):
+        self.pgn_string = ""
+
+    def visit_tag(self, label: str, value: str):
+        self.pgn_string += "[" + label + " " + json.dumps(str(value)) + "]"
+
+    def visit_event(self, event_name: str):
+        self.visit_tag("Event", event_name)
+
+    def visit_site(self, site: str):
+        self.visit_tag("Site", site)
+
+    def visit_start_time(self, date: datetime.datetime):
+        self.visit_tag("Date", date.strftime('%Y.%m.%d'))
+        self.visit_tag("Time", date.strftime('%H:%M:%S'))
+
+    def visit_is_online(self, is_online: bool):
+        self.visit_tag("Mode", "ICS" if is_online else "OTB")
+
+    def visit_fen(self, fen: str):
+        self.visit_tag("FEN", fen)
+
+    def visit_round(self, round: int | str):
+        self.visit_tag("Round", str(round))
+
+    def visit_players(self, white: str, black: str):
+        self.visit_tag("White", white)
+        self.visit_tag("Black", black)
+
+
diff --git a/server.py b/server.py
index c7abb20..340a09a 100644
--- a/server.py
+++ b/server.py
@@ -1,20 +1,24 @@
 import asyncio
+import datetime
 import json
 import os
 import pathlib
 import typing
-import weakref
+from collections import namedtuple
 
 import aiohttp
 import aiohttp_jinja2
 import chess
 import chess.engine
+import chess.pgn
 import jinja2
 from aiohttp import web
 
+RunningGame = namedtuple('RunningGame', 'websocket pgn')
+
 app = web.Application()
-app['websockets'] = weakref.WeakSet()
-app['chessgames'] = weakref.WeakSet()
+app['games'] = set()
+app['game_count'] = 0
 
 basepath = pathlib.Path(__file__).parent.absolute()
 print(f"Loading templates from {basepath}")
@@ -29,8 +33,23 @@ async def index(request: web.Request):
 async def handle_socket(request: web.Request):
     ws = web.WebSocketResponse()
     await ws.prepare(request)
-    request.app['websockets'].add(ws)
+
+    pgn = chess.pgn.GameBuilder()
+    app['game_count'] += 1
+    chess.pgn.Game()
+    pgn.begin_game()
+    pgn.begin_headers()
+    pgn.visit_header("White", "Player")
+    pgn.visit_header("Black", "Giri, Anish")
+    pgn.visit_header("Date", datetime.date.today().strftime('%Y.%m.%d'))
+    pgn.visit_header("Round", "-")
+    pgn.visit_header("Site", "https://chess.nea.moe")
+    pgn.end_headers()
+    running_game = RunningGame(ws, pgn)
+
+    request.app['games'].add(running_game)
     engine: typing.Optional[chess.engine.UciProtocol] = None
+
     try:
         board = chess.Board()
         transport, engine = await chess.engine.popen_uci('stockfish')
@@ -42,6 +61,18 @@ async def handle_socket(request: web.Request):
 
         await send_to_user(dict(event="ready", player_color='white'))
 
+        async def check_outcome():
+            outcome = board.outcome(claim_draw=True)
+            if outcome:
+                pgn.visit_result(outcome.result())
+                pgn.visit_header("Termination", "normal")
+                pgn.end_game()
+                pgn_str = pgn.result().accept(chess.pgn.StringExporter())
+                await send_to_user(dict(event="pgn", pgn=pgn_str))
+                await send_to_user(dict(event="game_over", result=outcome.result()))
+
+            return bool(outcome)
+
         async for msg in ws:
             msg: aiohttp.WSMessage
             if msg.type == aiohttp.WSMsgType.TEXT:
@@ -57,6 +88,7 @@ async def handle_socket(request: web.Request):
                     if user_move not in board.legal_moves:
                         await send_to_user(dict(event="reject_move"))
                         continue
+                    pgn.visit_move(board, user_move)
                     board.push(user_move)
                     await send_to_user(dict(event="accept_move", lastmove=user_move.uci()))
                     candidates: typing.List[chess.engine.InfoDict] = await engine.analyse(
@@ -67,26 +99,24 @@ async def handle_socket(request: web.Request):
                         numscore = score.relative.score(mate_score=100000)
                         return abs(numscore)
 
-                    if board.is_game_over():
-                        await send_to_user(dict(event="game_over", result=board.result()))
+                    if await check_outcome():
                         break
 
                     most_drawy_move: chess.engine.InfoDict = min(candidates, key=appraise)
                     my_move: chess.Move = (most_drawy_move['pv'][0])
+                    pgn.visit_move(board, my_move)
                     board.push(my_move)
                     await send_to_user(dict(
                         event="computer_moved",
                         lastmove=my_move.uci(),
                     ))
-                    if board.is_game_over(claim_draw=True):
-                        await send_to_user(dict(event="game_over", result=board.result(claim_draw=True)))
+                    if await check_outcome():
                         break
-
     finally:
         if not ws.closed:
             await ws.close()
         print("Cleaning up websocket")
-        request.app['websockets'].discard(ws)
+        request.app['games'].discard(running_game)
         if engine is not None:
             asyncio.create_task(engine.quit())
     return ws
@@ -94,8 +124,12 @@ async def handle_socket(request: web.Request):
 
 async def on_shutdown(app):
     print("On shutdown called")
-    for ws in set(app['websockets']):
-        ws: web.WebSocketResponse
+    for game in set(app['games']):
+        game: RunningGame
+        ws: web.WebSocketResponse = game.websocket
+        game.pgn.end_game()
+        pgn_str = game.pgn.result().accept(chess.pgn.StringExporter())
+        await ws.send_json(dict(event="pgn", pgn=pgn_str))
         await ws.close(code=aiohttp.WSCloseCode.GOING_AWAY,
                        message=b'Server shutdown')
         print("Closing websocket")
@@ -104,8 +138,8 @@ async def on_shutdown(app):
 async def handle_status(request: web.Request):
     return web.json_response(dict(
         all_tasks=len(asyncio.all_tasks(asyncio.get_running_loop())),
-        websockets=len(request.app['websockets']),
-        games=len(request.app['chessgames']),
+        games=len(request.app['games']),
+        total_recent_games=app['game_count'],
     ))
 
 
-- 
cgit