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

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

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


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


LEVELS = (
    Level("Group wood", (Job("a", "W", 1, 1, 2, 4),
                         Job("b", "W", 1, 1, 2, 5),
                         Job("c", "M", 1, 1, 3, 8)), 8, 7),
    Level("Long or two short", (Job("a", "W", 1, 1, 2, 4),
                               Job("L", "W", 3, 3, 4, 8),
                               Job("b", "M", 1, 1, 3, 7)), 9, 5),
    Level("Head of the rail", (Job("a", "W", 1, 2, 2, 4),
                              Job("b", "M", 2, 1, 3, 9),
                              Job("c", "W", 1, 1, 2, 6),
                              Job("d", "M", 1, 2, 3, 11)), 11, 7),
)
CAPACITY = 2


@dataclass(frozen=True)
class State:
    t: int = 0
    ctype: str = "W"
    ftype: str = "W"
    cjob: int = -1
    crem: int = 0
    fjob: int = -1
    frem: int = 0
    queue: tuple = ()
    started: tuple = ()
    done: tuple = ()
    value: int = 0
    cleans: int = 0
    outcome: str = "running"


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",), ("clean", "C"), ("clean", "F"), ("wait",))


def apply(state, level, action):
    if state.outcome != "running":
        return state, "terminal"
    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"
        if state.ctype not in ("-", job.kind):
            return state, "kind"
        started = list(state.started)
        started[index] = True
        next_state = replace(next_state, cjob=index, crem=job.cut,
                             ctype=job.kind, 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]
        if state.ftype not in ("-", job.kind):
            return state, "kind"
        next_state = replace(next_state, fjob=index, frem=job.finish,
                             ftype=job.kind, queue=state.queue[1:])
    elif kind == "clean":
        machine = action[1]
        busy = state.cjob if machine == "C" else state.fjob
        material = state.ctype if machine == "C" else state.ftype
        if busy >= 0:
            return state, "busy"
        if material == "-":
            return state, "clean"
        fields = {"ctype": "-"} if machine == "C" else {"ftype": "-"}
        next_state = replace(next_state, **fields, cleans=state.cleans + 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
        job = level.jobs[index]
        reward = job.value if state.t + 1 <= job.due else max(1, job.value - 1)
        next_state = replace(next_state, fjob=-1, done=tuple(done),
                             value=next_state.value + reward)
    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)
    if next_state.t >= level.end:
        next_state = replace(next_state, outcome="won" if next_state.value >= level.goal else "lost")
    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
    assert state.ctype in ("W", "M", "-") and state.ftype in ("W", "M", "-")
    assert state.t <= level.end
    assert state.value == sum((job.value if time <= job.due else max(1, job.value - 1))
                              for job, time in zip(level.jobs, state.done) if time >= 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))
        if key in ("1", "2", "3"):
            self.level_index = int(key) - 1
            self.reset()
            self.gated = True
            return
        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 4", 10)
        state = self.state
        p.text(12, 27, self.level.title + "  T=" + str(state.t), 7)
        p.text(330, 27, "END " + str(self.level.end) + "  VALUE " + str(state.value) + "/" + str(self.level.goal), 11)
        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.text(22, 111, "MATERIAL " + state.ctype, 6)
        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(346, 111, "MATERIAL " + state.ftype, 6)
        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)
        p.text(300, 144, "TYPE", 13)
        p.text(358, 144, "VALUE  DUE", 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)
            p.text(300, y, job.kind, 7)
            p.text(358, y, str(job.value) + "      " + str(job.due), 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, "DUE TURN IS FULL VALUE; LATE = MAX(1, VALUE-1).", 6)
        p.text(12, 340, self.message.upper() + " / " + state.outcome.upper(), 11)
        p.text(12, 359, "LEFT/RIGHT SELECT  Z/SPACE DO  R RESET  ESC QUIT", 7)
        p.text(12, 376, "1/2/3 LEVEL", 13)


DEMO_ACTIONS = (('cut', 0), ('finish',), ('cut', 1), ('finish',), ('clean', 'C'), ('cut', 2), ('clean', 'F'), ('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())