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

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 = ()
    actions: int = 0
    swaps: int = 0
    result: str = ''
    events: tuple = ()

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

def initial(board):
    return State(tuple(Unit(pos,3,i) for i,pos in enumerate(board.allies)))

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 legal(board,state):
    if state.result:
        return ()
    out=[]
    busy=occupied(state)|set(board.walls)|{board.flag}
    for i,unit in enumerate(state.allies):
        if unit.hp<=0:
            continue
        for dx,dy in DIRS:
            dest=(unit.pos[0]+dx,unit.pos[1]+dy)
            if inside(dest) and dest not in busy:
                out.append(Action('move',i,dest))
        for j,other in enumerate(state.allies):
            if i!=j and other.hp>0 and dist(unit.pos,other.pos)==1:
                out.append(Action('swap',i,(j,0)))
    return tuple(out)


def step(board,state,action):
    if action not in legal(board,state):
        return state
    allies=list(state.allies)
    enemies=list(state.enemies)
    swaps=state.swaps
    if action.kind=="move":
        allies[action.actor]=replace(allies[action.actor],pos=action.target)
    elif action.kind=="swap":
        j=action.target[0]
        left,right=allies[action.actor],allies[j]
        allies[action.actor]=replace(left,role=right.role)
        allies[j]=replace(right,role=left.role)
        swaps+=1
    result=""
    return replace(state,allies=tuple(allies),enemies=tuple(enemies),actions=state.actions+1,swaps=swaps,result=result,events=(action.kind,))

STAGE = 2
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])

    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,"PRACTICE / NO ENEMY TURN",6)
        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,"ARROWS SELECT THE LIST; THEY DO NOT DIRECTLY MOVE A BODY",7)

DEMO_ACTIONS = (('swap', 0, (2, 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)