# 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.
# 
# Third-party exception: bitmap glyph data derived from Noto Sans CJK is licensed
# under the SIL Open Font License 1.1. See assets/NOTO-LICENSE.txt.

"""からくり郵便室・教材5/5。単独で起動できます。"""
from dataclasses import dataclass, replace

CARD_NAMES = ("COLOR", "WEIGHT", "LEFT", "RIGHT")


@dataclass(frozen=True)
class Parcel:
    id: int
    color: str
    weight: int
    destination: str


# 原作Bの6便。番号の内部値は0始まり、画面だけ1始まり。
PARCELS = tuple(Parcel(i, color, weight, color+("H" if weight == 2 else "L"))
                for i, (color, weight) in enumerate(
                    (("R", 1), ("B", 2), ("R", 2), ("B", 1), ("B", 2), ("R", 1))))


def branch(card, parcel):
    # 出口は左=0、右=1。台番号や箱番号とは区別する。
    if card == 0:
        return 0 if parcel.color == "R" else 1
    if card == 1:
        return 0 if parcel.weight == 2 else 1
    return card-2


CAP = 3
LABELS = ("RL", "RH", "BL", "BH")
LABEL_NAMES = {"RL": "赤・軽", "RH": "赤・重", "BL": "青・軽", "BH": "青・重"}


@dataclass(frozen=True)
class Delivery:
    parcel: int
    box: int
    tick: int
    correct: bool
    path: tuple


@dataclass(frozen=True)
class State:
    cards: tuple = (2, 2, 2)
    labels: tuple = ("RL",)*4
    t: int = 0
    next_id: int = 1
    queue: tuple = (0,)
    tables: tuple = (-1, -1, -1)
    deliveries: tuple = ()
    paths: tuple = ((),)*6
    status: str = "edit"
    reason: str = "カードと箱札を編集して試運転。"
    moves: tuple = ()
    waits: tuple = ()


def condition(card, parcel):
    if card == 0:
        return "赤" if parcel.color == "R" else "青"
    if card == 1:
        return "重さ"+str(parcel.weight)
    return "そのまま"


def count_tests(state):
    return sum(card < 2 for card in state.cards)


def start(state):
    if state.status != "edit" or count_tests(state) > CAP:
        return state
    return replace(state, status="run", reason="試運転中。Spaceで一刻。",
                   moves=(), waits=())


def tick(state):
    if state.status != "run":
        return state
    t = state.t+1
    queue = list(state.queue)
    next_id = state.next_id
    if next_id < len(PARCELS) and next_id <= t:
        queue.append(next_id)
        next_id += 1
    # oldは読み取り専用。tablesだけを書き換えて最後に同時確定する。
    old = state.tables
    tables = list(old)
    paths = list(state.paths)
    deliveries = list(state.deliveries)
    moves, waits, faults = [], [], []
    for table in (1, 2):
        pid = old[table]
        if pid < 0:
            continue
        parcel = PARCELS[pid]
        exit_side = branch(state.cards[table], parcel)
        box = (table-1)*2+exit_side
        path = paths[pid]+(chr(65+table)+":"+condition(state.cards[table], parcel)+"→"+
                          ("左" if exit_side == 0 else "右"),
                          "箱"+str(box+1)+":"+LABEL_NAMES[state.labels[box]])
        paths[pid] = path
        correct = state.labels[box] == parcel.destination
        deliveries.append(Delivery(pid, box, t, correct, path))
        tables[table] = -1
        moves.append((pid, chr(65+table), "箱"+str(box+1)))
        if not correct:
            faults.append(pid)
    if old[0] >= 0:
        pid = old[0]
        parcel = PARCELS[pid]
        exit_side = branch(state.cards[0], parcel)
        target = 1+exit_side
        if old[target] < 0:
            tables[0] = -1
            tables[target] = pid
            paths[pid] += ("A:"+condition(state.cards[0], parcel)+"→"+
                           ("左" if exit_side == 0 else "右"),)
            moves.append((pid, "A", chr(65+target)))
        else:
            waits.append((pid, "この刻の初めは "+chr(65+target)+" が使用中"))
    if queue:
        if old[0] < 0:
            pid = queue.pop(0)
            tables[0] = pid
            moves.append((pid, "入口", "A"))
        else:
            waits.append((queue[0], "この刻の初めは A が使用中"))
    status = "fault" if faults else ("done" if len(deliveries) == 6 else
                                    ("timeout" if t >= 32 else "run"))
    reason = ("誤配: 荷物 "+",".join(str(pid+1) for pid in faults)+" の命令列を確認。" if faults else
              ("全6便の仕分け達成。空箱・重複札も正当な配置です。" if status == "done" else
               ("32刻で停止。" if status == "timeout" else "一刻の移動を同時に反映しました。")))
    return replace(state, t=t, next_id=next_id, queue=tuple(queue), tables=tuple(tables),
                   paths=tuple(paths), deliveries=tuple(deliveries), status=status,
                   reason=reason, moves=tuple(moves), waits=tuple(waits))


