# 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。原作Bの仕組みを段階的に読む。"""
from dataclasses import dataclass, replace, asdict

from collections import deque

import hashlib, json

VERSION = '085.1.0'

ROLES = ('槍', '弓', '盾')

NAMES = ('アキ', 'フミ', 'ルウ')

DIRS = ((0, -1), (-1, 0), (0, 1), (1, 0))

@dataclass(frozen=True)
class Board:
    id: str
    title: str
    flag: tuple
    allies: tuple
    enemies: tuple
    walls: tuple

BOARDS = (Board('A', '弓の場所', (0, 2), ((1, 2), (1, 3), (0, 3)), ((5, 1), (5, 2), (5, 3)), ((3, 1), (3, 3))), Board('B', '曲がった入口', (0, 0), ((1, 0), (2, 1), (1, 1)), ((0, 1), (5, 2), (4, 4)), ((3, 0), (3, 1), (1, 3), (2, 3))), Board('C', 'ふたつの門', (2, 4), ((2, 3), (3, 3), (3, 2)), ((1, 4), (3, 0), (5, 2)), ((1, 1), (2, 1), (4, 1), (4, 3))))

@dataclass(frozen=True)
class Unit:
    pos: tuple
    hp: int
    role: int = 0

@dataclass(frozen=True)
class State:
    allies: tuple
    enemies: tuple
    flag_hp: int = 3
    remaining: int = 2
    rounds: int = 0
    result: str = ''
    swaps: int = 0
    actions: int = 0
    events: tuple = ()

@dataclass(frozen=True)
class Action:
    kind: str
    actor: int = 0
    target: tuple = (0, 0)

def initial(b):
    return State(tuple((Unit(p, 3, i) for i, p in enumerate(b.allies))), tuple((Unit(p, 2) for p in b.enemies)))

def digest(s):
    return hashlib.sha256(json.dumps(asdict(s), ensure_ascii=False, sort_keys=True).encode('utf-8')).hexdigest()

def dist(a, b):
    return abs(a[0] - b[0]) + abs(a[1] - b[1])

def inside(p):
    return 0 <= p[0] < 6 and 0 <= p[1] < 5

def occupied(s):
    return {u.pos for u in s.allies + s.enemies if u.hp > 0}

def protected(s, pos):
    return any((u.hp > 0 and u.role == 2 and (dist(u.pos, pos) <= 1) for u in s.allies))

def can_attack(b, s, actor, target):
    u = s.allies[actor]
    d = dist(u.pos, target)
    if u.hp <= 0:
        return False
    if u.role != 1:
        return d == 1
    if not 2 <= d <= 3 or (u.pos[0] != target[0] and u.pos[1] != target[1]):
        return False
    dx = (target[0] > u.pos[0]) - (target[0] < u.pos[0])
    dy = (target[1] > u.pos[1]) - (target[1] < u.pos[1])
    return all(((u.pos[0] + dx * i, u.pos[1] + dy * i) not in set(b.walls) | occupied(s) for i in range(1, d)))

def legal(b, s, compare=False):
    if s.result:
        return ()
    out = [Action('wait'), Action('end')]
    busy = occupied(s) | set(b.walls) | {b.flag}
    for i, u in enumerate(s.allies):
        if u.hp <= 0:
            continue
        for dx, dy in DIRS:
            p = (u.pos[0] + dx, u.pos[1] + dy)
            if inside(p) and p not in busy:
                out.append(Action('move', i, p))
        for e in s.enemies:
            if e.hp > 0 and can_attack(b, s, i, e.pos):
                out.append(Action('attack', i, e.pos))
        for j, v in enumerate(s.allies):
            if j == i or dist(u.pos, v.pos) != 1:
                continue
            if v.hp > 0 and (not compare):
                out.append(Action('swap', i, (j, 0)))
            if v.hp == 0 and v.pos not in occupied(s):
                out.append(Action('rescue', i, (j, 0)))
    return tuple(out)

def enemy_move(b, s, idx):
    """Shortest path to flag, deterministic up/left/down/right tie order."""
    start = s.enemies[idx].pos
    blocked = occupied(s) | set(b.walls)
    blocked.discard(start)
    q = deque([(start, ())])
    seen = {start}
    while q:
        p, path = q.popleft()
        if p == b.flag:
            return path[0] if path else None
        for dx, dy in DIRS:
            n = (p[0] + dx, p[1] + dy)
            if inside(n) and n not in seen and (n not in blocked):
                seen.add(n)
                q.append((n, path + (n,)))
    return None

