from otree.api import *
import random
from constants import *

doc = """
Setup App — Gen1 only.
"""


class C(BaseConstants):
    NAME_IN_URL = 'AA_SetUpApp'
    PLAYERS_PER_GROUP = None
    NUM_ROUNDS = 1


class Subsession(BaseSubsession):
    gen1_threshold = models.FloatField(initial=0)


def creating_session(subsession: Subsession):
    if subsession.round_number != 1:
        return

    players = subsession.get_players()
    total_players = len(players)
    subsession.gen1_threshold = total_players / 2

    gen_counters = {1: 0, 2: 0}
    gen1_ids = []
    gen2_ids = []
    active_ids = []
    own_id_to_subsession_id = {}

    for player in players:
        generation = 1 if player.id_in_group <= subsession.gen1_threshold else 2
        player.Ana_generation = generation

        if generation == 1:
            player.participant.vars['flip_result_list'] = [
                random.choice(['Heads', 'Tails']) for _ in range(10)
            ]
            player.participant.vars['N_CoinFlip'] = 0
            player.participant.vars['N_TokenTake'] = 0
            player.participant.vars['first_task_order'] = random.randint(0, 1)

        gen_counters[generation] += 1
        player.Ana_own_id = gen_counters[generation]
        player.symbol = "1"

        player.participant.vars.update({
            'symbol': player.symbol,
            'Ana_id_in_group': player.id_in_group,
            'Ana_generation': generation,
            'Ana_own_id': player.Ana_own_id,
            'state': "active",
            'first_practice_order': random.randint(0, 1),
            'Ana_family_assignment': 'XX',
        })

        # Build ID lists
        if generation == 1:
            gen1_ids.append(player.id_in_subsession)
        else:
            gen2_ids.append(player.id_in_subsession)
        active_ids.append(player.id_in_subsession)

        # Build own_id lookup map — key: "gen{generation}_{own_id}"
        own_id_to_subsession_id[f"gen{generation}_{gen_counters[generation]}"] = player.id_in_subsession

    subsession.session.vars.update({
        # Existing vars — unchanged
        'arrival_checkins': {},
        'gen1_slot_counter': 0,
        'gen2_slot_counter': 0,
        'allowed_gen1': None,
        'allowed_gen2': None,
        'final_gen1_count': 0,
        # Pre-built caches — available instantly from first page load onwards
        'gen1_ids': gen1_ids,
        'gen2_ids': gen2_ids,
        'active_ids': active_ids,
        'all_ids': gen1_ids + gen2_ids,
        'gen2_count': len(gen2_ids),
        'own_id_to_subsession_id': own_id_to_subsession_id,
        '_caches_built': True,
    })


class Group(BaseGroup):
    pass


class Player(BasePlayer):
    manual_entry = models.StringField(blank=True, label="Enter code (or leave blank to continue):")
    state = models.StringField(initial="active")
    Ana_own_id = models.IntegerField()
    Ana_generation = models.IntegerField()
    symbol = models.StringField()
    Ana_family_assignment = models.StringField(initial="XX")
    N_CoinFlip = models.IntegerField(initial=0)
    N_TokenTake = models.IntegerField(initial=0)


class ArrivalPage(Page):
    @staticmethod
    def is_displayed(player):
        return player.round_number == 1

    @staticmethod
    def vars_for_template(player):
        return dict(total_players=player.subsession.session.num_participants)

    @staticmethod
    def js_vars(player):
        return dict(participant_code=player.participant.code)

    @staticmethod
    def live_method(player: Player, data: dict):
        if data.get('type') != 'checkin':
            return
        session = player.session
        session.vars['arrival_checkins'][player.participant.code] = True
        arrived = len(session.vars['arrival_checkins'])
        total = session.num_participants
        all_arrived = arrived >= total
        return {0: dict(arrived=arrived, all_arrived=all_arrived)}


class ManualDrop_G1(Page):
    form_model = 'player'
    form_fields = ['manual_entry']

    @staticmethod
    def is_displayed(player):
        return (player.round_number == 1
                and player.participant.vars.get('Ana_generation') == 1
                and player.participant.vars.get('state') == "active")

    @staticmethod
    def before_next_page(player: Player, timeout_happened=False):
        if player.manual_entry and player.manual_entry.strip().upper() == "QUIT":
            player.state = "inactive"
            player.participant.vars['state'] = "inactive"
            # Remove from active_ids cache
            active_ids = player.session.vars.get('active_ids', [])
            if player.id_in_subsession in active_ids:
                active_ids.remove(player.id_in_subsession)
            return

        session = player.session
        if session.vars.get('allowed_gen1') is None:
            _compute_allowed(player.subsession)

        allowed = session.vars['allowed_gen1']
        session.vars['gen1_slot_counter'] += 1
        slot = session.vars['gen1_slot_counter']

        if slot <= allowed:
            player.Ana_own_id = slot
            player.participant.vars['Ana_own_id'] = slot
            fam = FamilyName[slot - 1]
            player.Ana_family_assignment = fam
            player.participant.vars['Ana_family_assignment'] = fam
            player.state = "active"
            player.participant.vars['state'] = "active"
            session.vars['final_gen1_count'] = slot

            # Update own_id map with confirmed slot assignment
            session.vars['own_id_to_subsession_id'][f"gen1_{slot}"] = player.id_in_subsession
        else:
            player.state = "closed"
            player.participant.vars['state'] = "closed"
            # Remove closed players from active_ids cache
            active_ids = player.session.vars.get('active_ids', [])
            if player.id_in_subsession in active_ids:
                active_ids.remove(player.id_in_subsession)


class WaitForGen1(WaitPage):
    @staticmethod
    def is_displayed(player):
        return (player.round_number == 1
                and player.participant.vars.get('Ana_generation') == 1
                and player.participant.vars.get('state') in ("active", "closed"))


class Gen1RedirectToFirstTask(Page):
    @staticmethod
    def is_displayed(player):
        return (player.round_number == 1
                and player.participant.vars.get('Ana_generation') == 1)

    @staticmethod
    def app_after_this_page(player, upcoming_apps):
        state = player.participant.vars.get('state')
        if state == "inactive":
            if 'I_Post_Questionnaire_D' in upcoming_apps:
                return 'I_Post_Questionnaire_D'
            return upcoming_apps[-1] if upcoming_apps else None
        if state == "active":
            return 'A_Endowment_Creation_Gen1_D'
        return None


def _compute_allowed(subsession):
    all_players = subsession.get_players()
    active_gen1 = sum(1 for p in all_players
                      if p.participant.vars.get('Ana_generation') == 1
                      and p.participant.vars.get('state') == "active")
    active_gen2 = sum(1 for p in all_players
                      if p.participant.vars.get('Ana_generation') == 2
                      and p.participant.vars.get('state') == "active")
    allowed = min(active_gen1, active_gen2)
    allowed = (allowed // 3) * 3
    subsession.session.vars['allowed_gen1'] = allowed
    subsession.session.vars['allowed_gen2'] = allowed


class SessionFullPage_G1(Page):
    @staticmethod
    def is_displayed(player):
        return (player.round_number == 1
                and player.participant.vars.get('Ana_generation') == 1
                and player.participant.vars.get('state') == "closed")

    @staticmethod
    def app_after_this_page(player, upcoming_apps):
        if 'I_Post_Questionnaire_D' in upcoming_apps:
            return 'I_Post_Questionnaire_D'
        return upcoming_apps[-1] if upcoming_apps else None


page_sequence = [
    ArrivalPage,
    ManualDrop_G1,
    WaitForGen1,
    Gen1RedirectToFirstTask,
    SessionFullPage_G1,
]