# 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.

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

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

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

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

@dataclass(frozen=True)
class Level:
    title: str = 'A'
    outlets: tuple = (ix(1, 0), ix(3, 0))
    fires: tuple = (ix(0, 2),)
    winds: tuple = ('E',) * 6
    stock: int = 4
    heights: tuple = tuple(4 - i // 5 for i in range(25))
    stones: tuple = ()

LEVELS = (Level(),)

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

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

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

def flow_to(level, i):
    """低い隣へ一歩。同じ最低高度なら下・左・右・上。"""
    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 == N and 0 <= xx < N:
            options.append((-1, rank, -1))  # 南端の排水口
        elif 0 <= xx < N and 0 <= yy < N:
            j = ix(xx, yy)
            if j not in level.stones and level.heights[j] < level.heights[i]:
                options.append((level.heights[j], rank, j))
    return min(options)[2] if options else i

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

def apply(state, level, action):
    if state.outcome != 'playing':
        return state, 'terminal'
    if action not in ACTIONS:
        return state, 'invalid'
    outlet, amount = action
    if amount > state.stock:
        return state, 'stock'
    water, fire, burned = list(state.water), list(state.fire), list(state.burned)
    steam, drain = state.steam, state.drain
    if outlet >= 0:
        water[level.outlets[outlet]] += amount
    old_fire = {i for i, n in enumerate(state.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(level, i)
            if j < 0:
                drain += n
            else:
                moved[j] += n
    water = moved
    for i in range(25):
        if water[i] and fire[i]:
            water[i] -= 1
            steam += 1
            fire[i] = 0
    dx, dy = DIRS[wind(level, state.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 < N and 0 <= yy < N:
            j = ix(xx, yy)
            if not fire[j] and not burned[j] and not water[j] and j not in level.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
    outcome = 'playing' if any(fire) else 'clear'
    return State(state.t + 1, state.stock - amount, tuple(water), tuple(fire),
                 tuple(burned), steam, drain, outcome), 'ok'

def invariant(state, level):
    assert state.stock + sum(state.water) + state.steam + state.drain == level.stock
    assert 0 <= state.stock <= level.stock
    assert state.steam >= 0 and state.drain >= 0
    assert all(0 <= n <= level.stock for n in state.water)
    assert all(n in (0, 1, 2) for n in state.fire)
    assert not any(state.fire[i] and state.burned[i] for i in range(25))


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

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

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

    
    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 ('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.state, self.message = apply(self.state, 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 3', 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(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   NO REAL-WORLD SAFETY ADVICE', 6)
        p.text(10, 328, 'R RESET  ESC QUIT  RELEASE KEYS BETWEEN TURNS', 6)

KEYS = ('LEFT', 'RIGHT', 'UP', 'DOWN', 'Z', 'SPACE', 'R')
DEMO_KEYS = ('Z', 'SPACE')

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 3', 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())