# MIT License
# 
# Copyright (c) 2026
# 
# 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 applies to original game source and primitive artwork. Bitmap glyph data
# is derived from Noto Sans CJK under SIL Open Font License 1.1, and is covered
# by assets/NOTO-LICENSE.txt instead of MIT.

"""蜜蜂の引っ越し・教材5/5。解説はLESSON.mdを参照。"""
"""Immutable bee ecology; reservations are commitments, not duplicated nectar."""
from dataclasses import dataclass, replace, asdict
import hashlib, json

@dataclass(frozen=True)
class FlowerDef:
    x: int
    stock: int
    awake: bool

@dataclass(frozen=True)
class Level:
    title: str
    flowers: tuple
    links: tuple
    target: int

@dataclass(frozen=True)
class Flower:
    stock: int
    reserved: int = 0
    phase: str = 'sleep'
    timer: int = 0

@dataclass(frozen=True)
class Bee:
    x: int = 0
    target: int = -1
    mode: str = 'idle'
    load: int = 0

@dataclass(frozen=True)
class State:
    t: int
    hive: int
    nectar: int
    bees: tuple
    flowers: tuple
    spent: int = 0
    outcome: str = 'playing'

@dataclass(frozen=True)
class Timeline:
    state: State
    history: tuple = ()
LEVELS = (Level('隣の芽', (FlowerDef(1, 2, True), FlowerDef(3, 3, False), FlowerDef(5, 3, False)), ((0, 1), (1, 2)), 6), Level('少ない花を先に', (FlowerDef(1, 3, True), FlowerDef(2, 1, True), FlowerDef(4, 4, False)), ((1, 2),), 6), Level('両端の庭', (FlowerDef(1, 2, True), FlowerDef(2, 3, False), FlowerDef(5, 2, True), FlowerDef(4, 3, False)), ((0, 1), (2, 3)), 8))
REASONS = {'ready': '花を予約して、SPACEで一刻進めよう。', 'assign': '蜜1を予約。出発前ならCで取り消せます。', 'cancel': '出発前の予約を解きました。', 'tick': '蜂が一歩進み、芽の待ち時間が減りました。', 'move': '巣の蜜1を使い、巣台へ引っ越しました。', 'busy': '飛行中の行先は変更できません。', 'unavailable': '開花中の未予約蜜が必要です。', 'reserved': '先にCで予約を取り消してください。', 'need_idle': '2匹とも未予約で巣にいるときだけ移転できます。', 'cost': '巣に運搬用の蜜1が必要です。', 'same': '別の巣台を選んでください。', 'terminal': 'Xで直前へ戻すか、Rで再開できます。', 'invalid': 'その操作はできません。', 'undo': '直前の予約または一刻へ戻しました。', 'empty': '戻せる操作がまだありません。'}

def initial(l, open_c=False):
    return State(0, 0, 1, (Bee(), Bee()), tuple((Flower(f.stock, phase='open' if f.awake or (open_c and i == 2) else 'sleep') for i, f in enumerate(l.flowers))))

def potential(s):
    return s.nectar + sum((f.stock for f in s.flowers)) + sum((b.load for b in s.bees))

def finish(s, l):
    return replace(s, outcome='won' if s.nectar >= l.target else 'lost' if potential(s) < l.target else 'playing')

