from otree.api import *

class C(BaseConstants):
    NAME_IN_URL = 'split_the_pie'
    PLAYERS_PER_GROUP = 2
    NUM_ROUNDS = 2
    ENDOWMENT = cu(100)

SPLITCORRECT_CODE = '123'

class Subsession(BaseSubsession):
    def creating_session(player):
        player.group_randomly()


class Group(BaseGroup):
    piekept = models.CurrencyField(
        min=0,
        max=C.ENDOWMENT,
        label="How much do you keep?"
    )

class Player(BasePlayer):
    split_passcode = models.StringField(blank=True, label="")

class SplitIntroduction(Page):
    form_model = 'player'
    form_fields = ['split_passcode']

    def is_displayed(player):
        return player.round_number == 1

    def error_message(player, values):
        if values['split_passcode'].strip().upper() != SPLITCORRECT_CODE.upper():
            return 'Incorrect code. Please try again.'

class Decide(Page):
    form_model = 'group'
    form_fields = ['piekept']

    @staticmethod
    def is_displayed(player):
        return player.id_in_group == 1

    @staticmethod
    def vars_for_template(player: Player):
        return dict(endowment=int(C.ENDOWMENT))

    @staticmethod
    def js_vars(player: Player):
        return dict(endowment=int(C.ENDOWMENT))

class ResultsWaitPage(WaitPage):
    @staticmethod
    def is_displayed(player):
        return player.round_number == 2

    @staticmethod
    def after_all_players_arrive(group: Group):
        p1 = group.get_player_by_id(1)
        p2 = group.get_player_by_id(2)
        # Only assign real payoffs in round 2; zero out for practice round
        if group.round_number == 2:
            p1.payoff = group.piekept
            p2.payoff = C.ENDOWMENT - group.piekept
        else:
            p1.payoff = cu(0)
            p2.payoff = cu(0)

class Results(Page):
    @staticmethod
    def is_displayed(player: Player):
        # Only show results in round 2
        return player.round_number == 2

    @staticmethod
    def vars_for_template(player: Player):
        return dict(
            kept=player.group.piekept,
            given=C.ENDOWMENT - player.group.piekept,
        )

class PrisonerBetweenRounds(Page):
    def vars_for_template(player):
        return {
            'round_number': player.round_number,
        }

    def is_displayed(player):
        return player.round_number == 1


page_sequence = [SplitIntroduction,
                 Decide,
                 ResultsWaitPage,
                 Results,
                 PrisonerBetweenRounds]