# 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.
# 
# This applies to original game source and primitive artwork. Bitmap glyph data
# is derived from Noto Sans CJK under SIL Open Font License 1.1, and is covered
# by assets/NOTO-LICENSE.txt instead of MIT.

"""まちの停電係・独立教材 1/5。架空の整数配分ゲーム。"""
from dataclasses import dataclass, replace

@dataclass(frozen=True)
class Period:
    generation: int
    a: int
    b: int
    optional_a: int = 0
    optional_b: int = 0

@dataclass(frozen=True)
class Level:
    title: str
    periods: tuple

@dataclass(frozen=True)
class Plan:
    a: int = 0
    b: int = 0

@dataclass(frozen=True)
class Record:
    period: int
    plan: Plan
    missing_a: int
    missing_b: int
    optional: int
    curtailed: int

@dataclass(frozen=True)
class State:
    t: int
    records: tuple = ()
    missing: int = 0
    optional: int = 0
    curtailed: int = 0
    outcome: str = "playing"

LEVELS = (Level("PRACTICE: TWO DISTRICTS", (Period(4,1,1,1,1),)),)

def initial(level):
    return State(0)

def loads(plan):
    return plan.a + plan.b, plan.a, plan.b

def legal(state, level, plan):
    if state.outcome != "playing" or state.t >= len(level.periods):
        return "terminal"
    if not isinstance(plan, Plan) or any(
        type(value) is not int or not 0 <= value <= 2
        for value in (plan.a, plan.b)
    ):
        return "range"
    period = level.periods[state.t]
    trunk, west, east = loads(plan)
    if trunk > period.generation:
        return "generation"
    if plan.a > period.a + period.optional_a or plan.b > period.b + period.optional_b:
        return "demand"
    return "commit"

def apply(state, level, plan):
    reason = legal(state, level, plan)
    if reason != "commit":
        return state, reason
    period = level.periods[state.t]
    missing_a = max(0, period.a - plan.a)
    missing_b = max(0, period.b - plan.b)
    optional = max(0, plan.a - period.a) + max(0, plan.b - period.b)
    curtailed = period.generation - plan.a - plan.b
    record = Record(state.t, plan, missing_a, missing_b, optional, curtailed)
    after = State(state.t + 1, state.records + (record,),
                  state.missing + missing_a + missing_b,
                  state.optional + optional, state.curtailed + curtailed)
    if after.t == len(level.periods):
        after = replace(after, outcome="won" if after.missing == 0 else "lost")
    return after, reason

FIELDS = ('a', 'b')

