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

"""地下商人の帰還・教材5/5。ゲーム内の金額は架空の規則です。"""
from dataclasses import dataclass, replace

KINDS = ('C', 'M', 'O')  # 晶石、きのこ、鉄鉱
BASE = ((2, 2, 2), (2, 3, 3))
TOOLS = {'C': 'L', 'M': 'K', 'O': 'T'}  # 灯、冷石、磁石
BONUS = ((4, 0, 0), (0, 2, 2))
BOX = ('C1', 'M1', 'M2', 'O1', 'L1', 'K1', 'T1')
STOCK = ((('M3', 1),), ())  # (商品ID, 買値)。一品は一度だけ。
GOAL = 10
PHASES = ('LOAD', 'MARKET 1', 'REST STOP', 'MARKET 2', 'HOME')


@dataclass(frozen=True)
class State:
    phase: int = 0
    cargo: tuple = ('', '', '', '')
    box: tuple = BOX
    coins: int = 0
    available: tuple = ('M3',)
    sold: tuple = ()
    bought: tuple = ()
    paid: int = 0
    outcome: str = "journey"


def valid_slot(slot):
    return type(slot) is int and 0 <= slot < 4


def enhanced(cargo, slot):
    if not valid_slot(slot) or not cargo[slot] or cargo[slot][0] not in TOOLS:
        return False
    tool = TOOLS[cargo[slot][0]]
    # 端を越えない。anyなので左右に道具があっても一度だけ加算。
    return any(0 <= other < 4 and cargo[other] and cargo[other][0] == tool
               for other in (slot-1, slot+1))


def price(cargo, slot, market=0):
    if not valid_slot(slot) or type(market) is not int or market not in (0, 1):
        return 0
    item = cargo[slot]
    if not item or item[0] not in KINDS:
        return 0
    kind = KINDS.index(item[0])
    return BASE[market][kind] + (BONUS[market][kind] if enhanced(cargo, slot) else 0)


def market_index(state):
    return 0 if state.phase == 1 else 1 if state.phase == 3 else None


def step(state, action):
    if state.outcome != "journey":
        return state, "HOME: R TO RESTART"
    shapes = {'load': 3, 'sell': 2, 'swap': 3, 'buy': 3, 'next': 1}
    if not isinstance(action, (tuple, list)) or not action:
        return state, "INVALID ACTION"
    op = action[0]
    if not isinstance(op, str) or op not in shapes or len(action) != shapes[op]:
        return state, "INVALID ACTION"
    if op in ("load", "buy") and not isinstance(action[2], str):
        return state, "INVALID ITEM"
    cargo = list(state.cargo)
    if op == "next":
        if state.phase == 3:
            outcome = "won" if state.coins >= GOAL else "returned"
            return replace(state, phase=4, outcome=outcome), outcome.upper()
        return replace(state, phase=state.phase+1), "NEXT STOP"
    if op == "load":
        if state.phase != 0:
            return state, "LOAD ONLY BEFORE DEPARTURE"
        slot, item = action[1:]
        if not valid_slot(slot):
            return state, "INVALID SLOT"
        if item == cargo[slot]:
            return state, "SAME ITEM"
        if item and item not in state.box:
            return state, "NOT IN BOX"
        pool = set(state.box)
        if cargo[slot]:
            pool.add(cargo[slot])
        if item:
            pool.remove(item)
        cargo[slot] = item
        return replace(state, cargo=tuple(cargo), box=tuple(x for x in BOX if x in pool)), "LOADED"
    if op == "swap":
        if state.phase not in (0, 2):
            return state, "NO SWAP IN A MARKET"
        i, j = action[1:]
        if not valid_slot(i) or not valid_slot(j):
            return state, "INVALID SLOT"
        if cargo[i] == cargo[j]:
            return state, "SAME ARRANGEMENT"
        cargo[i], cargo[j] = cargo[j], cargo[i]
        return replace(state, cargo=tuple(cargo)), "SWAPPED"
    market = market_index(state)
    if market is None:
        return state, "TRADE ONLY IN A MARKET"
    slot = action[1]
    if not valid_slot(slot):
        return state, "INVALID SLOT"
    if op == "sell":
        if not cargo[slot]:
            return state, "EMPTY SLOT"
        item = cargo[slot]
        # 売る直前の配置で値を決め、その枠だけ空にする。
        value = price(state.cargo, slot, market)
        cargo[slot] = ''
        return replace(state, cargo=tuple(cargo), coins=state.coins+value,
                       sold=state.sold+((item, market, value),)), "SOLD FOR "+str(value)
    if op == "buy":
        item = action[2]
        stock = dict(STOCK[market])
        if item not in stock or item not in state.available:
            return state, "OUT OF STOCK"
        if cargo[slot]:
            return state, "NEED AN EMPTY SLOT"
        cost = stock[item]
        if state.coins < cost:
            return state, "NOT ENOUGH COINS"
        cargo[slot] = item
        return replace(state, cargo=tuple(cargo), coins=state.coins-cost,
                       available=tuple(x for x in state.available if x != item),
                       bought=state.bought+(item,), paid=state.paid+cost), "BOUGHT "+item
    return state, "INVALID ACTION"


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