def parcel_conservation(state):
    ids = list(state.queue)+[pid for pid in state.tables if pid >= 0]
    ids += [delivery.parcel for delivery in state.deliveries]
    ids += list(range(state.next_id, len(PARCELS)))
    return sorted(ids) == list(range(len(PARCELS)))


def apply(state, action):
    if action == "start":
        return start(state)
    if action == "tick":
        return tick(state)
    if state.status == "edit" and isinstance(action, tuple) and len(action) == 3:
        kind, index, value = action
        if type(index) is not int:
            return state
        if kind == "card" and 0 <= index < 3 and type(value) is int and 0 <= value < 4:
            return replace(state, cards=tuple(value if i == index else card
                                             for i, card in enumerate(state.cards)))
        if kind == "label" and 0 <= index < 4 and value in LABELS:
            return replace(state, labels=tuple(value if i == index else label
                                              for i, label in enumerate(state.labels)))
    return state


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


def act(timeline, action):
    if action == "undo":
        if not timeline.past:
            return timeline
        return Timeline(timeline.past[-1], timeline.past[:-1], timeline.actions+(action,))
    if len(timeline.actions) >= 512:
        return timeline
    state = apply(timeline.state, action)
    if state == timeline.state:
        return timeline
    return Timeline(state, timeline.past+(timeline.state,), timeline.actions+(action,))


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

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

    def reset(self):
        self.timeline = Timeline()
        self.cursor = 0
        self.section = 0
        self.locked = False
        self.message = "EDIT CARDS AND LABELS, THEN ENTER"

    def commit(self, action):
        before = self.state
        self.timeline = act(self.timeline, action)
        if self.state == before:
            self.message = "NO CHANGE / CHECK MODE, CAP OR HISTORY LIMIT"
        elif action == "undo":
            self.message = "UNDO WHOLE STATE / START CAN ALSO BE UNDONE"
        else:
            self.message = "OLD OCCUPANCY ONLY / ONE TABLE PER PARCEL PER TICK"

    def control(self, down):
        down = set(down)
        if self.locked:
            if not down:
                self.locked = False
            return
        if not down:
            return
        self.locked = True
        if len(down) != 1:
            self.message = "ONE KEY AT A TIME / RELEASE ALL"
            return
        key = next(iter(down))
        if key == "R":
            self.reset()
            self.locked = True
            return
        if key == "X":
            self.commit("undo")
            return
        if key == "RETURN":
            self.commit("start")
        elif key == "SPACE":
            self.commit("tick")
        elif self.state.status == "edit":
            if key == "TAB":
                self.section = 1-self.section
                self.cursor = 0
            elif key in ("LEFT", "RIGHT", "UP", "DOWN"):
                step = -1 if key in ("LEFT", "UP") else 1
                self.cursor = (self.cursor+step)%(3 if self.section == 0 else 4)
            elif key == "Z":
                if self.section == 0:
                    self.commit(("card", self.cursor, (self.state.cards[self.cursor]+1)%4))
                else:
                    value = LABELS[(LABELS.index(self.state.labels[self.cursor])+1)%4]
                    self.commit(("label", self.cursor, value))

    def draw(self, p):
        s = self.state
        p.cls(1)
        p.text(8, 8, "CLOCKWORK POST / STAGE 5 / ORIGINAL B", 10)
        p.text(8, 22, "TICK "+str(s.t)+"/32   STATUS "+s.status.upper()+"   DELIVERED "+str(len(s.deliveries))+"/6", 7)
        p.text(8, 34, "TEST CARDS "+str(count_tests(s))+"/"+str(CAP)+"   OLD-STATE SYNCHRONOUS TICKS", 6)
        for i, parcel in enumerate(PARCELS):
            p.text(8+i*65, 49, str(i+1)+":"+parcel.color+str(parcel.weight)+">"+parcel.destination, 8 if parcel.color == "R" else 12)
        queue = " ".join(str(pid+1) for pid in s.queue) or "EMPTY"
        p.text(8, 63, "QUEUE: "+queue+"   NOT ARRIVED: "+str(6-s.next_id), 7)
        points = ((200, 96), (100, 149), (300, 149))
        for x in (100, 300):
            p.line(200, 116, x, 137, 5)
        for i, (x, y) in enumerate(points):
            p.rect(x-49, y-12, 98, 39, 2)
            p.rectb(x-49, y-12, 98, 39, 10 if self.section == 0 and self.cursor == i else 5)
            p.text(x-40, y-5, chr(65+i)+": "+CARD_NAMES[s.cards[i]], 7)
            pid = s.tables[i]
            p.text(x-40, y+10, "EMPTY" if pid < 0 else "PARCEL "+str(pid+1), 6 if pid < 0 else 11)
        for i, x in enumerate((50, 150, 250, 350)):
            p.line(100 if i < 2 else 300, 176, x, 193, 5)
            p.rect(x-37, 193, 74, 35, 4)
            p.rectb(x-37, 193, 74, 35, 10 if self.section == 1 and self.cursor == i else 9)
            p.text(x-31, 199, "BOX "+str(i+1)+" "+s.labels[i], 7)
            received = [str(d.parcel+1)+("!" if not d.correct else "") for d in s.deliveries if d.box == i]
            p.text(x-31, 213, " ".join(received) or "EMPTY", 11)
        waits = ",".join(str(pid+1) for pid, _ in s.waits) or "NONE"
        p.text(8, 237, "WAITED THIS TICK: "+waits+" / ALL 6 CONSERVED: "+str(parcel_conservation(s)), 6)
        if s.deliveries:
            d = s.deliveries[-1]
            p.text(8, 249, "LAST: "+str(d.parcel+1)+" > A > "+("B" if d.box < 2 else "C")+" > BOX "+str(d.box+1)+" "+("OK" if d.correct else "WRONG / STOP"), 11 if d.correct else 8)
        else:
            p.text(8, 249, "ARRIVED NOW? READ ITS CARD NEXT TICK.", 6)
        p.text(8, 264, self.message, 10)
        p.text(8, 278, "ARROWS SELECT / TAB TABLE-BOX / Z CYCLE / ENTER START", 7)
        p.text(8, 289, "SPACE TICK / X UNDO / EDIT ONLY IN EDIT MODE", 7)
        p.text(8, 304, "COLOR R:LEFT B:RIGHT / WEIGHT 2:LEFT 1:RIGHT", 6)
        p.text(8, 315, "LEFT/RIGHT: UNCONDITIONAL / R:RED B:BLUE", 6)
        p.text(8, 326, "DEST RH:RED HEAVY RL:RED LIGHT BH/BL:BLUE", 6)
        p.text(8, 344, "R RESET / RELEASE ALL KEYS BETWEEN ACTIONS / ESC QUIT", 5)


