# 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.
# 教材画面は標準ASCIIと図形のみ。字形データは同梱しません。

"""返せる弾丸・教材4/5。日本語説明はLESSON.md。"""
from dataclasses import dataclass, field, asdict
from copy import deepcopy
import math
import json
import hashlib

@dataclass(frozen=True)
class Level:
    name: str
    start: float
    speed: float
    schedule: tuple
    ring_ticks: int = 12

# 原作A。一周180tick、中央へ30/100tickで発射。数値は変えない。
LEVEL = Level('一発を返す', 80.0, 1.4, ((30, 80), (100, 80)))

@dataclass
class Bullet:
    ident: int
    owner: str
    x: float
    y: float
    vx: float
    vy: float
    source: str = 'core'

@dataclass(frozen=True)
class Input:
    dx: int = 0
    z: bool = False
    x: bool = False

@dataclass(frozen=True)
class Observation:
    player_x: float
    advertised_lane: float | None
    fire_tick: int
    ident: int

@dataclass(frozen=True)
class Intent:
    lane: float
    fire_tick: int
    ident: int
    aimed: bool

@dataclass
class Memory:
    committed: list = field(default_factory=list)


def policy(obs, memory):
    # Aは固定レーン。Noneを使う原作Cも同じ関数で「予告時だけ」狙う。
    lane = obs.player_x if obs.advertised_lane is None else obs.advertised_lane
    return Intent(lane, obs.fire_tick, obs.ident, obs.advertised_lane is None)


def segment_distance(x1, y1, x2, y2, cx, cy):
    dx, dy = x2-x1, y2-y1
    length2 = dx*dx + dy*dy
    t = 0 if length2 == 0 else max(0, min(1, ((cx-x1)*dx+(cy-y1)*dy)/length2))
    return math.hypot(x1+t*dx-cx, y1+t*dy-cy)


def touches(x1, y1, x2, y2, cx, cy, radius):
    return segment_distance(x1, y1, x2, y2, cx, cy) <= radius

@dataclass
class WorldState:
    tick: int = 0
    x: float = 80.0
    hp: int = 3
    core_hp: int = 6
    invuln: int = 0
    charge: int = 3
    ring: int = 0
    cooldown: int = 0
    open: int = 0
    shot_cooldown: int = 0
    xheld: bool = False
    memory: Memory = field(default_factory=Memory)
    bullets: list = field(default_factory=list)
    next_id: int = 1
    status: str = 'playing'
    events: tuple = ()


def initial(level=LEVEL):
    return WorldState(x=level.start)