def enemy_phase(b, s):
    notes = []
    for i in range(len(s.enemies)):
        e = s.enemies[i]
        if e.hp <= 0:
            continue
        nearby = [j for j, u in enumerate(s.allies) if u.hp > 0 and dist(e.pos, u.pos) == 1]
        if dist(e.pos, b.flag) == 1:
            dmg = 0 if protected(s, b.flag) else 1
            s = replace(s, flag_hp=max(0, s.flag_hp - dmg))
            notes.append(f'敵{i + 1} → 旗 {dmg}ダメージ')
        elif nearby:
            j = min(nearby, key=lambda n: (s.allies[n].hp, n))
            u = s.allies[j]
            dmg = 0 if protected(s, u.pos) else 1
            units = list(s.allies)
            units[j] = replace(u, hp=max(0, u.hp - dmg))
            s = replace(s, allies=tuple(units))
            notes.append(f'敵{i + 1} → {NAMES[j]} {dmg}ダメージ')
        else:
            p = enemy_move(b, s, i)
            if p is not None:
                es = list(s.enemies)
                es[i] = replace(e, pos=p)
                s = replace(s, enemies=tuple(es))
                notes.append(f'敵{i + 1} → ({p[0] + 1},{p[1] + 1})へ移動・攻撃なし')
            else:
                notes.append(f'敵{i + 1} → 道がなく待機')
        if s.flag_hp == 0:
            return replace(s, result='旗が倒れた', events=tuple(notes))
    s = replace(s, rounds=s.rounds + 1, remaining=2)
    return replace(s, result='守り成功' if s.rounds >= 6 else '', events=tuple(notes))

def preview(b, s):
    if s.result:
        return ()
    return enemy_phase(b, s).events

def step(b, s, a, compare=False):
    if a not in legal(b, s, compare):
        return s
    units = list(s.allies)
    es = list(s.enemies)
    notes = []
    swaps = s.swaps
    if a.kind == 'move':
        units[a.actor] = replace(units[a.actor], pos=a.target)
        notes = [NAMES[a.actor] + 'が移動']
    elif a.kind == 'attack':
        j = next((j for j, e in enumerate(es) if e.hp > 0 and e.pos == a.target))
        damage = 2 if units[a.actor].role == 0 else 1
        es[j] = replace(es[j], hp=max(0, es[j].hp - damage))
        notes = [NAMES[a.actor] + 'の' + ROLES[units[a.actor].role] + f' → 敵{j + 1} {damage}ダメージ']
    elif a.kind == 'swap':
        j = a.target[0]
        u, v = (units[a.actor], units[j])
        units[a.actor] = replace(u, role=v.role)
        units[j] = replace(v, role=u.role)
        swaps += 1
        notes = [NAMES[a.actor] + 'と' + NAMES[j] + 'が役目交換 / HP・場所はそのまま']
    elif a.kind == 'rescue':
        j = a.target[0]
        units[j] = replace(units[j], hp=1)
        notes = [NAMES[a.actor] + 'が' + NAMES[j] + 'をHP1で救起']
    else:
        notes = ['残りの行動を待つ' if a.kind == 'end' else '一行動待つ']
    n = replace(s, allies=tuple(units), enemies=tuple(es), swaps=swaps, remaining=0 if a.kind == 'end' else s.remaining - 1, actions=s.actions + 1, events=tuple(notes))
    if all((e.hp == 0 for e in es)):
        return replace(n, result='全敵撃破')
    if n.remaining == 0:
        e = enemy_phase(b, n)
        n = replace(e, events=tuple(notes) + e.events)
    return n

STAGE = 4
BOARD = BOARDS[1]
LABELS = ("AKI","FUMI","RUU")
ROLE_LABELS = ("SPEAR","BOW","SHIELD")
ROLE_COLORS = (8,10,11)