DEMO_KEYS = ('Z', 'Z', 'RIGHT', 'Z', 'Z', 'Z', 'RIGHT', 'Z', 'Z', 'Z', 'TAB', 'Z', 'RIGHT', 'RIGHT', 'Z', 'Z', 'Z', 'RIGHT', 'Z', 'Z', 'RETURN', 'SPACE', 'SPACE', 'SPACE', 'SPACE', 'SPACE', 'SPACE', 'SPACE', 'SPACE', 'SPACE', 'SPACE', 'SPACE', 'SPACE', 'SPACE', 'X')


def main(headless=False, screenshot=None, demo=False):
    from pathlib import Path
    import pyxel as p
    # Pyxelの初期化より先に、指定した保存先を絶対化する。
    screenshot = str(Path(screenshot).resolve()) if screenshot else None
    if screenshot:
        Path(screenshot).parent.mkdir(parents=True, exist_ok=True)
    app = App()
    p.init(400, 360, title="Clockwork Post", fps=30, display_scale=2,
           headless=headless)
    keys = {name: getattr(p, "KEY_"+name) for name in
            ("UP", "DOWN", "LEFT", "RIGHT", "TAB", "Z", "P", "RETURN", "SPACE", "X", "R")}
    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)
            from dataclasses import asdict
            import json
            print(json.dumps({"state": asdict(app.state), "message": app.message},
                             ensure_ascii=True, sort_keys=True), flush=True)
            p.quit()
            return
        if headless and demo:
            index = p.frame_count//2
            down = {DEMO_KEYS[index]} if p.frame_count%2 == 0 and index < len(DEMO_KEYS) else set()
        else:
            down = {key for key, value in keys.items() if p.btn(value)}
        app.control(down)

    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)