def step(level, old, inp=Input()):
    if old.status != 'playing':
        return deepcopy(old)
    s = deepcopy(old)  # 前の状態を残す。テストでも比較しやすい。
    events = []
    oldx = s.x
    for name in ('invuln', 'ring', 'cooldown', 'open', 'shot_cooldown'):
        setattr(s, name, max(0, getattr(s, name)-1))
    # 使用時刻と独立。tick=90,180,... の更新開始時に回復する。
    if s.tick and s.tick % 90 == 0:
        s.charge = min(3, s.charge+1)
        events.append(('recharge', s.charge))
    cycle = s.tick // 180
    for k, (fire, lane) in enumerate(level.schedule):
        if s.tick % 180 == fire-30:
            observation = Observation(old.x, lane, cycle*180+fire, cycle*len(level.schedule)+k)
            intent = policy(observation, s.memory)
            s.memory.committed.append(intent)
            events.append(('warning', intent.ident, intent.lane, intent.fire_tick, intent.aimed))
    for intent in s.memory.committed:
        if intent.fire_tick == s.tick:
            s.bullets.append(Bullet(s.next_id, 'enemy', intent.lane, 18.0, 0.0, level.speed))
            s.next_id += 1
            events.append(('fire', intent.ident, intent.lane, s.next_id-1))
    s.memory.committed = [intent for intent in s.memory.committed if intent.fire_tick >= s.tick]
    s.x = max(16, min(144, s.x + max(-1, min(1, inp.dx))*1.8))
    if inp.x and not s.xheld and s.charge > 0 and s.cooldown == 0:
        s.ring = level.ring_ticks
        s.cooldown = 36
        s.charge -= 1
        events.append(('ring', s.charge))
    s.xheld = inp.x  # 押し続けても新しい押下ではない。
    if inp.z and s.shot_cooldown == 0:
        s.bullets.append(Bullet(s.next_id, 'player', s.x, 93.0, 0.0, -3.0, 'player'))
        s.next_id += 1
        s.shot_cooldown = 24
        events.append(('shot', s.next_id-1))
    alive = []
    core_damage = 0
    player_hit = False
    reflected = set()
    moved = [(b, b.x, b.y, b.x+b.vx, b.y+b.vy) for b in s.bullets]
    # 全弾の反射を先に決める。リストの順で防壁判定を変えない。
    for b, x, y, nx, ny in moved:
        if b.owner == 'enemy' and s.ring > 0 and touches(x-oldx, y-100, nx-s.x, ny-100, 0, 0, 20):
            b.owner = 'reflected'
            b.x, b.y = nx, ny
            dx, dy = 80-nx, 18-ny
            norm = math.hypot(dx, dy)
            b.vx, b.vy = 3*dx/norm, 3*dy/norm
            s.open = 75
            reflected.add(b.ident)
            events.append(('reflect', b.ident, b.source))
    for b, x, y, nx, ny in moved:
        if b.ident in reflected:
            alive.append(b)
            continue  # 今tickで向きが変わった弾を二度動かさない。
        if b.owner == 'enemy':
            if touches(x-oldx, y-100, nx-s.x, ny-100, 0, 0, 8):
                if s.invuln == 0:
                    player_hit = True
                events.append(('enemy_contact', b.ident))
                continue
        else:
            if b.owner == 'player' and s.open == 0 and y >= 32 and ny <= 32:
                events.append(('blocked', b.ident))
                continue
            if touches(x, y, nx, ny, 80, 18, 15 if b.owner == 'reflected' else 14):
                core_damage += 1
                events.append(('core_hit', b.ident, b.owner))
                continue
        b.x, b.y = round(nx, 6), round(ny, 6)
        if -8 <= b.x <= 168 and -8 <= b.y <= 116:
            alive.append(b)
        else:
            events.append(('missed', b.ident, b.owner))
    s.bullets = alive
    s.core_hp = max(0, s.core_hp-core_damage)
    if player_hit:
        s.hp -= 1
        s.invuln = 30
        events.append(('hurt',))
    if s.hp <= 0:
        s.status = 'lost'
        events.append(('lost',))
    elif s.core_hp <= 0:
        s.status = 'won'
        events.append(('won',))
    s.tick += 1
    s.events = tuple(events)
    return s


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

STAGE = 4
STAGE_TITLE = 'AUTHENTICATE THE BARRIER'

CONTROLS = ('LEFT', 'RIGHT', 'Z', 'X', 'R')

class Controller:
    def __init__(self):
        self.state = initial()
        self.gate = False
        self.reason = 'ready'
        self.last_events = ()

    def reset(self):
        self.state = initial()
        self.gate = True
        self.last_events = ()

    def tick(self, pressed=(), held=()):
        pressed = set(pressed) & set(CONTROLS)
        held = set(held) & set(CONTROLS)
        # 停止中/全解放待ちでもXを離した事実は覚える。原作と同じ順序。
        if 'X' not in held:
            self.state.xheld = False
        if self.gate:
            if not held:
                self.gate = False
            return  # 全解放フレーム自体も世界を進めない。
        if 'R' in pressed:
            self.reset()
            self.reason = 'reset'
            return
        if self.state.status != 'playing':
            return
        # 左右同時は0。Zは保持で連射し、Xの新押下はモデル側で判断する。
        inp = Input(int('RIGHT' in held)-int('LEFT' in held), 'Z' in held, 'X' in held)
        self.state = step(LEVEL, self.state, inp)
        self.reason = 'play'
        if self.state.events:
            self.last_events = self.state.events

    def snapshot(self):
        result = {'state': asdict(self.state), 'hash': state_hash(self.state),
                  'gate': self.gate, 'reason': self.reason, 'last_events': self.last_events}
        return result