def apply(timeline, action):
    if action == ('undo',):
        if not timeline.history:
            return timeline, "NOTHING TO UNDO"
        return Timeline(timeline.history[-1], timeline.history[:-1]), "UNDO"
    state, message = step(timeline.state, action)
    if state == timeline.state:
        return timeline, message
    return Timeline(state, timeline.history+(timeline.state,)), message


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

    def reset(self):
        self.timeline = Timeline()
        self.selected = 0
        self.locked = False
        self.message = "CHOOSE A SLOT"
        self.anchor = None
        self.option = 0

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

    def options(self):
        state = self.state
        if state.phase == 0:
            return [('load', self.selected, item) for item in state.box]+[('load', self.selected, '')]
        if state.phase in (1, 3):
            return [('sell', self.selected)]+[('buy', self.selected, item) for item, cost in STOCK[market_index(state)]]
        return [('select_swap', self.selected)] if state.phase == 2 else [('finished',)]

    def label(self, action):
        if action[0] == 'load':
            return 'LOAD '+action[2] if action[2] else 'UNLOAD TO BOX'
        if action[0] == 'sell':
            return 'SELL / TOOLS GIVE 0'
        if action[0] == 'buy':
            item = action[2]
            return 'BUY '+item+' COST 1'+(' / SOLD OUT' if item not in self.state.available else '')
        if action[0] == 'select_swap':
            return 'CHOOSE SWAP END' if self.anchor is not None else 'CHOOSE SWAP START'
        return self.state.outcome.upper()

    def forecast(self):
        action = self.options()[self.option]
        if action[0] == 'select_swap':
            if self.anchor is None:
                return self.state
            action = ('swap', self.anchor, self.selected)
        return step(self.state, action)[0]

    def act(self, action):
        self.timeline, self.message = apply(self.timeline, action)
        self.anchor = None
        self.option = min(self.option, len(self.options())-1)

    def control(self, pressed, held):
        # 確定後も複数キー拒否後も全解放を待つ。Rもゲート中は待機。
        if self.locked:
            if not held:
                self.locked = False
            return
        if 'R' in pressed:
            self.reset()
            self.locked = True
            return
        if len(held) > 1 or len(pressed) > 1:
            self.locked = bool(held)
            return
        if len(pressed) != 1 or pressed != held:
            return
        key = next(iter(pressed))
        if key == 'X':
            if self.anchor is not None:
                self.anchor = None
                self.message = "SWAP CANCELLED"
            else:
                self.act(('undo',))
            self.locked = True
            return
        if self.state.outcome != 'journey':
            return
        if key in ('LEFT', 'RIGHT'):
            self.selected = (self.selected+(1 if key == 'RIGHT' else -1)) % 4
            self.option = min(self.option, len(self.options())-1)
        elif key == 'Z':
            if self.state.phase == 2:
                if self.anchor is None:
                    self.anchor = self.selected
                    self.message = "SELECT ANOTHER SLOT, THEN Z"
                else:
                    self.act(('swap', self.anchor, self.selected))
            else:
                self.act(self.options()[self.option])
        elif key == 'SPACE':
            self.act(('next',))
        elif key in ('UP', 'DOWN'):
            self.option = (self.option+(1 if key == 'DOWN' else -1)) % len(self.options())
        else:
            return
        self.locked = True

    def draw(self, p):
        state = self.state
        p.cls(1)
        p.text(8, 8, "CARGO MERCHANT / STAGE 5", 10)
        p.text(8, 22, PHASES[state.phase]+"   COINS "+str(state.coins)+" / GOAL "+str(GOAL), 7)
        p.text(8, 34, "BOX "+" ".join(state.box), 6)
        future = self.forecast()
        for slot, item in enumerate(state.cargo):
            x = 8+slot*90
            p.rect(x, 52, 84, 94, 2)
            p.rectb(x, 52, 84, 94, 10 if slot == self.selected else 5)
            if slot == self.anchor:
                p.rectb(x+2, 54, 80, 90, 11)
            p.text(x+5, 59, str(slot+1)+": "+(item or 'EMPTY'), 7)
            if item:
                color = {'C': 12, 'M': 8, 'O': 13, 'L': 10, 'K': 6, 'T': 9}[item[0]]
                p.circ(x+42, 87, 12, color)
                p.text(x+40, 85, item[0], 1)
                if enhanced(state.cargo, slot):
                    p.text(x+5, 103, "ENHANCED", 11)
            for market in range(2):
                quote = str(price(state.cargo, slot, market))
                if item and future.cargo[slot] == item:
                    after = price(future.cargo, slot, market)
                    if after != price(state.cargo, slot, market):
                        quote += " > "+str(after)
                p.text(x+5, 117+12*market, "M"+str(market+1)+" $"+quote, 7)
        p.text(8, 156, "L -> C   K -> M   T -> O / ADJACENT ONLY", 6)
        p.text(8, 168, "MARKET 1: "+"   ".join(k+" "+str(BASE[0][i])+"+"+str(BONUS[0][i]) for i, k in enumerate(KINDS)), 7)
        p.text(8, 180, "MARKET 2: "+"   ".join(k+" "+str(BASE[1][i])+"+"+str(BONUS[1][i]) for i, k in enumerate(KINDS)), 7)
        p.text(8, 198, 'UP/DOWN: '+self.label(self.options()[self.option]), 10)
        p.text(8, 210, 'OPTION '+str(self.option+1)+'/'+str(len(self.options()))+'   SPACE: NEXT STOP', 6)
        p.text(8, 224, "FORECAST COINS "+str(state.coins)+" > "+str(future.coins)+"   UNDO "+str(len(self.timeline.history)), 6)
        p.text(8, 242, self.message, 11)
        p.text(8, 256, "RESULT: "+state.outcome.upper()+" / UNSOLD CARGO IS KEPT", 7)
        p.text(8, 276, "LEFT/RIGHT SELECT   Z CONFIRM   R RESET", 7)
        p.text(8, 290, "X UNDO (OR CANCEL SWAP)   ESC QUIT / FICTIONAL GAME PRICES", 6)