def action_label(action):
    if action.kind in ('wait','end'):
        return 'WAIT ONE ACTION' if action.kind=='wait' else 'END TEAM TURN'
    actor=LABELS[action.actor]
    if action.kind in ('swap','rescue'):
        return action.kind.upper()+' '+actor+' / '+LABELS[action.target[0]]
    return action.kind.upper()+' '+actor+' -> '+str(action.target)


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

    def reset(self):
        self.state=initial(BOARD)
        self.selected=0
        self.gated=False
        self.message="READY"

    def choices(self):
        return legal(BOARD,self.state)

    def commit(self,action):
        before=self.state
        self.state=step(BOARD,self.state,action)
        self.selected=0
        self.message=action_label(action) if self.state!=before else "REFUSED"

    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))
        self.gated=True
        choices=self.choices()
        if not choices:
            return
        if key in ('LEFT','RIGHT'):
            self.selected=(self.selected+(-1 if key=='LEFT' else 1))%len(choices)
        elif key=='Z':
            self.commit(choices[self.selected])
        elif key=='E':
            self.commit(Action('end'))

    def draw(self,p):
        state=self.state
        p.cls(1)
        p.text(12,8,'ROLE FORTRESS / STAGE '+str(STAGE)+' / BOARD '+BOARD.id,10)
        p.text(12,23,'ACTIONS '+str(state.actions)+' / SWAPS '+str(state.swaps),7)
        for y in range(5):
            for x in range(6):
                px=24+x*34;py=60+y*34
                p.rect(px,py,32,32,5 if (x,y) in BOARD.walls else 1)
                p.rectb(px,py,32,32,13)
                if protected(state,(x,y)):
                    p.rectb(px+2,py+2,28,28,11)
        for x in range(6):p.text(38+x*34,46,str(x),6)
        for y in range(5):p.text(12,74+y*34,str(y),6)
        fx,fy=BOARD.flag;p.tri(36+fx*34,66+fy*34,36+fx*34,83+fy*34,49+fx*34,75+fy*34,10)
        for n,unit in enumerate(state.enemies):
            x,y=unit.pos;p.rect(33+x*34,69+y*34,16,16,8 if unit.hp>0 else 5)
            p.text(37+x*34,74+y*34,str(n+1),7)
        for n,unit in enumerate(state.allies):
            x,y=unit.pos;color=ROLE_COLORS[unit.role] if unit.hp>0 else 6
            if unit.hp>0:
                p.circ(40+x*34,76+y*34,10,color)
                p.text(38+x*34,74+y*34,'AFR'[n],1)
            else:
                # A fallen body never hides a living unit occupying this cell.
                p.text(26+x*34+n*6,62+y*34,'AFR'[n],6)
            p.text(12,248+n*16,LABELS[n]+' '+str(unit.pos)+' HP'+str(unit.hp)+' '+ROLE_LABELS[unit.role],color)
        for n,enemy in enumerate(state.enemies):p.text(12,307+n*14,'ENEMY '+str(n+1)+' '+str(enemy.pos)+' HP'+str(enemy.hp),8 if enemy.hp>0 else 6)
        choices=self.choices()
        p.text(256,46,'LEGAL ACTIONS '+str(len(choices)),7)
        start=(self.selected//7)*7
        for n,action in enumerate(choices[start:start+7],start):
            y=65+(n-start)*22
            if n==self.selected:p.rect(252,y-4,328,19,2);p.rectb(252,y-4,328,19,10)
            p.text(260,y,action_label(action),7)
        p.text(256,230,'TEAM LEFT '+str(state.remaining)+' / ROUND '+str(state.rounds)+' / FLAG HP '+str(state.flag_hp),11)
        if not state.result:
            future=enemy_phase(BOARD,state)
            p.text(256,247,'IF ENEMIES NOW: FLAG '+str(future.flag_hp)+' / ALLY HP '+','.join(str(u.hp) for u in future.allies),6)
        result={'守り成功':'HELD SIX ROUNDS','全敵撃破':'ALL ENEMIES DEFEATED','旗が倒れた':'FLAG LOST'}.get(state.result,'PLAY')
        p.text(256,264,result,10)
        p.text(12,363,self.message,14)
        p.text(12,382,'LEFT/RIGHT: CHOOSE ACTION   Z: DO IT   R: RESET   ESC: QUIT',7)
        p.text(12,398,"E: END TEAM TURN   TWO ACTIONS ARE SHARED BY THE WHOLE TEAM",7)

DEMO_ACTIONS = (('swap', 0, (2, 0)), ('end', 0, (0, 0)), ('attack', 2, (0, 1)), ('end', 0, (0, 0)), ('end', 0, (0, 0)), ('end', 0, (0, 0)), ('end', 0, (0, 0)), ('end', 0, (0, 0)))


def make_demo_keys():
    app=App();keys=[]
    def tap(key):
        keys.append(key)
        app.control({key},{key});app.control(set(),set())
    for kind,actor,target in DEMO_ACTIONS:
        desired=Action(kind,actor,target)
        choices=app.choices()
        if desired not in choices:
            raise ValueError('The QA route does not fit this modified board.')
        index=choices.index(desired)
        while app.selected!=index:tap('RIGHT')
        tap('Z')
    return tuple(keys)


def main(headless=False,screenshot=None,demo=False):
    import pyxel as p
    p.init(592,420,title='Role fortress lesson '+str(STAGE),fps=30,display_scale=2,headless=headless)
    app=App();sequence=make_demo_keys() if demo else ()
    keys={'LEFT':p.KEY_LEFT,'RIGHT':p.KEY_RIGHT,'Z':p.KEY_Z,'X':p.KEY_X,'R':p.KEY_R,'E':p.KEY_E}
    limit=max(20,2*len(sequence)+8)
    def update():
        if p.btnp(p.KEY_ESCAPE) or headless and p.frame_count>=limit:
            if screenshot:
                app.draw(p);p.screen.save(screenshot,1)
            p.quit()
        if headless:
            if demo:
                index=p.frame_count//2
                pressed={sequence[index]} if p.frame_count%2==0 and index<len(sequence) 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))


if __name__=='__main__':
    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')
    args=parser.parse_args()
    if args.demo_checkpoint and not args.headless_smoke:parser.error('demo requires headless mode')
    if args.screenshot and not args.headless_smoke:parser.error('screenshot requires headless mode')
    if args.screenshot and (Path(args.screenshot).suffix.lower()!='.png' or not Path(args.screenshot).parent.is_dir() or Path(args.screenshot).exists()):parser.error('use a new .png in an existing directory')
    if args.screenshot:args.screenshot=str(Path(args.screenshot).resolve())
    main(args.headless_smoke,args.screenshot,args.demo_checkpoint)