# MIT License
#
# Copyright (c) 2026 OpenAI-assisted game contributors
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
#
# This MIT license covers original source, documentation and primitive artwork.
# Bitmap glyph data is derived from Noto Sans CJK and remains under SIL Open
# Font License 1.1. See assets/NOTO-LICENSE.txt for the separate font license.

"""消防水路・教材5/5。架空の格子規則。現実の消火手順ではありません。"""
from dataclasses import dataclass, replace

WAIT = (-1, 0)
ACTIONS = (WAIT, (0, 1), (0, 2), (1, 1), (1, 2))

N = 5
DIRS = {'E': (1, 0), 'W': (-1, 0), 'S': (0, 1), 'N': (0, -1)}

def ix(x, y):
    return y * N + x

def xy(i):
    return (i % N, i // N)

@dataclass(frozen=True)
class Level:
    title: str
    outlets: tuple
    fires: tuple
    houses: tuple
    winds: tuple
    goal: int
    stock: int = 4
    heights: tuple = tuple((4 - i // 5 for i in range(25)))
    stones: tuple = ()
LEVELS = (Level('下へ残る水', (ix(1, 0), ix(3, 0)), (ix(0, 2),), (ix(2, 2), ix(4, 2), ix(3, 4)), ('E',) * 6, 3), Level('近い火は一手待てる', (ix(1, 0), ix(3, 1)), (ix(0, 0), ix(2, 3)), (ix(4, 3), ix(4, 4), ix(2, 2)), ('E',) * 6, 3), Level('曲がる風', (ix(1, 0), ix(3, 2)), (ix(0, 1), ix(0, 3)), (ix(3, 3), ix(2, 4), ix(4, 4)), ('E', 'E', 'S', 'E', 'E', 'E'), 2))

@dataclass(frozen=True)
class State:
    t: int
    stock: int
    water: tuple
    wet: tuple
    fire: tuple
    burned: tuple
    damaged: tuple
    steam: int = 0
    drain: int = 0
    outcome: str = 'playing'

def initial(l):
    fire = [0] * 25
    for i in l.fires:
        fire[i] = 2
    return State(0, l.stock, (0,) * 25, (0,) * 25, tuple(fire), (False,) * 25, tuple((i in l.fires for i in l.houses)))

def flow_to(l, i):
    """Strictly lower neighbor. Equal lowest: down,left,right,up. South edge drain."""
    x, y = xy(i)
    options = []
    for rank, (dx, dy) in enumerate(((0, 1), (-1, 0), (1, 0), (0, -1))):
        xx, yy = (x + dx, y + dy)
        if yy == 5 and 0 <= xx < 5:
            options.append((-1, rank, -1))
        elif 0 <= xx < 5 and 0 <= yy < 5:
            j = ix(xx, yy)
            if j not in l.stones and l.heights[j] < l.heights[i]:
                options.append((l.heights[j], rank, j))
    return min(options)[2] if options else i

def wind(l, t):
    return l.winds[min(t, len(l.winds) - 1)]

def saved(s, l):
    return len(l.houses) - sum(s.damaged)

def apply(s, l, a, wet_life=2):
    if s.outcome != 'playing':
        return (s, 'terminal')
    if a not in ((-1, 0), (0, 1), (0, 2), (1, 1), (1, 2)):
        return (s, 'invalid')
    outlet, amount = a
    if amount > s.stock:
        return (s, 'stock')
    water = list(s.water)
    wet = [max(0, n - 1) for n in s.wet]
    fire = list(s.fire)
    burned = list(s.burned)
    damaged = list(s.damaged)
    steam = s.steam
    drain = s.drain
    if outlet >= 0:
        water[l.outlets[outlet]] += amount
    touched = {i for i, n in enumerate(water) if n}
    old_fire = {i for i, n in enumerate(s.fire) if n}
    for i in range(25):
        if water[i] and fire[i]:
            water[i] -= 1
            steam += 1
            fire[i] = 0
    moved = [0] * 25
    for i, n in enumerate(water):
        if n:
            j = flow_to(l, i)
            if j < 0:
                drain += n
            else:
                moved[j] += n
                touched.add(j)
    water = moved
    for i in range(25):
        if water[i] and fire[i]:
            water[i] -= 1
            steam += 1
            fire[i] = 0
    for i in touched:
        wet[i] = wet_life
    dx, dy = DIRS[wind(l, s.t)]
    ignitions = set()
    for i in old_fire:
        if not fire[i]:
            continue
        x, y = xy(i)
        xx, yy = (x + dx, y + dy)
        if 0 <= xx < 5 and 0 <= yy < 5:
            j = ix(xx, yy)
            if not fire[j] and (not burned[j]) and (not wet[j]) and (not water[j]) and (j not in l.stones):
                ignitions.add(j)
    for i in old_fire:
        if fire[i]:
            fire[i] -= 1
            if not fire[i]:
                burned[i] = True
    for j in ignitions:
        fire[j] = 2
        if j in l.houses:
            damaged[l.houses.index(j)] = True
    outcome = 'playing' if any(fire) else 'won' if len(l.houses) - sum(damaged) >= l.goal else 'lost'
    ns = State(s.t + 1, s.stock - amount, tuple(water), tuple(wet), tuple(fire), tuple(burned), tuple(damaged), steam, drain, outcome)
    return (ns, 'ok')

def forecast(s, l, a, wet_life=2):
    return apply(s, l, a, wet_life)

def invariant(s, l, wet_life=2):
    assert s.stock + sum(s.water) + s.steam + s.drain == l.stock
    assert 0 <= s.stock <= l.stock and s.steam >= 0 and (s.drain >= 0)
    assert all((0 <= n <= l.stock for n in s.water))
    assert all((0 <= n <= wet_life for n in s.wet))
    assert all((n in (0, 1, 2) for n in s.fire))
    assert not any((s.fire[i] and s.burned[i] for i in range(25)))
    assert len(s.damaged) == len(l.houses)

@dataclass(frozen=True)
class Timeline:
    state: State
    history: tuple = ()

def step(timeline, level, action, wet_life=2):
    if action == ('undo',):
        if not timeline.history:
            return timeline, 'no_undo'
        return Timeline(timeline.history[-1], timeline.history[:-1]), 'undo'
    state, reason = apply(timeline.state, level, action, wet_life)
    if state == timeline.state:
        return timeline, reason
    return Timeline(state, timeline.history + (timeline.state,)), reason


class App:
    def __init__(self):
        self.level_index = 0
        self.reset()

    @property
    def level(self):
        return LEVELS[self.level_index]

    def reset(self):
        self.timeline = Timeline(initial(self.level))
        self.outlet = 0
        self.amount = 1
        self.gated = False
        self.message = 'READY'

    @property
    def state(self):
        return self.timeline.state

    def control(self, pressed, held):
        if 'R' in pressed:
            self.reset()
            self.gated = True
            return
        if self.gated:
            if not pressed and not held:
                self.gated = False
            return
        if len(pressed | held) > 1:
            self.gated = True
            self.message = 'RELEASE ALL KEYS'
            return
        if len(pressed) != 1:
            return
        key = next(iter(pressed))
        if key in ('1', '2', '3'):
            self.level_index = int(key)-1
            self.reset()
            self.gated = True
            return
        if key == 'X':
            self.timeline, self.message = step(self.timeline, self.level, ('undo',))
            self.gated = True
            return
        if key in ('LEFT', 'RIGHT', 'UP', 'DOWN'):
            if key == 'LEFT':self.outlet = 0
            if key == 'RIGHT':self.outlet = 1
            if key == 'UP':self.amount = 2
            if key == 'DOWN':self.amount = 1
            self.message = 'SELECTED'
            self.gated = True
            return
        if key == 'SPACE':
            action = WAIT
        elif key == 'Z':
            action = (self.outlet, self.amount)
        else:
            return
        self.timeline, self.message = step(self.timeline, self.level, action)
        self.gated = True

    def draw(self, p):
        p.cls(1)
        s, level = self.state, self.level
        p.text(10, 8, 'FIRE WATERWAYS / STAGE 5', 10)
        outcome = getattr(s, 'outcome', 'practice').upper()
        p.text(10, 22, 'T '+str(s.t)+'  LEVEL '+chr(65+self.level_index)+'  '+outcome, 7)
        p.text(10, 36, 'FICTIONAL GRID RULES ONLY', 6)
        for i in range(25):
            x, y = xy(i)
            px, py = 12+x*40, 62+y*40
            p.rect(px, py, 38, 38, 5)
            if i in getattr(level, 'stones', ()):
                p.rect(px+2, py+2, 34, 34, 13)
            if getattr(s, 'burned', (False,)*25)[i]:
                p.rect(px+2, py+2, 34, 34, 0)
            wet = getattr(s, 'wet', (0,)*25)[i]
            if wet:
                for dx in (8,18,28):p.pset(px+dx, py+32, 12)
            water = s.water[i]
            if water:
                p.rect(px+3, py+17, 31, 12, 12)
                p.text(px+5, py+20, 'W'+str(water), 7)
            fire = getattr(s, 'fire', (0,)*25)[i]
            if fire:
                p.tri(px+18, py+8, px+7, py+28, px+30, py+28, 9)
                p.text(px+15, py+20, str(fire), 0)
            houses = getattr(level, 'houses', ())
            if i in houses:
                damage = s.damaged[houses.index(i)]
                p.rectb(px+24, py+3, 11, 10, 8 if damage else 11)
                p.text(px+27, py+5, 'H', 8 if damage else 11)
            if hasattr(level, 'heights'):
                p.text(px+2, py+2, str(level.heights[i]), 6)
            if i in level.outlets:
                label = 'W' if level.outlets.index(i)==0 else 'E'
                p.text(px+14, py+4, label, 10)
        p.text(12, 266, 'SOUTH EDGE -> DRAIN', 6)
        p.text(232, 62, 'WATER LEDGER', 10)
        ledger = (('STOCK', getattr(s, 'stock', 0)), ('GRID', sum(s.water)),
                  ('STEAM', getattr(s, 'steam', 0)), ('DRAIN', s.drain))
        for row, (label, value) in enumerate(ledger):
            p.text(232, 78+row*14, label+' '+str(value), 7)
        p.text(232, 136, 'TOTAL '+str(sum(v for _, v in ledger))+' / '+str(level.stock), 11)
        p.text(232, 152, 'WET DOTS ARE NOT WATER', 6)
        p.text(232, 176, 'SELECT '+('W' if self.outlet==0 else 'E')+' / '+str(self.amount), 10)
        after, reason = apply(s, level, (self.outlet, self.amount))
        p.text(232, 192, 'Z -> T'+str(after.t)+' STOCK '+str(after.stock), 7)
        p.text(232, 206, 'PREVIEW '+reason.upper(), 6)
        p.text(232, 222, 'WIND '+wind(level,s.t)+'  FIRE '+str(sum(v>0 for v in s.fire)), 9)
        p.text(232, 238, 'SAVED '+str(saved(s,level))+'/'+str(len(level.houses))+'  GOAL '+str(level.goal), 11)
        p.text(232, 254, 'UNDO '+str(len(self.timeline.history)), 6)
        p.text(10, 284, self.message.upper(), 7)
        p.text(10, 300, 'LEFT/RIGHT OUTLET   UP=2 DOWN=1   Z RELEASE', 6)
        p.text(10, 314, 'SPACE WAIT   1/2/3 LEVEL   X UNDO', 6)
        p.text(10, 328, 'R RESET  ESC QUIT  RELEASE KEYS BETWEEN TURNS', 6)

KEYS = ('LEFT', 'RIGHT', 'UP', 'DOWN', 'Z', 'SPACE', 'R', 'X', '1', '2', '3')
DEMO_KEYS = ('2', 'RIGHT', 'Z', 'SPACE', 'SPACE', 'SPACE', 'SPACE', 'SPACE', 'X')

def parse_args(argv=None):
    import argparse
    parser = argparse.ArgumentParser(description='Five-stage fictional grid tutorial')
    parser.add_argument('--headless-smoke', action='store_true')
    parser.add_argument('--demo-checkpoint', action='store_true')
    parser.add_argument('--screenshot')
    args = parser.parse_args(argv)
    if (args.demo_checkpoint or args.screenshot is not None) and not args.headless_smoke:
        parser.error('--demo-checkpoint/--screenshot require --headless-smoke')
    if args.screenshot is not None and not args.screenshot.lower().endswith('.png'):
        parser.error('--screenshot requires a .png path')
    return args

def main(args):
    from pathlib import Path
    import pyxel as p  # 引数を検証した後で初めて読み込む
    screenshot = str(Path(args.screenshot).resolve()) if args.screenshot else None
    if screenshot:
        Path(screenshot).parent.mkdir(parents=True, exist_ok=True)
    app = App()
    p.init(480, 352, title='Fire Waterways 5', fps=30,
           display_scale=2, headless=args.headless_smoke)
    keys = {k:getattr(p, 'KEY_'+k) for k in KEYS}
    limit = max(12, 2*len(DEMO_KEYS)+3) if args.demo_checkpoint else 12
    def update():
        if args.headless_smoke and p.frame_count >= limit:
            if screenshot:
                p.screenshot(screenshot, scale=2)
            p.quit()
            return
        if args.headless_smoke and args.demo_checkpoint:
            index = p.frame_count // 2
            pressed = {DEMO_KEYS[index]} if p.frame_count%2==0 and index<len(DEMO_KEYS) else set()
            app.control(pressed, pressed)
        else:
            app.control({k for k,v in keys.items() if p.btnp(v)},
                        {k for k,v in keys.items() if p.btn(v)})
    p.run(update, lambda:app.draw(p))

if __name__ == '__main__':
    main(parse_args())