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

"""町工場の一日・教材 2/5。単体起動可能、外部素材なし。"""
from dataclasses import dataclass, replace

@dataclass(frozen=True)
class Job:
    id: str
    cut: int
    finish: int


@dataclass(frozen=True)
class Level:
    title: str
    jobs: tuple


LEVELS = (Level("Parallel work", (Job("a", 1, 2), Job("b", 2, 1),
                                       Job("c", 1, 1), Job("d", 1, 2))),)
CAPACITY = 2


@dataclass(frozen=True)
class State:
    t: int = 0
    cjob: int = -1
    crem: int = 0
    fjob: int = -1
    frem: int = 0
    queue: tuple = ()
    started: tuple = ()
    done: tuple = ()


def initial(level):
    return State(started=(False,) * len(level.jobs), done=(-1,) * len(level.jobs))


def actions(level):
    cuts = tuple(("cut", i) for i in range(len(level.jobs)))
    return cuts + (("finish",), ("wait",))


def apply(state, level, action):
    if action not in actions(level):
        return state, "invalid"
    next_state = state
    kind = action[0]
    if kind == "cut":
        index = action[1]
        job = level.jobs[index]
        if state.started[index]:
            return state, "started"
        if state.cjob >= 0:
            return state, "busy"
        started = list(state.started)
        started[index] = True
        next_state = replace(next_state, cjob=index, crem=job.cut,
                             started=tuple(started))
    elif kind == "finish":
        if state.fjob >= 0:
            return state, "busy"
        if not state.queue:
            return state, "empty"
        index = state.queue[0]
        job = level.jobs[index]
        next_state = replace(next_state, fjob=index, frem=job.finish,
                             queue=state.queue[1:])
    # 指示を反映した後で進める。開始した刻も1刻目に数える。
    next_state = replace(next_state, crem=max(0, next_state.crem - 1),
                         frem=max(0, next_state.frem - 1))
    # Fの取り出しは指示時。Cの搬出は同じ刻の最後。
    if next_state.fjob >= 0 and next_state.frem == 0:
        index = next_state.fjob
        done = list(next_state.done)
        done[index] = state.t + 1
        next_state = replace(next_state, fjob=-1, done=tuple(done))
    if next_state.cjob >= 0 and next_state.crem == 0 and len(next_state.queue) < CAPACITY:
        next_state = replace(next_state, queue=next_state.queue + (next_state.cjob,), cjob=-1)
    next_state = replace(next_state, t=state.t + 1)
    check(next_state, level)
    return next_state, "ok"


def check(state, level):
    assert state.t >= 0 and state.crem >= 0
    assert len(state.started) == len(level.jobs) == len(state.done)
    assert 0 <= len(state.queue) <= CAPACITY and state.frem >= 0
    owners = [i for i in (state.cjob, state.fjob) if i >= 0] + list(state.queue)
    assert state.fjob >= 0 or state.frem == 0
    owners += [i for i, time in enumerate(state.done) if time >= 0]
    assert len(owners) == len(set(owners))
    assert set(owners) == {i for i, started in enumerate(state.started) if started}
    assert all(0 <= i < len(level.jobs) for i in owners)
    assert state.cjob >= 0 or state.crem == 0


