# MIT License
# 
# Copyright (c) 2026 Pyxel 100 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.
# 
# The bundled Noto-derived bitmap glyphs are separately licensed under SIL OFL.
# See assets/NOTO-LICENSE.txt.

"""お弁当列車・教材5/5。原作Bを段階的に組み立てる。"""
from dataclasses import dataclass, replace

STOCK = ("F", "F", "V", "V", "E", "E")
KINDS = "FVE"
PRESET = ((4, 0, 1), (3, 2, 5))
ORDERS = ((1, 1, 0), (1, 0, 1), (0, 1, 1))
WORK = (3, 2, 3)
GOAL = 5
CAPACITY = 3

@dataclass(frozen=True)
class State:
    station: int
    remaining: int
    shelves: tuple
    hand: int
    warehouse: tuple
    sold: tuple
    orders: tuple
    status: str = 'load'
    transfers: int = 0
    actions: int = 0
    departed: tuple = ()

def initial():
    return State(0, WORK[0], PRESET, -1, (), (), ORDERS, status='load', departed=())

def advance(s):
    if s.station == len(WORK) - 1:
        return replace(s, remaining=0, status='won' if len(s.sold) >= GOAL else 'short')
    return replace(s, station=s.station + 1, remaining=WORK[s.station + 1])

def step(s, a):
    if s.status in ('won', 'short') or not isinstance(a, tuple) or (not a):
        return s
    if s.status == 'load':
        if a == ('depart',):
            return replace(s, status='play', departed=s.shelves)
        if a[0] == 'place' and len(a) == 3:
            _, item, dest = a
            if not isinstance(item, int) or not 0 <= item < len(STOCK) or dest not in (0, 1, 2):
                return s
            owner = next((j for j, sh in enumerate(s.shelves) if item in sh), 2)
            if owner == dest:
                return s
            if dest < 2 and len(s.shelves[dest]) >= CAPACITY:
                return s
            shelves = [list(sh) for sh in s.shelves]
            warehouse = list(s.warehouse)
            if owner == 2:
                warehouse.remove(item)
            else:
                shelves[owner].remove(item)
            if dest == 2:
                warehouse.append(item)
            else:
                shelves[dest].insert(0, item)
            return replace(s, shelves=tuple((tuple(sh) for sh in shelves)), warehouse=tuple(warehouse))
        if a[0] == 'swap_next' and len(a) == 2:
            item = a[1]
            for j, sh in enumerate(s.shelves):
                if item in sh and len(sh) > 1:
                    row = list(sh)
                    k = row.index(item)
                    n = (k + 1) % len(row)
                    row[k], row[n] = (row[n], row[k])
                    shelves = list(s.shelves)
                    shelves[j] = tuple(row)
                    return replace(s, shelves=tuple(shelves))
            return s
        return s
    if a == ('depart',):
        return advance(s)
    shelves = [list(sh) for sh in s.shelves]
    hand = s.hand
    sold = s.sold
    orders = [list(order) for order in s.orders]
    transfer = 0
    if a[0] in ('sell', 'take') and len(a) in (2, 3):
        src = a[1]
        depth = a[2] if len(a) == 3 else 0
        if src not in (0, 1, 2) or not isinstance(depth, int) or depth < 0:
            return s
        if src == 2:
            if hand < 0 or a[0] == 'take' or depth != 0:
                return s
            item = hand
        else:
            if depth >= len(shelves[src]) or depth != 0:
                return s
            item = shelves[src][depth]
        if a[0] == 'take':
            if hand >= 0:
                return s
            shelves[src].pop(depth)
            hand = item
            transfer = 1
        else:
            kind = KINDS.index(STOCK[item])
            if orders[s.station][kind] <= 0:
                return s
            orders[s.station][kind] -= 1
            sold = sold + (item,)
            if src == 2:
                hand = -1
            else:
                shelves[src].pop(depth)
    elif a[0] == 'put' and len(a) == 2:
        dest = a[1]
        if dest not in (0, 1) or hand < 0 or len(shelves[dest]) >= CAPACITY:
            return s
        shelves[dest].insert(0, hand)
        hand = -1
        transfer = 1
    else:
        return s
    n = replace(s, shelves=tuple((tuple(sh) for sh in shelves)), hand=hand, sold=sold, orders=tuple((tuple(row) for row in orders)), remaining=s.remaining - 1, actions=s.actions + 1, transfers=s.transfers + transfer)
    return advance(n) if n.remaining == 0 else n

def invariant(s):
    owned = sum(s.shelves, ()) + s.warehouse + s.sold + ((s.hand,) if s.hand >= 0 else ())
    assert sorted(owned) == list(range(len(STOCK))) and len(set(owned)) == len(owned)
    assert all((len(sh) <= CAPACITY for sh in s.shelves))
    assert all((n >= 0 for row in s.orders for n in row))
    assert sum((sum(row) for row in ORDERS)) - sum((sum(row) for row in s.orders)) == len(s.sold)
    assert 0 <= s.station < len(WORK) and 0 <= s.remaining <= WORK[s.station]
    assert s.transfers + s.sold.__len__() == s.actions
    assert s.status in ('load', 'play', 'won', 'short')

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


def apply(timeline, action):
    state = step(timeline.state, action)
    return timeline if state == timeline.state else Timeline(state, timeline.history+(timeline.state,))


