import random
import string
import json
import time
from otree.api import *


class Constants(BaseConstants):
    name_in_url = 'letter_search'
    players_per_group = None
    num_rounds = 3
    round_time = 240
    showup_fee = Currency(5.00)
    reward_per_grid = Currency(0.25)


def generate_letter_search():
    target = random.choice(string.ascii_uppercase)
    alphabet = [c for c in string.ascii_uppercase if c != target]

    # Start with zero instances of the target
    grid = [[random.choice(alphabet) for _ in range(10)] for _ in range(10)]

    count = random.randint(1, 8)
    placed = 0
    while placed < count:
        i, j = random.randrange(10), random.randrange(10)
        if grid[i][j] != target:
            grid[i][j] = target
            placed += 1

    # Now actual count = intended count, guaranteed
    return dict(letter=target, count=count, grid=grid)


class Subsession(BaseSubsession):
    def creating_session(subsession):
        if subsession.round_number == 1:
            for p in subsession.get_players():
                p.participant.vars['total_correct'] = 0

        for p in subsession.get_players():
            # Prolific PID from participant_label
            p.prolific_pid = p.participant.label

            # Optional: study & session from URL params
            # (only if your Prolific link includes them with these exact names)
            vars = p.participant.vars
            p.study_id = vars.get('study_id', '')
            p.prolific_session_id = vars.get('prolific_session_id', '')


class Group(BaseGroup):
    pass


class Player(BasePlayer):
    prolific_pid = models.StringField()
    study_id = models.StringField()
    prolific_session_id = models.StringField()
    consent = models.BooleanField(blank=True, initial=False)
    current_letter = models.StringField()
    current_count = models.IntegerField()
    grid_json = models.LongStringField()
    total_correct = models.IntegerField(initial=0)
    R1_correct = models.IntegerField(initial=0)
    R2_correct = models.IntegerField(initial=0)
    R3_correct = models.IntegerField(initial=0)
    selected_pay_round = models.IntegerField()
    payoff_from_task = models.CurrencyField()
    familiarity = models.IntegerField(
        label="How familiar are you with letter search tasks?",
        choices=[
            [1, "I have never heard of them"],
            [2, ""],
            [3, ""],
            [4, "I have done a few before"],
            [5, ""],
            [6, ""],
            [7, "I do them frequently"]
        ],
        widget=widgets.RadioSelectHorizontal
    )

    enjoyment = models.IntegerField(
        label="To what extent did you enjoy completing the task?",
        choices=[
            [1, "Not at all"],
            [2, ""],
            [3, ""],
            [4, "Somewhat"],
            [5, ""],
            [6, ""],
            [7, "Very much"]
        ],
        widget=widgets.RadioSelectHorizontal
    )

    gender = models.StringField(
        label="Gender",
        choices=["Male", "Female", "Non-binary / Third gender", "Prefer not to say"],
        widget=widgets.RadioSelect
    )

    age = models.IntegerField(label="Age")


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




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


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


class NextRound(Page):
    pass


