# 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。日本語の解説はLESSON.md。1ファイルで動作。"""
from dataclasses import dataclass, replace, asdict
from collections import deque
from functools import lru_cache

# マスの順序も仕様：現在、上、下、左、右。同点なら先を選ぶ。
DIRS = ((0, -1), (0, 1), (-1, 0), (1, 0))
ACTIONS = dict(zip("UDLR", DIRS))

@dataclass(frozen=True)
class Source:
    id: int
    pos: tuple
    birth: int
    life: int
    kind: str

@dataclass(frozen=True)
class Level:
    title: str
    width: int
    height: int
    walls: tuple
    start: tuple
    dog: tuple
    patrol: tuple
    doors: tuple
    bread: int
    old: tuple = ()

# 原作Cと同じ配置。最初の灰色パンは別の店の残りで、所持数に含めない。
BOARD = Level("近い古いパン、遠い新しいパン", 11, 7,
              ((5, 2), (5, 3), (5, 4)), (4, 1), (7, 2),
              ((7, 2), (8, 2), (8, 4)), ((2, 1), (8, 1), (9, 5)), 5,
              (Source(0, (7, 4), -3, 8, "old"),))

@dataclass(frozen=True)
class Memory:
    investigated: tuple = ()
    patrol_index: int = 0

@dataclass(frozen=True)
class World:
    player: tuple
    dog: tuple
    bread: int
    delivered: int = 0
    sources: tuple = ()
    turn: int = 0
    next_id: int = 1
    chew: int = 0
    memory: Memory = Memory()
    status: str = "play"
    reason: str = "ready"
    picked: int = -1
    bait_used: int = 0
    gifts: int = 0

def initial(b=BOARD):
    return World(b.start, b.dog, b.bread, sources=b.old)

def floor(b, p):
    return 0 < p[0] < b.width-1 and 0 < p[1] < b.height-1 and p not in b.walls

@lru_cache(None)
def distances(b, start):
    # 幅優先探索。壁を通らず、床の最短歩数を辞書へ記録する。
    distance = {start: 0}
    queue = deque([start])
    while queue:
        x, y = queue.popleft()
        for dx, dy in DIRS:
            cell = (x+dx, y+dy)
            if floor(b, cell) and cell not in distance:
                distance[cell] = distance[(x, y)]+1
                queue.append(cell)
    return distance

def active(s):
    # 年齢が寿命と同じになった手で消える。フレーム数は使わない。
    return tuple(src for src in s.sources if s.turn-src.birth < src.life)

@dataclass(frozen=True)
class Record:
    birth: int
    distance: int
    source_id: int

def scent(b, s, cell):
    records = []
    for src in active(s):
        if src.id in s.memory.investigated:
            continue
        distance = distances(b, src.pos).get(cell, 999)
        if distance <= 4:
            records.append(Record(src.birth, distance, src.id))
    # 各セルで、新しい誕生日→短い距離→小さいID。
    return min(records, key=lambda r: (-r.birth, r.distance, r.source_id), default=None)

@dataclass(frozen=True)
class Observation:
    cells: tuple
    records: tuple
    patrol_step: tuple
    chew: int

@dataclass(frozen=True)
class Intent:
    cell: tuple
    source_id: int = -1
    reason: str = "patrol"

def observation(b, s):
    cells = (s.dog,)+tuple((s.dog[0]+dx, s.dog[1]+dy) for dx, dy in DIRS
                          if floor(b, (s.dog[0]+dx, s.dog[1]+dy)))
    target = b.patrol[s.memory.patrol_index]
    distance = distances(b, target)
    patrol_step = min(cells, key=lambda cell: distance.get(cell, 999))
    return Observation(cells, tuple(scent(b, s, cell) for cell in cells), patrol_step, s.chew)

def policy(o):
    if o.chew:
        return Intent(o.cells[0], reason="chew")
    # 追跡と巡回を分ける。追跡先は観測にある現在セル・隣接セルのみ。
    found = [(index, r) for index, r in enumerate(o.records) if r is not None]
    if found:
        index, r = min(found, key=lambda item: (-item[1].birth, item[1].distance,
                                               item[0], item[1].source_id))
        return Intent(o.cells[index], r.source_id, "scent")
    return Intent(o.patrol_step)

def preview(b, s):
    return policy(observation(b, s))