def draw_board(pyxel, controller):
    s = controller.state
    pyxel.cls(1)
    pyxel.text(12, 10, f'RETURN BULLETS 036 / STEP {STAGE}', 10)
    pyxel.text(12, 24, STAGE_TITLE, 7)
    pyxel.rect(20, 44, 320, 236, 0)
    pyxel.rectb(20, 44, 320, 236, 5)
    def xy(x, y):
        return 20+x*2, 44+y*2
    for lane in (40, 80, 120):
        x, _ = xy(lane, 0)
        for y in range(84, 276, 12):
            pyxel.line(x, y, x, y+3, 1)
    if STAGE >= 2:
        for intent in s.memory.committed:
            x, _ = xy(intent.lane, 0)
            for y in range(110, 269, 10):
                pyxel.line(x, y, x, y+4, 9)
            pyxel.text(x+6, 146, 'FIRE IN '+str(max(0, intent.fire_tick-s.tick+1)), 9)
    cx, cy = xy(80, 18)
    pyxel.circ(cx, cy, 22, 2)
    pyxel.rect(cx-12, cy-12, 24, 24, 8 if STAGE>=4 else 5)
    pyxel.text(cx-8, cy-2, 'CORE', 7)
    if STAGE>=4:
        for j in range(6):
            pyxel.rect(138+j*14, 49, 10, 4, 8 if j<s.core_hp else 1)
        _, wy = xy(0, 32)
        pyxel.line(28, wy, 332, wy, 11 if s.open else 12)
        if not s.open:
            for x in range(28, 330, 10):
                pyxel.line(x, wy-3, x+4, wy+3, 12)
    pyxel.clip(21,45,318,234)
    for b in s.bullets:
        x,y = xy(b.x,b.y)
        if b.owner=='enemy':
            pyxel.circ(x,y,5,8)
            pyxel.circb(x,y,5,10)
        elif b.owner=='player':
            pyxel.tri(x,y-5,x-3,y+3,x+3,y+3,7)
        else:
            pyxel.tri(x,y-7,x-6,y,x,y+7,11)
            pyxel.tri(x,y-7,x+6,y,x,y+7,11)
    x,y = xy(s.x,100)
    if STAGE>=3 and s.ring:
        pyxel.circb(x,y,34,11)  # 原作の可視半径17。衝突判定半径は20のまま。
    if not getattr(s,'invuln',0) or s.tick%4<2:
        pyxel.tri(x,y-10,x-9,y+8,x+9,y+8,7)
        pyxel.rect(x-3,y-4,6,8,12)
    pyxel.clip()
    lines=[('WORLD TICK',str(s.tick),7),('STATUS',s.status.upper(),10),
           ('X POSITION',f'{s.x:.1f}',6),('SHOT WAIT',str(s.shot_cooldown),7)]
    if STAGE>=2: lines += [('PLAYER HP',str(s.hp)+'/3',8),('INVULNERABLE',str(s.invuln),6)]
    if STAGE>=3: lines += [('CHARGE',str(s.charge)+'/3',11),('RING ACTIVE',str(s.ring),11),
                           ('REUSE WAIT',str(s.cooldown),6),('RECHARGE IN',str(90-s.tick%90),6)]
    if STAGE>=4: lines += [('BARRIER',('OPEN '+str(s.open)) if s.open else 'CLOSED',11 if s.open else 12),('CORE HP',str(s.core_hp)+'/6',8)]
    for i,(label,value,color) in enumerate(lines):
        pyxel.text(356,48+i*17,label,13)
        pyxel.text(446,48+i*17,value,color)
    pyxel.text(12,292,'WHITE: PLAYER SHOT   RED: ENEMY   GREEN: AUTHENTICATED RETURN',7)
    if STAGE==1: message='LAB: movement + shots only. No damage, wall, enemies, or ending.'
    elif STAGE==2: message='LAB: telegraphs + enemy contact. No reflection, wall, or victory.'
    elif STAGE==3: message='LAB: reflection only. Returned bullets cannot damage the core yet.'
    else: message='GOAL: destroy core AND survive. Simultaneous zero HP is a loss.'
    pyxel.text(12,308,message,6)
    controls='LEFT/RIGHT: move  Z: hold fire  R: reset  ESC: quit'
    if STAGE>=3: controls='ARROWS: move  Z: hold fire  X: press reflect  R: reset  ESC: quit'
    pyxel.text(12,327,controls,7)
    if STAGE==5:
        pyxel.text(12,342,'SPACE: start  P: pause  M: title / after changes: release ALL keys',10)
    else:
        pyxel.text(12,342,'After reset: release ALL control keys. Fixed 30 ticks/second.',13)
    if s.status!='playing':
        pyxel.rect(62,160,236,44,1)
        pyxel.rectb(62,160,236,44,10)
        pyxel.text(92,172,'CORE DESTROYED / YOU SURVIVED' if s.status=='won' else 'PLAYER LOST / PRESS R',10)
        pyxel.text(92,188,'R: reset the same original case A',7)
    if getattr(controller,'title',False) or getattr(controller,'pause',False):
        pyxel.rect(38,122,284,94,1)
        pyxel.rectb(38,122,284,94,11)
        pyxel.text(52,134,'RETURN BULLETS' if controller.title else 'PAUSED: ALL GAME CLOCKS STOP',10)
        pyxel.text(52,153,'Return one red bullet to open the wall.',7)
        pyxel.text(52,170,'X lasts 12 ticks. Z fires while held.',7)
        pyxel.text(52,191,'SPACE: start' if controller.title else 'P: resume / release all keys after',11)


