# HG changeset patch
# User Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
# Date 1344177446 -7200
# Node ID 6deab6ad8be8d42fbf538f1de9763fbe1f2324d9
# Parent  704bea2e43605a5d137b43b5d923744ee9307245
Add the ability to save a replay.

diff --git a/eosd b/eosd
--- a/eosd
+++ b/eosd
@@ -27,7 +27,7 @@ from pytouhou.ui.gamerunner import GameR
 from pytouhou.games.eosd import EoSDGame
 from pytouhou.game.game import GameOver
 from pytouhou.game.player import PlayerState
-from pytouhou.formats.t6rp import T6RP
+from pytouhou.formats.t6rp import T6RP, Level
 from pytouhou.utils.random import Random
 from pytouhou.vm.msgrunner import NextStage
 
@@ -62,7 +62,9 @@ class EoSDGameBossRush(EoSDGame):
 
 
 
-def main(path, data, stage_num, rank, character, replay, boss_rush, fps_limit, single_buffer, debug, fixed_pipeline):
+def main(path, data, stage_num, rank, character, replay, save_filename,
+         boss_rush, fps_limit, single_buffer, debug, fixed_pipeline):
+
     resource_loader = Loader(path)
 
     try:
@@ -89,6 +91,12 @@ def main(path, data, stage_num, rank, ch
         rank = replay.rank
         character = replay.character
 
+    save_keystates = None
+    if save_filename:
+        save_replay = T6RP()
+        save_replay.rank = rank
+        save_replay.character = character
+
     difficulty = 16
     default_power = [0, 64, 128, 128, 128, 128, 0][stage_num - 1]
     states = [PlayerState(character=character, power=default_power)]
@@ -110,12 +118,25 @@ def main(path, data, stage_num, rank, ch
                 previous_level = replay.levels[stage_num - 1]
                 states[0].score = previous_level.score
                 states[0].effective_score = previous_level.score
+            states[0].point_items = level.point_items
             states[0].power = level.power
             states[0].lives = level.lives
             states[0].bombs = level.bombs
             difficulty = level.difficulty
         else:
-            prng = None
+            prng = Random()
+
+        if save_filename:
+            if not replay:
+                save_replay.levels[stage_num - 1] = level = Level()
+                level.score = states[0].score
+                level.random_seed = prng.seed
+                level.point_items = states[0].points
+                level.power = states[0].power
+                level.lives = states[0].lives
+                level.bombs = states[0].bombs
+                level.difficulty = difficulty
+            save_keystates = []
 
         # Load stage data
         stage = resource_loader.get_stage('stage%d.std' % stage_num)
@@ -126,7 +147,7 @@ def main(path, data, stage_num, rank, ch
         background = Background(stage, background_anm_wrapper)
 
         # Main loop
-        runner.load_game(game, background, stage.bgms, replay)
+        runner.load_game(game, background, stage.bgms, replay, save_keystates)
         try:
             runner.start()
             break
@@ -139,6 +160,17 @@ def main(path, data, stage_num, rank, ch
         except GameOver:
             print('Game over')
             break
+        finally:
+            if save_filename:
+                last_key = -1
+                for time, key in enumerate(save_keystates):
+                    if key != last_key:
+                        level.keys.append((time, key, 0))
+                    last_key = key
+
+    if save_filename:
+        with open(save_filename, 'wb+') as file:
+            save_replay.write(file)
 
 
 pathsep = os.path.pathsep
@@ -157,6 +189,7 @@ parser.add_argument('-s', '--stage', met
 parser.add_argument('-r', '--rank', metavar='RANK', type=int, default=0, help='Rank, from 0 (Easy, default) to 3 (Lunatic).')
 parser.add_argument('-c', '--character', metavar='CHARACTER', type=int, default=0, help='Select the character to use, from 0 (ReimuA, default) to 3 (MarisaB).')
 parser.add_argument('--replay', metavar='REPLAY', help='Select a replay')
+parser.add_argument('--save-replay', metavar='REPLAY', help='Save the upcoming game into a replay file')
 parser.add_argument('-b', '--boss-rush', action='store_true', help='Fight only bosses')
 parser.add_argument('--single-buffer', action='store_true', help='Disable double buffering')
 parser.add_argument('--fps-limit', metavar='FPS', default=60, type=int, help='Set fps limit')
@@ -166,5 +199,5 @@ parser.add_argument('--fixed-pipeline', 
 args = parser.parse_args()
 
 main(args.path, tuple(args.data), args.stage, args.rank, args.character,
-     args.replay, args.boss_rush, args.fps_limit, args.single_buffer,
-     args.debug, args.fixed_pipeline)
+     args.replay, args.save_replay, args.boss_rush, args.fps_limit,
+     args.single_buffer, args.debug, args.fixed_pipeline)
diff --git a/pytouhou/formats/t6rp.py b/pytouhou/formats/t6rp.py
--- a/pytouhou/formats/t6rp.py
+++ b/pytouhou/formats/t6rp.py
@@ -20,8 +20,9 @@ a game of EoSD. Since the EoSD engine is
 replay file is sufficient to unfold a full game.
 """
 
-from struct import unpack
+from struct import unpack, pack
 from io import BytesIO
+from time import strftime
 
 from pytouhou.utils.helpers import read_string, get_logger
 
@@ -47,7 +48,17 @@ class T6RP(object):
         self.version = 0x102
         self.character = 0
         self.rank = 0
+        self.unknown1 = 0
+        self.unknown2 = 0
         self.key = 0
+        self.unknown3 = 0
+        self.date = strftime('%d/%m/%y')
+        self.name = 'PyTouhou'
+        self.unknown4 = 0
+        self.score = 0
+        self.unknown5 = 0
+        self.slowdown = 0.
+        self.unknown6 = 0
         self.levels = [None] * 7
 
 
@@ -84,8 +95,9 @@ class T6RP(object):
         if verify:
             data = file.read()
             file.seek(15)
-            if checksum != (sum(ord(c) for c in data) + 0x3f000318 + replay.key) & 0xffffffff:
-                raise Exception #TODO
+            real_sum = (sum(ord(c) for c in data) + 0x3f000318 + replay.key) & 0xffffffff
+            if checksum != real_sum:
+                raise Exception('Checksum mismatch: %d ≠ %d.' % (checksum, real_sum))
 
         replay.unknown3, = unpack('<B', file.read(1))
         replay.date = file.read(9) #read_string(file, 9, 'ascii')
@@ -114,3 +126,64 @@ class T6RP(object):
                 level.keys.append((time, keys, unknown))
 
         return replay
+
+
+    def write(self, file, encrypt=True):
+        if encrypt:
+            encrypted_file = file
+            file = BytesIO()
+
+        file.write(b'T6RP')
+        file.write(pack('<HBB', self.version, self.character, self.rank))
+
+        checksum_offset = file.tell()
+        file.seek(4, 1) # For checksum
+        file.write(pack('<BBB', self.unknown1, self.unknown2, self.key))
+
+        file.write(pack('<B', self.unknown3))
+
+        #TODO: find a more elegant method.
+        n = 9 - len(self.date)
+        file.write(self.date)
+        file.write('\0' * n)
+        n = 9 - len(self.name)
+        file.write(self.name)
+        file.write('\0' * n)
+
+        file.write(pack('<HIIfI', self.unknown4, self.score, self.unknown5, self.slowdown, self.unknown6))
+
+        stages_offsets_offset = file.tell()
+        file.seek(7*4, 1) # Skip the stages offsets.
+
+        stages_offsets = []
+        for level in self.levels:
+            if not level:
+                stages_offsets.append(0)
+                continue
+
+            stages_offsets.append(file.tell())
+            file.write(pack('<IHHBbbBI', level.score, level.random_seed,
+                            level.point_items, level.power, level.lives,
+                            level.bombs, level.difficulty, level.unknown))
+
+            for time, keys, unknown in level.keys:
+                file.write(pack('<IHH', time, keys, unknown))
+
+            file.write(pack('<IHH', 9999999, 0, 0))
+
+        file.seek(stages_offsets_offset)
+        file.write(pack('<7I', *stages_offsets))
+
+        # Write checksum
+        file.seek(15)
+        data = file.read()
+        checksum = (sum(ord(c) for c in data) + 0x3f000318 + self.key) & 0xffffffff
+        file.seek(checksum_offset)
+        file.write(pack('<I', checksum))
+
+        # Encrypt
+        if encrypt:
+            file.seek(0)
+            encrypted_file.write(file.read(15))
+            encrypted_file.write(b''.join(chr((ord(c) + self.key + 7*i) & 0xff) for i, c in enumerate(file.read())))
+
diff --git a/pytouhou/game/game.py b/pytouhou/game/game.py
--- a/pytouhou/game/game.py
+++ b/pytouhou/game/game.py
@@ -14,8 +14,6 @@
 
 from itertools import chain
 
-from pytouhou.utils.random import Random
-
 from pytouhou.vm.eclrunner import ECLMainRunner
 from pytouhou.vm.msgrunner import MSGRunner
 
@@ -72,7 +70,7 @@ class Game(object):
         self.msg_wait = False
         self.bonus_list = [0, 0, 1, 0, 1, 0, 0, 1, 1, 1, 0, 0, 0, 1, 1, 0,
                            1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 0, 1, 1, 1, 0, 2]
-        self.prng = prng or Random()
+        self.prng = prng
         self.frame = 0
 
         self.enm_anm_wrapper = resource_loader.get_anm_wrapper2(('stg%denm.anm' % stage,
diff --git a/pytouhou/ui/gamerunner.py b/pytouhou/ui/gamerunner.py
--- a/pytouhou/ui/gamerunner.py
+++ b/pytouhou/ui/gamerunner.py
@@ -61,7 +61,7 @@ class GameRunner(pyglet.window.Window, G
         self.clock = pyglet.clock.get_default()
 
 
-    def load_game(self, game=None, background=None, bgms=None, replay=None):
+    def load_game(self, game=None, background=None, bgms=None, replay=None, save_keystates=None):
         GameRenderer.load_game(self, game, background)
         self.replay_level = None
         if not replay or not replay.levels[game.stage-1]:
@@ -75,6 +75,8 @@ class GameRunner(pyglet.window.Window, G
             game.players[0].state.bombs = self.replay_level.bombs
             game.difficulty = self.replay_level.difficulty
 
+        self.save_keystates = save_keystates
+
         game.music = MusicPlayer(game.resource_loader, bgms)
         game.music.play(0)
 
@@ -166,6 +168,9 @@ class GameRunner(pyglet.window.Window, G
                     else:
                         keystate = _keystate
 
+            if self.save_keystates is not None:
+                self.save_keystates.append(keystate)
+
             self.game.run_iter(keystate)