DEMO_KEYS = ('Z', 'RIGHT', 'DOWN', 'DOWN', 'DOWN', 'Z', 'RIGHT', 'UP', 'UP', 'UP', 'Z', 'RIGHT', 'DOWN', 'DOWN', 'Z', 'SPACE', 'RIGHT', 'UP', 'Z', 'DOWN', 'Z', 'SPACE', 'RIGHT', 'Z', 'RIGHT', 'RIGHT', 'Z', 'SPACE', 'RIGHT', 'UP', 'Z', 'RIGHT', 'RIGHT', 'Z', 'SPACE')


def main(headless=False, screenshot=None, demo=False, preview=False):
    from pathlib import Path
    from dataclasses import asdict
    import json
    import pyxel as p
    target = Path(screenshot).resolve() if screenshot else None
    if target:
        target.parent.mkdir(parents=True, exist_ok=True)
    app = App()
    p.init(368, 304, title='Cargo Merchant', fps=30, display_scale=2, headless=headless)
    keys = {key: getattr(p, 'KEY_'+key) for key in ('LEFT', 'RIGHT', 'UP', 'DOWN', 'Z', 'SPACE', 'X', 'R')}
    demo_keys = DEMO_KEYS
    if preview:
        demo_keys = DEMO_KEYS[:18]+('RIGHT',)  # 積載→市場1→灯を選択。まだ売らない。
    limit = max(12, 2*len(demo_keys)+4) if demo or preview else 12

    def update():
        if headless and p.frame_count >= limit:
            if target:
                p.screenshot(str(target), scale=2)
                record = {'stage': 5, 'mode': 'preview' if preview else 'checkpoint' if demo else 'initial',
                          'input_kind': 'synthetic controller keys; not OS/human input',
                          'frames': p.frame_count, 'state': asdict(app.state),
                          'selected': app.selected, 'message': app.message}
                target.with_suffix('.json').write_text(json.dumps(record, ensure_ascii=False, indent=2), encoding='utf-8')
            p.quit()
            return
        if headless and (demo or preview):
            index = p.frame_count//2
            pressed = {demo_keys[index]} if p.frame_count%2 == 0 and index < len(demo_keys) 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)})
    p.run(update, lambda: app.draw(p))
    return app


if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--headless-smoke', action='store_true')
    parser.add_argument('--demo-checkpoint', action='store_true')
    parser.add_argument('--demo-preview', action='store_true')
    parser.add_argument('--screenshot')
    args = parser.parse_args()
    main(args.headless_smoke, args.screenshot, args.demo_checkpoint, args.demo_preview)