def apply(s, l, a):
    if s.outcome != 'playing':
        return (s, 'terminal')
    if not isinstance(a, tuple) or not a:
        return (s, 'invalid')
    k = a[0]
    bees = list(s.bees)
    fs = list(s.flowers)
    if k == 'assign' and len(a) == 3 and (type(a[1]) == int) and (type(a[2]) == int) and (0 <= a[1] < 2) and (0 <= a[2] < len(fs)):
        bi, fi = a[1:]
        b = bees[bi]
        f = fs[fi]
        if b.mode == 'reserved':
            return (s, 'reserved')
        if b.mode != 'idle':
            return (s, 'busy')
        if f.phase != 'open' or f.stock <= f.reserved:
            return (s, 'unavailable')
        bees[bi] = replace(b, target=fi, mode='reserved')
        fs[fi] = replace(f, reserved=f.reserved + 1)
        return (replace(s, bees=tuple(bees), flowers=tuple(fs)), k)
    if k == 'cancel' and len(a) == 2 and (type(a[1]) == int) and (0 <= a[1] < 2):
        b = bees[a[1]]
        if b.mode != 'reserved':
            return (s, 'busy')
        fs[b.target] = replace(fs[b.target], reserved=fs[b.target].reserved - 1)
        bees[a[1]] = Bee(s.hive)
        return (replace(s, bees=tuple(bees), flowers=tuple(fs)), k)
    if k == 'move' and len(a) == 2 and (type(a[1]) == int) and (a[1] in (0, 3, 6)):
        if any((b.mode != 'idle' for b in bees)):
            return (s, 'need_idle')
        if a[1] == s.hive:
            return (s, 'same')
        if s.nectar < 1:
            return (s, 'cost')
        fs = tuple((replace(f, timer=f.timer - 1, phase='open' if f.timer == 1 else 'wait') if f.phase == 'wait' else f for f in fs))
        return (finish(replace(s, t=s.t + 1, hive=a[1], nectar=s.nectar - 1, spent=s.spent + 1, bees=(Bee(a[1]), Bee(a[1])), flowers=fs), l), k)
    if a != ('tick',):
        return (s, 'invalid')
    fs = [replace(f, timer=f.timer - 1, phase='open' if f.timer == 1 else 'wait') if f.phase == 'wait' else f for f in fs]
    depleted = []
    nectar = s.nectar
    for i, b in enumerate(bees):
        if b.mode in ('reserved', 'out'):
            dest = l.flowers[b.target].x
            x = b.x + (dest > b.x) - (dest < b.x)
            if x == dest:
                f = fs[b.target]
                fs[b.target] = replace(f, stock=f.stock - 1, reserved=f.reserved - 1, phase='spent' if f.stock == 1 else 'open')
                bees[i] = Bee(x, b.target, 'back', 1)
                if f.stock == 1:
                    depleted.append(b.target)
            else:
                bees[i] = replace(b, x=x, mode='out')
        elif b.mode == 'back':
            x = b.x + (s.hive > b.x) - (s.hive < b.x)
            if x == s.hive:
                nectar += b.load
                bees[i] = Bee(s.hive)
            else:
                bees[i] = replace(b, x=x)
    for source in depleted:
        for a, b in l.links:
            target = b if a == source else a if b == source else -1
            if target >= 0 and fs[target].phase == 'sleep':
                fs[target] = replace(fs[target], phase='wait', timer=2)
    return (finish(replace(s, t=s.t + 1, nectar=nectar, bees=tuple(bees), flowers=tuple(fs)), l), 'tick')

def step(tl, l, a):
    if a == ('undo',):
        return (Timeline(tl.history[-1], tl.history[:-1]), 'undo') if tl.history else (tl, 'empty')
    q, r = apply(tl.state, l, a)
    return (Timeline(q, tl.history + (tl.state,)), r) if q != tl.state else (tl, r)

def forecast(s, l, a):
    return apply(s, l, a)

def invariant(s, l):
    if not (s.t >= 0 and s.hive in (0, 3, 6) and (s.nectar >= 0) and (s.spent >= 0) and (potential(s) + s.spent == 1 + sum((f.stock for f in l.flowers)))):
        return False
    for i, f in enumerate(s.flowers):
        if not 0 <= f.reserved <= f.stock <= l.flowers[i].stock:
            return False
        if f.reserved != sum((b.target == i and b.mode in ('out', 'reserved') for b in s.bees)):
            return False
        if f.phase not in ('open', 'sleep', 'wait', 'spent') or (f.phase == 'wait') != (f.timer > 0) or (f.phase == 'spent') != (f.stock == 0):
            return False
    for b in s.bees:
        if not 0 <= b.x <= 6 or b.mode not in ('idle', 'reserved', 'out', 'back') or b.load != (1 if b.mode == 'back' else 0):
            return False
        if b.mode in ('idle', 'reserved') and b.x != s.hive:
            return False
        if b.mode == 'idle' and b.target != -1:
            return False
    return True

def state_hash(s):
    return hashlib.sha256(json.dumps(asdict(s), sort_keys=True).encode('utf-8')).hexdigest()