def step(b, s, action):
    if s.status != "play":
        return s
    player, turn, bread = s.player, s.turn+1, s.bread
    sources, delivered, next_id = s.sources, s.delivered, s.next_id
    bait_used, reason = s.bait_used, "walk"
    if action in ACTIONS:
        dx, dy = ACTIONS[action]
        player = (player[0]+dx, player[1]+dy)
        if not floor(b, player):
            return s
    elif action == "Z":
        if bread <= 0:
            return s
        delivery = player in b.doors and not delivered & (1 << b.doors.index(player))
        bread -= 1
        if delivery:
            delivered |= 1 << b.doors.index(player)
            reason = "delivery"
        else:
            bait_used += 1
            reason = "bait"
        sources += (Source(next_id, player, turn, 6 if delivery else 8,
                           "delivery" if delivery else "bait"),)
        next_id += 1
    elif action == "W":
        reason = "wait"
    else:
        return s
    s = replace(s, player=player, turn=turn, bread=bread, delivered=delivered,
                sources=sources, next_id=next_id, bait_used=bait_used, reason=reason)
    # 時刻更新→期限切れを除去→旧犬位置で観測→判断→1歩移動。
    s = replace(s, sources=active(s))
    o = observation(b, s)
    intent = policy(o)
    dog, chew, memory, gifts = intent.cell, max(0, s.chew-1), s.memory, s.gifts
    if not o.chew:  # 1→0になる手も、まだ食事中。追加のパンを取らない。
        here = [src for src in s.sources if src.pos == dog and src.id not in memory.investigated]
        edible = [src for src in here if src.kind != "delivery"]
        if edible:
            src = min(edible, key=lambda src: (-src.birth, src.id))
            s = replace(s, sources=tuple(other for other in s.sources if other.id != src.id))
            chew, reason = 2, "eat"
        elif dog == player and bread > 0:
            bread -= 1
            gifts += 1
            chew, reason = 2, "contact"
        elif here:
            src = min(here, key=lambda src: (-src.birth, src.id))
            memory = replace(memory, investigated=memory.investigated+(src.id,))
            chew, reason = 2, "inspect"
        if dog == b.patrol[memory.patrol_index]:
            memory = replace(memory, patrol_index=(memory.patrol_index+1) % len(b.patrol))
    status = "won" if delivered == 7 else "short" if bread < 3-delivered.bit_count() else "play"
    if status == "won":
        reason = "complete"
    elif status == "short":
        reason = "short"
    return replace(s, dog=dog, bread=bread, chew=chew, memory=memory, status=status,
                   reason=reason, picked=intent.source_id, gifts=gifts)

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

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

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

KEY_ACTION = {"UP": "U", "DOWN": "D", "LEFT": "L", "RIGHT": "R", "SPACE": "W", "Z": "Z"}
CONTROL_KEYS = tuple(KEY_ACTION)+("R", "X")

class App:
    def __init__(self):
        self.locked = False
        self.reset()

    def reset(self):
        self.timeline = Timeline(initial())
        self.message = "READY"

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

    def control(self, pressed, held):
        # R直後は、全操作キーを離したフレームを1回通ってから再受付。
        if self.locked:
            if not held.intersection(CONTROL_KEYS):
                self.locked = False
            return
        if "R" in pressed:
            self.reset()
            self.locked = True
            return
        # 原作通りXは行動キーより優先。X+矢印も一手戻すだけ。
        if "X" in pressed:
            self.timeline = undo(self.timeline)
            self.message = "UNDO / ALL STATE RESTORED"
            return
        if self.state.status != "play":
            return
        chosen = [KEY_ACTION[key] for key in KEY_ACTION if key in pressed]
        action_held = sum(key in held for key in KEY_ACTION)
        # 新しく押したキーだけでなく、押し続けている行動キーも数える。
        if len(chosen) == 1 and action_held == 1:
            before = self.state
            self.timeline = apply(BOARD, self.timeline, chosen[0])
            self.message = self.state.reason.upper() if self.state != before else "BLOCKED / NO TURN"

    def update(self, p):
        # 実Pyxelもテストも、このポーリング入口とcontrolを通る。
        pressed = {key for key in CONTROL_KEYS if p.btnp(getattr(p, "KEY_"+key))}
        held = {key for key in CONTROL_KEYS if p.btn(getattr(p, "KEY_"+key))}
        self.control(pressed, held)

    def draw(self, p):
        s, b = self.state, BOARD
        p.cls(1)
        p.text(12, 9, "WATCHDOG BAKERY / STAGE 5", 10)
        p.text(12, 23, "FULL STATE HISTORY / UNDO", 6)
        p.text(12, 38, "TURN "+str(s.turn)+"   BREAD "+str(s.bread), 7)
        p.text(164, 38, "DELIVERED "+str(s.delivered.bit_count())+"/3  "+s.status.upper(), 11)
        for y in range(b.height):
            for x in range(b.width):
                cell = (x, y)
                px, py = 12+x*28, 61+y*28
                if not floor(b, cell):
                    p.rect(px, py, 27, 27, 5)
                    p.line(px+2, py+9, px+24, py+9, 13)
                    p.line(px+2, py+19, px+24, py+19, 13)
                    continue
                p.rectb(px, py, 27, 27, 13)
                rec = scent(b, s, cell)
                if rec:
                    for dot in range(max(1, 5-rec.distance)):
                        p.pset(px+3+dot*4, py+23, 10)
                if cell in b.doors:
                    bit = 1 << b.doors.index(cell)
                    color = 11 if s.delivered & bit else 9
                    p.rect(px+6, py+5, 15, 14, color)
                    p.text(px+12, py+7, str(b.doors.index(cell)+1), 1)
                    p.line(px+3, py+20, px+24, py+20, 7)
        for src in active(s):
            x, y = 26+src.pos[0]*28, 75+src.pos[1]*28
            color = 13 if src.kind == "old" else 10
            p.circ(x-5, y+3, 4, color)
            p.rect(x-9, y+3, 9, 4, color)
            p.text(x-11, y-11, str(src.life-(s.turn-src.birth)), color)
        x, y = 26+s.player[0]*28, 75+s.player[1]*28
        p.circ(x, y-6, 4, 12)
        p.rect(x-4, y-1, 9, 9, 12)
        p.rect(x+4, y+1, 5, 5, 4)
        x, y = 26+s.dog[0]*28, 75+s.dog[1]*28
        p.rect(x-7, y-2, 13, 9, 9)
        p.circ(x+5, y-5, 5, 9)
        p.tri(x+2, y-8, x+4, y-13, x+7, y-7, 4)
        p.pset(x+7, y-6, 1)
        intent = preview(b, s)
        mode = "CHASE" if intent.reason == "scent" else intent.reason.upper()
        p.text(12, 282, "DOG: "+mode+"  SOURCE: "+str(intent.source_id), 6)
        p.text(12, 266, self.message, 10)
        p.text(12, 298, "ARROWS MOVE  SPACE WAIT  Z DELIVER/BAIT", 7)
        p.text(12, 310, "R RESET   X UNDO   ESC QUIT", 6)

