# 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

WEATHER = (0, 0, 0, 2, 2, 2)
WEATHER_NAMES = ("WET", "DRY", "COLD")
BASE_HARVEST = ((2, 0, 0), (3, 2, 1))
BASE_USE = (2, 3, 4)
ENERGY_CAP = 18
TRAITS = {"": "EMPTY", "F": "FIN", "C": "FUR", "T": "CLAW"}
UPKEEP = {"F": (0, 1, 1), "C": (1, 2, 0), "T": (1, 1, 1)}


@dataclass(frozen=True)
class State:
    day: int = 0
    energy: int = 6
    traits: tuple = ("", "")
    changed: bool = False
    changes: int = 0
    tick: int = 0
    outcome: str = "running"
    history: tuple = ()

def forecast(state, field):
    if state.day >= 6:
        return {"harvest": 0, "base": 0, "upkeep": 0, "change": 0, "after": state.energy}
    weather = WEATHER[state.day]
    harvest = BASE_HARVEST[field][weather]
    base = BASE_USE[weather]
    if "F" in state.traits and field == 0:
        harvest += (3, 1, 0)[weather]
    if "T" in state.traits and field == 1:
        harvest += 2
    if "C" in state.traits and weather == 2:
        base -= 2
    upkeep = sum(UPKEEP[t][weather] for t in state.traits if t)
    paid_change = 2 if state.changed else 0
    # paid_changeは表示する支払済み費用。ここでは二重に引かない。
    return {"harvest": harvest, "base": base, "upkeep": upkeep,
            "change": paid_change,
            "after": min(ENERGY_CAP, state.energy+harvest-base-upkeep)}