def policy(l, kind='near', open_c=False, max_steps=300):
    """Bounded, transparent priorities. Does not certify human enjoyment."""
    s = initial(l, open_c)
    path = []
    seen = set()
    moved = set()
    for _ in range(max_steps):
        if s.outcome != 'playing':
            break
        a = None
        if kind in ('move', 'twice') and all((b.mode == 'idle' for b in s.bees)):
            dest = 3 if 3 not in moved and s.flowers[0].stock == 0 else 6 if kind == 'twice' and 3 in moved and (6 not in moved) and (s.flowers[1].stock == 0) else None
            if dest is not None and s.nectar >= 1 and (potential(s) > l.target):
                a = ('move', dest)
                moved.add(dest)
        if a is None:
            for bi, b in enumerate(s.bees):
                if b.mode != 'idle':
                    continue
                opts = [i for i, f in enumerate(s.flowers) if f.phase == 'open' and f.stock > f.reserved]
                if opts:
                    if kind == 'small':
                        opts.sort(key=lambda i: (s.flowers[i].stock, abs(l.flowers[i].x - s.hive), i))
                    elif kind == 'right':
                        opts.sort(reverse=True)
                    else:
                        opts.sort(key=lambda i: (abs(l.flowers[i].x - s.hive), i))
                    a = ('assign', bi, opts[0])
                    break
        if a is None:
            a = ('tick',)
        q, r = apply(s, l, a)
        if q == s:
            raise AssertionError((kind, a, r))
        s = q
        path.append(a)
    return (s, tuple(path))

STAGE = 5
STAGE_NAMES = ('ONE BEE DELIVERY', 'TWO BEES / RESERVATIONS', 'LINKED BLOOMING', 'MOVING HIVE / CONSERVATION', 'UNDO / FORECAST')

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

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

    @property
    def state(self):
        return self.timeline.state if STAGE==5 else self.current

    def reset(self):
        if STAGE==5:self.timeline=Timeline(initial(self.board))
        else:self.current=initial(self.board)
        self.selection=0
        self.target=0
        self.platform=1
        self.message='READY'
        self.gated=False

    def selected(self):
        if self.selection==2:return ('move',(0,3,6)[self.platform])
        return ('assign',self.selection,self.target)

    def action(self,a):
        if STAGE==5:self.timeline,self.message=step(self.timeline,self.board,a)
        else:self.current,self.message=apply(self.current,self.board,a)

    def press(self,key):
        if key=='R':self.reset();return
        if STAGE>=3 and key in ('1','2','3'):
            self.level=int(key)-1;self.reset();return
        if STAGE==5 and key=='X':self.action(('undo',));return
        if self.state.outcome!='playing':return
        if key=='TAB':self.selection=(self.selection+1)%(3 if STAGE>=4 else 2)
        if key in ('LEFT','RIGHT'):
            delta=1 if key=='RIGHT' else -1
            if self.selection<2:self.target=(self.target+delta)%len(self.board.flowers)
            else:self.platform=(self.platform+delta)%3
        if key=='Z':self.action(self.selected())
        if key=='C' and self.selection<2:self.action(('cancel',self.selection))
        if key=='SPACE':self.action(('tick',))

    def control(self,pressed,held):
        # 一度押したら全キーの解放を待つ。長押しで世界を進めない。
        if self.gated:
            if not pressed and not held:self.gated=False
            return
        if len(pressed|held)>1:
            self.message='RELEASE ALL KEYS';self.gated=True;return
        if len(pressed)==1:
            self.press(next(iter(pressed)));self.gated=True

    def draw(self,p):
        s=self.state;l=self.board;p.cls(1)
        p.text(12,12,'MOVING BEES / STAGE '+str(STAGE)+' / '+STAGE_NAMES[STAGE-1],10)
        p.text(12,30,'GARDEN '+('ABC'[self.level] if STAGE>=3 else 'PRACTICE')+'  TICK '+str(min(s.t,9999))+('+' if s.t>9999 else ''),7)
        p.rect(12,52,616,169,3)
        for a,b in l.links:
            p.line(52+l.flowers[a].x*88,177,52+l.flowers[b].x*88,177,4)
            for j in (a,b):p.line(52+l.flowers[j].x*88,159,52+l.flowers[j].x*88,177,4)
        for x in range(7):p.text(50+x*88,205,str(x),7)
        for x in (0,3,6):p.rectb(41+x*88,187,23,10,9)
        hx=52+s.hive*88;p.rect(hx-10,180,21,18,9);p.rect(hx-3,188,7,10,4)
        for i,(f,d) in enumerate(zip(s.flowers,l.flowers)):
            x=52+d.x*88;c=10 if f.phase=='open' else (11 if f.phase=='wait' else 5)
            p.line(x,123,x,166,11);p.circ(x,120,12,c);p.circ(x,120,4,4)
            p.text(x-3,95,chr(97+i),7)
            p.text(x-22,141,'STOCK '+str(f.stock),7)
            p.text(x-22,152,'BOOK '+str(f.reserved),7)
            p.text(x-22,164,f.phase.upper()+(' '+str(f.timer) if f.timer else ''),c if f.phase in ('open','wait') else 7)
            if self.selection<2 and i==self.target:p.rectb(x-28,88,58,92,10)
        for i,b in enumerate(s.bees):
            x=52+b.x*88;y=64+i*14
            p.circ(x-4,y-2,3,6);p.circ(x+4,y-2,3,6);p.rect(x-5,y,11,5,10);p.line(x,y,x,y+4,0)
            p.text(x+8,y,str(i+1)+('+' if b.load else ''),7)
        p.text(12,235,'HIVE NECTAR '+str(s.nectar)+' / TARGET '+str(l.target)+' / SPENT '+str(s.spent),10)
        p.text(12,251,'POTENTIAL '+str(potential(s))+' = HIVE + FLOWER STOCK (SLEEP INCLUDED) + LOAD',6)
        p.text(12,266,'BOOK IS PART OF STOCK. POTENTIAL IS NOT A REACHABILITY PROMISE.',13)
        for i,b in enumerate(s.bees):
            p.text(12,287+i*17,'BEE '+str(i+1)+' '+b.mode.upper()+' / X '+str(b.x)+' / LOAD '+str(b.load),11 if self.selection==i else 7)
        p.text(12,328,'SELECT '+('BEE '+str(self.selection+1)+' / FLOWER '+chr(97+self.target) if self.selection<2 else 'HIVE / PLATFORM '+str((0,3,6)[self.platform])),10)
        if STAGE==5:
            q,reason=forecast(s,l,self.selected())
            p.text(12,348,'NEXT Z: '+reason.upper()+' / TICK '+str(min(q.t,9999))+' / NECTAR '+str(q.nectar)+' / '+q.outcome.upper(),6)
            p.text(12,365,'HISTORY '+str(len(self.timeline.history))+' / FORECAST NEVER COMMITS',6)
        p.text(12,386,'RESULT '+s.outcome.upper()+' / '+self.message.upper(),11 if s.outcome=='won' else 7)
        p.text(12,410,'TAB: SELECT / LEFT-RIGHT: FLOWER OR PLATFORM / Z: BOOK OR MOVE' if STAGE>=4 else 'TAB: SELECT BEE / LEFT-RIGHT: FLOWER / Z: BOOK',7)
        p.text(12,427,'SPACE: ONE TICK / C: CANCEL BEFORE DEPARTURE / R: RESET / ESC: QUIT',7)
        p.text(12,444,('1-3: GARDEN / ' if STAGE>=3 else 'PRACTICE: ALL FLOWERS OPEN / NO LINKS / ')+('X: UNDO' if STAGE==5 else 'NO UNDO YET'),13)