class LetterSearch(Page):
    timeout_seconds = Constants.round_time

    @staticmethod
    def vars_for_template(player: Player):
        puzzle = generate_letter_search()
        player.current_letter = puzzle['letter']
        player.current_count = puzzle['count']
        player.grid_json = json.dumps(puzzle['grid'])

        # Set up tracking if not yet initialized
        if 'total_correct' not in player.participant.vars:
            player.participant.vars['total_correct'] = 0
        round_key = f'round_{player.round_number}_correct'
        if round_key not in player.participant.vars:
            player.participant.vars[round_key] = 0

        return dict(
            letter=player.current_letter,
            grid=puzzle['grid'],
            total_correct=player.participant.vars['total_correct'],
            round_correct=player.participant.vars[round_key],
            round_time=Constants.round_time,
        )

    @staticmethod
    def live_method(player: Player, data):
        if data.get('type') == 'init':
            # Only generate puzzle if not already generated for this round
            if not player.current_letter:
                puzzle = generate_letter_search()
                player.current_letter = puzzle['letter']
                player.current_count = puzzle['count']
                player.grid_json = json.dumps(puzzle['grid'])

            # Set round start time once per round
            start_time_key = f'round_start_time_{player.round_number}'
            if start_time_key not in player.participant.vars:
                player.participant.vars[start_time_key] = time.time()

            # Calculate remaining time
            elapsed = time.time() - player.participant.vars[start_time_key]
            time_left = max(0, int(Constants.round_time - elapsed))

            return {
                player.id_in_group: dict(
                    letter=player.current_letter,
                    grid=json.loads(player.grid_json),
                    total_correct=player.participant.vars['total_correct'],
                    round_correct=player.participant.vars[round_key],
                    time_left=time_left
                )
            }

        if data.get('type') == 'submit':
            round_key = f'round_{player.round_number}_correct'

            if data['response'] == player.current_count:
                player.participant.vars['total_correct'] += 1
                player.participant.vars[round_key] += 1

                puzzle = generate_letter_search()
                player.current_letter = puzzle['letter']
                player.current_count = puzzle['count']
                player.grid_json = json.dumps(puzzle['grid'])

                return {
                    player.id_in_group: dict(
                        correct=True,
                        letter=player.current_letter,
                        grid=puzzle['grid'],
                        total_correct=player.participant.vars['total_correct'],
                        round_correct=player.participant.vars[round_key]
                    )
                }
            else:
                return {player.id_in_group: dict(correct=False)}

    @staticmethod
    def before_next_page(player: Player, timeout_happened):
        round_key = f'round_{player.round_number}_correct'
        round_correct = player.participant.vars.get(round_key, 0)
        total_correct = player.participant.vars.get('total_correct', 0)

        if player.round_number == 1:
            player.R1_correct = round_correct
        elif player.round_number == 2:
            player.R2_correct = round_correct
        elif player.round_number == 3:
            player.R3_correct = round_correct

        player.total_correct = total_correct


class SurveyStart(Page):
    @staticmethod
    def is_displayed(player: Player):
        return player.round_number == Constants.num_rounds


class PEQ(Page):
    form_model = 'player'
    form_fields = ['familiarity', 'enjoyment', 'gender', 'age']

    @staticmethod
    def is_displayed(player: Player):
        return player.round_number == Constants.num_rounds

    @staticmethod
    def before_next_page(player: Player, timeout_happened):
        # Choose a random paying round once
        if 'paying_round' not in player.participant.vars:
            paying_round = random.randint(1, Constants.num_rounds)
            player.participant.vars['paying_round'] = paying_round
        else:
            paying_round = player.participant.vars['paying_round']

        # Get correct counts for that round from participant.vars
        round_key = f'round_{paying_round}_correct'
        correct_in_pay_round = player.participant.vars.get(round_key, 0)

        # Store on the player for convenience
        player.selected_pay_round = paying_round
        player.payoff_from_task = correct_in_pay_round * Constants.reward_per_grid

        # Final payoff: showup fee + payoff from chosen round
        player.participant.payoff = Constants.showup_fee + player.payoff_from_task


class Results(Page):
    @staticmethod
    def is_displayed(player: Player):
        return player.round_number == Constants.num_rounds

    @staticmethod
    def vars_for_template(player: Player):
        paying_round = player.participant.vars.get('paying_round', None)

        round_1 = player.participant.vars.get('round_1_correct', 0)
        round_2 = player.participant.vars.get('round_2_correct', 0)
        round_3 = player.participant.vars.get('round_3_correct', 0)
        total = player.participant.vars.get('total_correct', 0)

        if paying_round is not None:
            correct_in_pay_round = player.participant.vars.get(
                f'round_{paying_round}_correct', 0
            )
        else:
            correct_in_pay_round = 0

        return dict(
            total_correct=total,
            showup_fee=Constants.showup_fee,
            reward_per_grid=Constants.reward_per_grid,
            payoff=player.participant.payoff,
            round_1=round_1,
            round_2=round_2,
            round_3=round_3,
            total=total,
            paying_round=paying_round,
            correct_in_pay_round=correct_in_pay_round,
            payoff_from_task=player.payoff_from_task,
        )


class Final(Page):
    @staticmethod
    def is_displayed(player: Player):
        return player.round_number == Constants.num_rounds




page_sequence = [Consent,
                 Introduction,
                 Instructions,
                 NextRound,
                 LetterSearch,
                 SurveyStart,
                 PEQ,
                 Results,
                 Final
                 ]