# 合成キー列を実updateへ渡す検証用。OSのキー操作や人の試遊ではない。
DEMO_ACTIONS = "ZLLZRRRRRRZDDDDRZ"

def main(headless=False, screenshot=None, checkpoint=False, actions=None, state_json=None):
    from pathlib import Path
    import json
    import pyxel as p
    if screenshot:
        screenshot = str(Path(screenshot).resolve())
        Path(screenshot).parent.mkdir(parents=True, exist_ok=True)
    state_json = str(Path(state_json).resolve()) if state_json else None
    app = App()
    plan = actions if actions is not None else DEMO_ACTIONS if checkpoint else ""
    inverse = {action: key for key, action in KEY_ACTION.items()}
    inverse.update({"X": "X", "!": "R"})
    if any(action not in inverse or inverse[action] not in CONTROL_KEYS for action in plan):
        raise ValueError("Unsupported action in --actions")
    frames = max(4, 2*len(plan)+3)
    p.init(332, 328, title="039 Watchdog Bakery / 5", fps=30, display_scale=2,
           headless=headless)

    class FrameInput:
        def __init__(self, pressed):
            self.pressed = pressed
        def __getattr__(self, name):
            if name.startswith("KEY_"):
                return name[4:]
            raise AttributeError(name)
        def btnp(self, key):
            return key in self.pressed
        def btn(self, key):
            return key in self.pressed

    def update():
        if headless:
            index = p.frame_count//2
            pressed = {inverse[plan[index]]} if p.frame_count%2 == 0 and index < len(plan) else set()
            app.update(FrameInput(pressed))
        else:
            app.update(p)

    def draw():
        app.draw(p)
        if headless and p.frame_count >= frames:
            if screenshot:
                p.screenshot(screenshot, scale=2)
            if state_json:
                path = Path(state_json)
                path.parent.mkdir(parents=True, exist_ok=True)
                path.write_text(json.dumps({"stage": 5, "actions": plan,
                    "state": asdict(app.state), "message": app.message,
                    "evidence": "Real Pyxel headless drawing; synthetic update input; no OS input"},
                    ensure_ascii=False, indent=2), encoding="utf-8")
            p.quit()
    p.run(update, draw)
    return app

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description="Independent Pyxel tutorial / Python 3.11+")
    parser.add_argument("--headless-smoke", action="store_true")
    parser.add_argument("--demo-checkpoint", action="store_true")
    parser.add_argument("--actions")
    parser.add_argument("--screenshot")
    parser.add_argument("--state-json")
    args = parser.parse_args()
    main(args.headless_smoke, args.screenshot, args.demo_checkpoint, args.actions, args.state_json)