def location(state, index):
    if state.done[index] >= 0:
        return "DONE @" + str(state.done[index])
    if state.cjob == index:
        return "C REM " + str(state.crem) if state.crem else "C HOLD"
    if state.fjob == index:
        return "F REM " + str(state.frem)
    if index in state.queue:
        return "QUEUE " + str(state.queue.index(index) + 1)
    return "NOT STARTED"


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.selection = 0
        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))
        choices = actions(self.level)
        if key in ("LEFT", "RIGHT"):
            self.selection = (self.selection + (1 if key == "RIGHT" else -1)) % len(choices)
        elif key in ("Z", "SPACE"):
            self.state, self.message = apply(self.state, self.level, choices[self.selection])
        else:
            return
        self.gated = True

    def draw(self, p):
        p.cls(1)
        p.text(12, 10, "SMALL FACTORY / STAGE 2", 10)
        state = self.state
        p.text(12, 27, self.level.title + "  T=" + str(state.t), 7)
        p.rect(12, 48, 164, 78, 2)
        p.rectb(12, 48, 164, 78, 6)
        p.text(22, 58, "C / CUTTING", 10)
        p.text(22, 77, "IDLE" if state.cjob < 0 else "JOB " + self.level.jobs[state.cjob].id, 7)
        p.text(22, 94, "REM " + str(state.crem), 11)
        p.rect(336, 48, 164, 78, 2)
        p.rectb(336, 48, 164, 78, 6)
        p.text(346, 58, "F / FINISHING", 10)
        p.text(346, 77, "IDLE" if state.fjob < 0 else "JOB " + self.level.jobs[state.fjob].id, 7)
        p.text(346, 94, "REM " + str(state.frem), 11)
        p.text(194, 58, "RAIL / FIFO " + str(len(state.queue)) + "/" + str(CAPACITY), 6)
        for slot in range(CAPACITY):
            x = 197 + slot * 50
            p.rectb(x, 77, 38, 26, 5)
            if slot < len(state.queue):
                p.text(x + 12, 87, self.level.jobs[state.queue[slot]].id, 10)
        p.text(194, 115, "HEAD -> TAIL", 13)
        p.text(12, 144, "JOB  TIMES   LOCATION", 13)
        for index, job in enumerate(self.level.jobs):
            y = 162 + index * 17
            if state.done[index] >= 0:
                p.rect(10, y - 2, 491, 13, 3)
            times = str(job.cut)
            times += "/" + str(job.finish)
            p.text(12, y, job.id, 10)
            p.text(36, y, times, 7)
            p.text(96, y, location(state, index), 7)
        for index, action in enumerate(actions(self.level)):
            x = 12 + (index % 4) * 124
            y = 244 + (index // 4) * 23
            label = "CUT " + self.level.jobs[action[1]].id if action[0] == "cut" else " ".join(action)
            p.rectb(x, y, 118, 19, 10 if index == self.selection else 5)
            p.text(x + 5, y + 6, label.upper(), 10 if index == self.selection else 7)
        p.text(12, 305, "FULL RAIL? C HOLDS THE COMPLETED JOB. NOTHING IS LOST.", 6)
        p.text(12, 340, self.message.upper(), 11)
        p.text(12, 359, "LEFT/RIGHT SELECT  Z/SPACE DO  R RESET  ESC QUIT", 7)


DEMO_ACTIONS = (('cut', 0), ('cut', 2), ('cut', 3), ('finish',))
KEYS = ("LEFT", "RIGHT", "Z", "SPACE", "X", "R", "1", "2", "3", "B")


def demo_keys():
    selection = 0
    choices = actions(LEVELS[0])
    result = []
    for action in DEMO_ACTIONS:
        if action == ("undo",):
            result.append("X")
            continue
        target = choices.index(action)
        while selection != target:
            result.append("RIGHT")
            selection = (selection + 1) % len(choices)
        result.append("Z")
    return tuple(result)


def parse_args(argv=None):
    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")
    parser.add_argument("--frames", type=int)
    args = parser.parse_args(argv)
    if (args.demo_checkpoint or args.screenshot or args.frames is not None) and not args.headless_smoke:
        parser.error("--demo-checkpoint / --screenshot / --frames require --headless-smoke")
    minimum = 2 * len(demo_keys()) + 2 if args.demo_checkpoint else 2
    if args.frames is not None and not minimum <= args.frames <= 600:
        parser.error("--frames must be between " + str(minimum) + " and 600")
    if args.screenshot:
        path = Path(args.screenshot)
        if path.suffix.lower() != ".png" or path.is_dir():
            parser.error("--screenshot must name a PNG file")
        args.screenshot = str(path.resolve())
    args.frames = args.frames or max(8, minimum)
    return args


def main(args):
    from pathlib import Path
    import pyxel as p
    if args.screenshot:
        Path(args.screenshot).parent.mkdir(parents=True, exist_ok=True)
    app = App()
    demo = demo_keys()
    frame = 0
    p.init(512, 396, title="Small Factory", fps=30, display_scale=2, headless=args.headless_smoke)
    keys = {key: getattr(p, "KEY_" + key) for key in KEYS}

    def update():
        nonlocal frame
        frame += 1
        if args.headless_smoke:
            if args.demo_checkpoint:
                index = (frame - 1) // 2
                pressed = {demo[index]} if frame % 2 == 1 and index < len(demo) else set()
                app.control(pressed, pressed)
        else:
            app.control({key for key, code in keys.items() if p.btnp(code)},
                        {key for key, code in keys.items() if p.btn(code)})

    def draw():
        app.draw(p)
        if args.headless_smoke and frame >= args.frames:
            if args.screenshot:
                p.screenshot(args.screenshot, scale=2)
            p.quit()
    p.run(update, draw)


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