class App:
    def __init__(self):
        self.level_index = 0
        self.reset()

    @property
    def level(self):
        return LEVELS[self.level_index]

    def reset(self):
        self.state = initial(self.level)
        self.plan = Plan()
        self.selection = 0
        self.gated = False
        self.message = "READY"

    def commit(self):
        self.state, self.message = apply(self.state, self.level, self.plan)
        if self.message == "commit":
            self.plan = Plan()

    def control(self, pressed, held):
        # A new R wins even during release gating or after the result.
        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.message = "RELEASE ALL KEYS"
            self.gated = True
            return
        if len(pressed) != 1:
            return
        key = next(iter(pressed))
        self.gated = True
        if self.state.outcome != "playing":
            return
        if key == "C":
            self.plan = Plan()
        elif key == "UP":
            self.selection = (self.selection - 1) % len(FIELDS)
        elif key == "DOWN":
            self.selection = (self.selection + 1) % len(FIELDS)
        elif key in ("LEFT", "RIGHT"):
            field = FIELDS[self.selection]
            value = getattr(self.plan, field) + (1 if key == "RIGHT" else -1)
            self.plan = replace(self.plan, **{field: max(0, min(2, value))})
        elif key in ("Z", "SPACE"):
            self.commit()

    def draw(self, p):
        p.cls(1)
        s, level, plan = self.state, self.level, self.plan
        period = level.periods[min(s.t, len(level.periods)-1)]
        trunk, west, east = loads(plan)
        after, reason = apply(s, level, plan)
        p.text(12,10,"TOWN POWER / STAGE 1",10)
        p.text(12,25,level.title,7)
        p.text(400,10,"PERIOD "+str(s.t)+"/"+str(len(level.periods)),7)
        p.text(400,25,s.outcome.upper(),11 if s.outcome=="won" else 7)
        p.rect(12,44,552,112,2)
        p.rectb(28,76,76,46,6)
        p.text(38,86,"GENERATOR",7)
        p.text(38,102,"G "+str(period.generation),10)
        p.line(104,98,235,98,10)
        p.circ(239,98,4,10)
        p.line(239,98,239,66,10)
        p.line(239,66,356,66,10)
        p.line(239,98,239,130,10)
        p.line(239,130,356,130,10)
        p.rectb(356,52,70,29,6)
        p.rectb(356,116,70,29,6)
        p.text(363,58,"WEST A "+str(plan.a),7)
        p.text(363,124,"EAST B "+str(plan.b),7)
        p.text(115,82,"TOTAL "+str(trunk),7)
        for x, label in ((16,"PERIOD"),(112,"G"),(200,"A NEED/+"),(290,"B NEED/+"),
                         (390,"MISS"),(442,"OPT"),(498,"CUT")):
            p.text(x,167,label,6)
        for index, d in enumerate(level.periods):
            y = 184 + 16*index
            if index == s.t:
                p.rect(12,y-2,552,13,3)
            for x, value in ((16,index),(112,d.generation),(200,f"{d.a}/{d.optional_a}"),
                             (290,f"{d.b}/{d.optional_b}")):
                p.text(x,y,str(value),7)
            if index < len(s.records):
                r = s.records[index]
                for x, value in ((390,r.missing_a+r.missing_b),(442,r.optional),(498,r.curtailed)):
                    p.text(x,y,str(value),11)
        for index, field in enumerate(FIELDS):
            x = 12 + 138*index
            p.rectb(x,254,128,24,10 if index==self.selection else 5)
            p.text(x+7,263,field.upper()+" "+str(getattr(plan,field)),7)
        p.text(12,289,"PREVIEW: "+reason.upper(),11 if reason=="commit" else 8)
        if reason=="commit":
            p.text(220,289,f"MISS +{after.missing-s.missing}  CUT +{after.curtailed-s.curtailed}",7)
        p.text(12,322,f"TOTAL MISS {s.missing}  OPTIONAL {s.optional}  CURTAILED {s.curtailed}",7)
        p.text(12,339,"LAST: "+self.message.upper(),10)
        p.text(12,360,"UP/DOWN ROUTE  LEFT/RIGHT AMOUNT  SPACE/Z COMMIT",6)
        p.text(12,376,"C CLEAR PLAN  R RESET  ESC QUIT",6)
        p.text(12,392,"1 PERIOD / NO WIRE LIMITS OR BATTERY YET",6)

DEMO_KEYS = ('RIGHT', 'RIGHT', 'DOWN', 'RIGHT', 'RIGHT', 'SPACE')
KEYS = ('UP', 'DOWN', 'LEFT', 'RIGHT', 'SPACE', 'Z', 'C', 'R')

def main(headless=False, screenshot=None, demo=False, frames=None):
    # Import Pyxel only after command-line validation and only for rendering.
    from pathlib import Path
    import json
    import pyxel as p
    app = App()
    destination = Path(screenshot).resolve() if screenshot else None
    if destination:
        destination.parent.mkdir(parents=True, exist_ok=True)
    limit = frames or (max(12, 2*len(DEMO_KEYS)+4) if demo else 12)
    p.init(576,416,title="Town Power lesson 1",fps=30,display_scale=2,headless=headless)
    keys = {name:getattr(p,"KEY_"+name) for name in KEYS}
    def update():
        if headless and p.frame_count >= limit:
            if destination:
                p.screenshot(str(destination),scale=1)
            from dataclasses import asdict
            print(json.dumps({"stage":1,"state":asdict(app.state),
                              "plan":asdict(app.plan),"message":app.message},sort_keys=True),flush=True)
            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))


def parse_args(argv=None):
    import argparse
    parser = argparse.ArgumentParser(description="Independent Town Power lesson 1")
    parser.add_argument("--headless-smoke",action="store_true")
    parser.add_argument("--demo-checkpoint",action="store_true")
    parser.add_argument("--screenshot",metavar="FILE.png")
    parser.add_argument("--qa-frames",type=int,metavar="12..600")
    args = parser.parse_args(argv)
    if not args.headless_smoke and (args.demo_checkpoint or args.screenshot or args.qa_frames is not None):
        parser.error("QA options require --headless-smoke")
    if args.qa_frames is not None and not 12 <= args.qa_frames <= 600:
        parser.error("--qa-frames must be 12..600")
    if args.demo_checkpoint and args.qa_frames is not None and args.qa_frames < 2*len(DEMO_KEYS)+4:
        parser.error("--qa-frames is shorter than the demo")
    if args.screenshot and not args.screenshot.lower().endswith(".png"):
        parser.error("--screenshot must end in .png")
    return args


if __name__ == "__main__":
    args = parse_args()
    main(args.headless_smoke, args.screenshot, args.demo_checkpoint, args.qa_frames)