def move(state, action):
    if state.outcome != "running":
        return state, "ENDED"
    if isinstance(action, str) and action.startswith("set:"):
        try:
            _, raw, trait = action.split(":")
            slot = int(raw)
        except (ValueError, TypeError):
            return state, "INVALID"
        if slot not in (0, 1) or trait not in TRAITS:
            return state, "INVALID"
        if trait == state.traits[slot]:
            return state, "SAME TRAIT"
        if trait and trait in state.traits:
            return state, "NO DUPLICATES"
        if state.day > 0 and state.changed:
            return state, "ONE CHANGE PER MORNING"
        cost = 0 if state.day == 0 else 2
        if state.energy-cost <= 0:
            return state, "NEED ENERGY AFTER PAYMENT"
        traits = list(state.traits)
        traits[slot] = trait
        return replace(state, traits=tuple(traits), energy=state.energy-cost,
                       changed=state.day>0, changes=state.changes+(state.day>0),
                       tick=state.tick+1), "CHANGED / PAID "+str(cost)
    if action not in ("forage:0", "forage:1"):
        return state, "INVALID"
    field = int(action[-1])
    detail = forecast(state, field)
    energy = detail["after"]
    day = state.day+1
    outcome = "failed" if energy <= 0 else "won" if day == 6 else "running"
    entry = (day, WEATHER[state.day], state.traits, field, state.energy,
             detail["harvest"], detail["base"], detail["upkeep"], detail["change"], energy)
    return replace(state, day=day, energy=energy, changed=False, tick=state.tick+1,
                   outcome=outcome, history=state.history+(entry,)), "DAY CONFIRMED"

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


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

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

    def reset(self):
        self.timeline = Timeline()
        self.field = 0
        self.slot = 0
        self.gated = False
        self.message = "UP/DOWN FIELD, SPACE CONFIRM"

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

    def control(self, pressed, held):
        # 新しいRはいつでも優先。ただし、その後は全解放を待つ。
        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 = "ONE KEY AT A TIME / RELEASE ALL"
            return
        if len(pressed) != 1:
            return
        key = next(iter(pressed))
        if key == "X":
            self.timeline, self.message = step(self.timeline, "undo")
            self.gated = True
            return
        if self.state.outcome != "running":
            return
        if key in ("UP", "DOWN"):
            self.field = 0 if key == "UP" else 1
            self.gated = True
            return
        if key in ("LEFT", "RIGHT"):
            self.slot = 0 if key == "LEFT" else 1
            self.gated = True
            return
        action = "forage:"+str(self.field) if key == "SPACE" else None
        if key in ("0", "1", "2", "3"):
            action = "set:"+str(self.slot)+":"+{"0": "", "1": "F", "2": "C", "3": "T"}[key]
        if action is None:
            return
        self.timeline, self.message = step(self.timeline, action)
        self.gated = True

    def draw(self, p):
        p.cls(1)
        p.text(8, 8, "TRAIT ISLAND / STAGE 5", 10)
        p.text(8, 22, "DAY "+str(self.state.day)+" / 6  ENERGY "+str(self.state.energy), 7)
        for i in range(ENERGY_CAP):
            p.rect(8+i*10, 34, 8, 10, 11 if i<self.state.energy else 5)
        for i, weather in enumerate(WEATHER):
            p.text(8+i*60, 55, str(i+1)+":"+WEATHER_NAMES[weather], 10 if i==self.state.day else 6)
        for slot, trait in enumerate(self.state.traits):
            p.text(8+slot*165, 74, ("> " if slot==self.slot else "  ")+"SLOT "+str(slot+1)+": "+TRAITS[trait], 10 if slot==self.slot else 7)
        p.text(8, 87, "1 FIN  2 FUR  3 CLAW  0 EMPTY", 7)
        detail = forecast(self.state, self.field)
        p.text(8, 107, "FIELD: "+("POND" if self.field==0 else "TREE"), 10)
        p.text(8, 121, "HARVEST +"+str(detail["harvest"])+"  BASE -"+str(detail["base"]), 7)
        p.text(8, 134, "UPKEEP -"+str(detail["upkeep"]), 7)
        p.text(8, 148, "CHANGE ALREADY PAID: "+str(detail["change"]), 6)
        p.text(8, 164, "NEXT ENERGY: "+str(detail["after"])+" / CAP "+str(ENERGY_CAP), 11)
        if self.state.outcome != "running":
            result = "FAILED" if self.state.outcome=="failed" else "SURVIVED / "+("EGG" if self.state.energy>=5 else "NO EGG")
            p.text(8, 182, result, 10)
        p.text(8, 200, self.message, 7)
        p.text(8, 220, "UP/DOWN FIELD  SPACE FORAGE  R RESET", 6)
        p.text(8, 232, "LEFT/RIGHT SLOT  0/1/2/3 SET TRAIT", 6)
        p.text(8, 244, "X UNDO  ESC QUIT", 6)

DEMO_KEYS = ('1', 'SPACE', 'SPACE', 'SPACE', '2', 'DOWN', 'SPACE', 'SPACE', 'SPACE')

CONTROL_KEYS = ('UP', 'DOWN', 'SPACE', 'R', 'LEFT', 'RIGHT', '0', '1', '2', '3', 'X')

def main(headless=False, screenshot=None, demo=False):
    from pathlib import Path
    import pyxel as p
    screenshot = str(Path(screenshot).resolve()) if screenshot else None
    if screenshot:
        Path(screenshot).parent.mkdir(parents=True, exist_ok=True)
    app = App()
    p.init(384, 260, title="Trait Island", fps=30, display_scale=3, headless=headless)
    keys = {key: getattr(p, "KEY_"+key) for key in CONTROL_KEYS}
    limit = max(12, 2*len(DEMO_KEYS)+3) if demo else 12

    def update():
        if headless and p.frame_count >= limit:
            if screenshot:
                p.screenshot(screenshot, scale=3)
            p.quit()
            return
        if headless and demo:
            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({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))
    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("--screenshot")
    args = parser.parse_args()
    main(args.headless_smoke, args.screenshot, args.demo_checkpoint)