def demo_keys():
    if STAGE==2:return ('Z','TAB','Z','SPACE','SPACE')
    if STAGE==3:return ('Z','TAB','Z','SPACE','SPACE','SPACE')
    return ('Z','TAB','Z','SPACE','SPACE','TAB','Z')

def main(headless=False,screenshot=None,demo=False):
    # Pyxel may change cwd. Resolve and validate all destinations before importing it.
    from pathlib import Path
    if screenshot:
        target=Path(screenshot).resolve()
        if target.suffix.lower()!='.png' or not target.parent.is_dir() or target.exists():
            raise ValueError('use a new .png in an existing directory')
        screenshot=str(target)
    import pyxel as p
    p.init(640,464,title='Moving bees lesson '+str(STAGE),fps=30,display_scale=1,headless=headless)
    app=App();sequence=demo_keys() if demo else ()
    keys={name:getattr(p,'KEY_'+name) for name in ('LEFT','RIGHT','TAB','SPACE','Z','C','X','R','1','2','3')}
    limit=max(20,2*len(sequence)+8)
    def update():
        if p.btnp(p.KEY_ESCAPE) or headless and p.frame_count>=limit:
            if screenshot:app.draw(p);p.screen.save(screenshot,1)
            p.quit();return
        if headless:
            if demo:
                index=p.frame_count//2
                pressed={sequence[index]} if p.frame_count%2==0 and index<len(sequence) 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__':
    import argparse
    from pathlib import Path
    parser=argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--headless-smoke',action='store_true')
    parser.add_argument('--demo-checkpoint',action='store_true')
    parser.add_argument('--screenshot')
    args=parser.parse_args()
    if args.demo_checkpoint and not args.headless_smoke:parser.error('demo requires headless mode')
    if args.screenshot and not args.headless_smoke:parser.error('screenshot requires headless mode')
    if args.screenshot:
        output=Path(args.screenshot).resolve()
        if output.suffix.lower()!='.png' or not output.parent.is_dir() or output.exists():parser.error('use a new .png in an existing directory')
        args.screenshot=str(output)
    main(args.headless_smoke,args.screenshot,args.demo_checkpoint)