def demo_keys(controller, frame):
    """公開状態だけを見る技術QA用入力。人間の試遊ではない。"""
    if STAGE==5:
        if frame==1: return ('SPACE',)
        if frame==2: return ()
    s=controller.state
    if STAGE==1: return ('Z','RIGHT') if s.tick<10 else ('Z',)
    if STAGE==2: return ()
    keys=['Z']
    if any(b.owner=='enemy' and 73<=b.y<101 and abs(b.x-s.x)<21
           for b in s.bullets) and s.charge>0 and s.cooldown<=1 and not s.xheld:
        keys.append('X')
    return tuple(keys)


def main(headless=False, frames=180, demo=False, replay=None, screenshot=None, report=None):
    from pathlib import Path
    import pyxel  # 通常テスト/インポートにPyxelや画面を要求しない。
    controller=Controller()
    codes={name:getattr(pyxel,'KEY_'+name) for name in CONTROLS}
    frame=0
    old_held=set()
    changes=[]
    screenshot=str(Path(screenshot).resolve()) if screenshot else None
    report=str(Path(report).resolve()) if report else None
    def update():
        nonlocal frame,old_held
        frame+=1
        if headless:
            held=set(replay[frame-1]) if replay is not None and frame<=len(replay) else set(demo_keys(controller,frame) if demo else ())
            pressed=held-old_held
            old_held=held
        else:
            held={name for name,code in codes.items() if pyxel.btn(code)}
            pressed={name for name,code in codes.items() if pyxel.btnp(code)}
        before=controller.snapshot()
        controller.tick(pressed,held)
        if headless and (controller.state.events or before['reason']!=controller.reason):
            changes.append({'frame':frame,'pressed':sorted(pressed),'held':sorted(held),'snapshot':controller.snapshot()})
    def draw():
        draw_board(pyxel,controller)
        if headless and frame>=frames:
            if screenshot: pyxel.screenshot(screenshot,scale=1)
            if report:
                Path(report).write_text(json.dumps({'label':'Actual Pyxel headless; synthetic controller inputs; silent; not human play',
                    'stage':STAGE,'frames':frame,'demo':demo,'changes':changes,
                    'snapshot':controller.snapshot()},ensure_ascii=False,indent=2)+'\n',encoding='utf-8')
            pyxel.quit()
    pyxel.init(520,360,title=f'Return Bullets / Step {STAGE}',fps=30,headless=headless)
    pyxel.run(update,draw)


if __name__=='__main__':
    import argparse
    from pathlib import Path
    parser=argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--headless',action='store_true')
    parser.add_argument('--frames',type=int,default=180)
    parser.add_argument('--demo',action='store_true')
    parser.add_argument('--replay-json',help='UTF-8 JSON list of held-key lists, one per frame')
    parser.add_argument('--screenshot',help='PNG output; headless only, relative to current directory')
    parser.add_argument('--report',help='UTF-8 JSON output; headless only')
    args=parser.parse_args()
    if args.frames<12: parser.error('Use at least 12 frames')
    if not args.headless and (args.demo or args.replay_json or args.screenshot or args.report):
        parser.error('QA input and output options require --headless')
    if args.demo and args.replay_json: parser.error('Choose demo or replay, not both')
    replay=None
    if args.replay_json:
        try: replay=json.loads(Path(args.replay_json).read_text(encoding='utf-8'))
        except (OSError,ValueError) as error: parser.error(str(error))
        if not isinstance(replay,list) or any(not isinstance(keys,list) or any(not isinstance(k,str) or k not in CONTROLS for k in keys) for keys in replay):
            parser.error('Replay must be a list of lists of supported key names')
    main(args.headless,args.frames,args.demo,replay,args.screenshot,args.report)