def undo(timeline):
    return Timeline(timeline.history[-1], timeline.history[:-1]) if timeline.history else timeline


class App:
    def __init__(self):
        self.reset()

    def reset(self):
        self.timeline = Timeline(initial())
        self.selection = 0
        self.item = 0
        self.gated = False
        self.message = "ONE KEY AT A TIME"

    @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 == "X":
            before = self.state
            self.timeline = undo(self.timeline)
            self.message = "UNDO" if before != self.state else "NOTHING TO UNDO"
            self.gated = True
            return
        if self.state.status in ("won", "short"):
            return
        if key in ("UP", "DOWN"):
            self.selection = (self.selection+(-1 if key=="UP" else 1))%3
            self.gated = True
            return
        if self.state.status == "load" and key in ("LEFT", "RIGHT"):
            self.item = (self.item+(-1 if key=="LEFT" else 1))%len(STOCK)
            self.gated = True
            return
        if self.state.status == "load":
            action = ("place", self.item, self.selection) if key=="Z" else ("swap_next", self.item) if key=="C" else ("depart",) if key=="SPACE" else None
        else:
            action = ("take", self.selection) if key=="C" else ("put", self.selection) if key=="V" else ("sell", self.selection) if key=="Z" else ("depart",) if key=="SPACE" else None
        if action is None:
            return
        before = self.state
        self.timeline = apply(self.timeline, action)
        self.message = "CONFIRMED" if before != self.state else "BLOCKED / NO COST"
        self.gated = True

    def draw(self, p):
        p.cls(1)
        p.text(8, 8, "LUNCH TRAIN / STAGE 5", 10)
        loading = self.state.status == "load"
        p.text(8, 23, "LOADING / NO TIME COST" if loading else "STATION "+str(self.state.station+1)+" / WORK "+str(self.state.remaining), 7)
        for row, items in enumerate(self.state.shelves):
            y = 65+row*47
            p.text(8, y+10, ("> " if self.selection==row else "  ")+("A" if row==0 else "B"), 10 if self.selection==row else 7)
            for depth in range(CAPACITY):
                x = 65+depth*75
                p.rectb(x, y, 64, 30, 6)
                if depth<len(items):
                    item = items[depth]
                    p.text(x+6, y+11, str(item+1)+":"+STOCK[item], 7)
            p.text(65, y-9, "FRONT -> BACK", 6)
        p.text(8, 170, ("> " if self.selection==2 else "  ")+"HAND: "+("-" if self.state.hand<0 else str(self.state.hand+1)+":"+STOCK[self.state.hand]), 10 if self.selection==2 else 7)
        if loading:
            p.text(8, 39, "SELECT ITEM "+str(self.item+1)+":"+STOCK[self.item]+" / DEST "+("A", "B", "WAREHOUSE")[self.selection], 11)
            p.text(8, 186, "WAREHOUSE: "+" ".join(str(i+1)+":"+STOCK[i] for i in self.state.warehouse), 7)
        p.text(8, 202, "ORDERS F/V/E: "+" / ".join(map(str,self.state.orders[self.state.station])), 11)
        p.text(8, 218, "SOLD: "+str(len(self.state.sold))+" / GOAL "+str(GOAL)+"  "+self.state.status.upper(), 7)
        for i, order in enumerate(ORDERS):
            p.text(8+i*122, 237, str(i+1)+": "+"/".join(map(str,order))+" W"+str(WORK[i]), 6)
        p.text(8, 255, self.message, 7)
        if loading:
            p.text(8, 274, "LEFT/RIGHT ITEM  UP/DOWN DEST", 6)
            p.text(8, 286, "Z PLACE  C SWAP NEXT  SPACE DEPART", 6)
        else:
            p.text(8, 274, "UP/DOWN SOURCE: A / B / HAND", 6)
            p.text(8, 286, "FRONT ITEMS ONLY / HAND PERSISTS", 6)
        p.text(8, 301, "PLAY: Z SELL  C TAKE  V PUT  SPACE DEPART", 6)
        p.text(8, 313, "X UNDO  R RESET  ESC QUIT", 6)

DEMO_KEYS = ('SPACE', 'C', 'Z', 'DOWN', 'Z', 'UP', 'Z', 'DOWN', 'DOWN', 'Z', 'UP', 'Z', 'Z', 'SPACE')
CONTROL_KEYS = ("UP", "DOWN", "LEFT", "RIGHT", "Z", "C", "V", "SPACE", "X", "R")

def main(headless=False, screenshot=None, demo=False):
    from pathlib import Path
    import pyxel as p
    screenshot = str(Path(screenshot).resolve()) if screenshot else None
    if screenshot:
        Path(screenshot).parent.mkdir(parents=True, exist_ok=True)
    app = App()
    p.init(384, 330, title="Lunch Train", fps=30, display_scale=2, headless=headless)
    keys = {k: getattr(p, "KEY_"+k) for k in CONTROL_KEYS}
    limit = max(12, 2*len(DEMO_KEYS)+3) if demo else 12
    def update():
        if headless and p.frame_count >= limit:
            if screenshot:
                p.screenshot(screenshot, scale=2)
            p.quit()
            return
        if headless and demo:
            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))
    return app

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("--headless-smoke", action="store_true")
    parser.add_argument("--demo-checkpoint", action="store_true")
    parser.add_argument("--screenshot")
    args = parser.parse_args()
    main(args.headless_smoke, args.screenshot, args.demo_checkpoint)