from otree.api import *
import random
import json
from openai import OpenAI
import os

client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))


doc = 'Narratives with Krupka & Weber Social Norms Elicitation'


# -----------------CLASSES-----------------#
class C(BaseConstants):
    NAME_IN_URL = 'NarrativeKW'
    PLAYERS_PER_GROUP = None
    NUM_ROUNDS = 1

    # Krupka & Weber appropriateness scale
    APPROPRIATENESS_CHOICES = [
        [1, 'Very socially inappropriate'],
        [2, 'Socially inappropriate'],
        [3, 'Somewhat socially inappropriate'],
        [4, 'Neither socially appropriate nor inappropriate'],
        [5, 'Somewhat socially appropriate'],
        [6, 'Socially appropriate'],
        [7, 'Very socially appropriate'],
    ]

class Subsession(BaseSubsession):
    pass


class Group(BaseGroup):
    pass


class Player(BasePlayer):
    consent = models.BooleanField(
        label='<b>I have read and understood the information above, and I give my consent to participate in this study.</b>')
    narrative_treatment = models.StringField()
    social_context = models.StringField()

    # For Krupka & Weber payment mechanism
    selected_allocation = models.IntegerField()  # Randomly selected allocation (0-10)
    modal_response = models.IntegerField()  # Modal response for that allocation
    player_response = models.IntegerField()  # Player's response for selected allocation
    matches_modal = models.BooleanField()  # Whether player matched modal response

    # Krupka & Weber appropriateness ratings for each allocation (0-10)
    appropriateness_0 = models.IntegerField(
        label='Keep 0 for yourself, give 10 to the organization',
        widget=widgets.RadioSelectHorizontal,
        choices=C.APPROPRIATENESS_CHOICES
    )
    appropriateness_1 = models.IntegerField(
        label='Keep 1 for yourself, give 9 to the organization',
        widget=widgets.RadioSelectHorizontal,
        choices=C.APPROPRIATENESS_CHOICES
    )
    appropriateness_2 = models.IntegerField(
        label='Keep 2 for yourself, give 8 to the organization',
        widget=widgets.RadioSelectHorizontal,
        choices=C.APPROPRIATENESS_CHOICES
    )
    appropriateness_3 = models.IntegerField(
        label='Keep 3 for yourself, give 7 to the organization',
        widget=widgets.RadioSelectHorizontal,
        choices=C.APPROPRIATENESS_CHOICES
    )
    appropriateness_4 = models.IntegerField(
        label='Keep 4 for yourself, give 6 to the organization',
        widget=widgets.RadioSelectHorizontal,
        choices=C.APPROPRIATENESS_CHOICES
    )
    appropriateness_5 = models.IntegerField(
        label='Keep 5 for yourself, give 5 to the organization',
        widget=widgets.RadioSelectHorizontal,
        choices=C.APPROPRIATENESS_CHOICES
    )
    appropriateness_6 = models.IntegerField(
        label='Keep 6 for yourself, give 4 to the organization',
        widget=widgets.RadioSelectHorizontal,
        choices=C.APPROPRIATENESS_CHOICES
    )
    appropriateness_7 = models.IntegerField(
        label='Keep 7 for yourself, give 3 to the organization',
        widget=widgets.RadioSelectHorizontal,
        choices=C.APPROPRIATENESS_CHOICES
    )
    appropriateness_8 = models.IntegerField(
        label='Keep 8 for yourself, give 2 to the organization',
        widget=widgets.RadioSelectHorizontal,
        choices=C.APPROPRIATENESS_CHOICES
    )
    appropriateness_9 = models.IntegerField(
        label='Keep 9 for yourself, give 1 to the organization',
        widget=widgets.RadioSelectHorizontal,
        choices=C.APPROPRIATENESS_CHOICES
    )
    appropriateness_10 = models.IntegerField(
        label='Keep 10 for yourself, give 0 to the organization',
        widget=widgets.RadioSelectHorizontal,
        choices=C.APPROPRIATENESS_CHOICES
    )

    # Manipulation checks
    en_q1 = models.BooleanField(
        label='The main character of The OceanMaker is a female pilot who flies a small red airplane.',
        choices=[[True, 'True'], [False, 'False']]
    )
    en_q2 = models.BooleanField(
        label='The world of the film shows a dry ocean floor—ships are stranded on land with no water around them.',
        choices=[[True, 'True'], [False, 'False']]
    )
    en_q3 = models.BooleanField(
        label='The water in the film is created by a giant underground machine that pumps water from deep wells.',
        choices=[[True, 'True'], [False, 'False']]
    )
    en_q4 = models.BooleanField(
        label='The pilot battles sky pirates who try to steal clouds using their own aircraft.',
        choices=[[True, 'True'], [False, 'False']]
    )
    en_q5 = models.BooleanField(
        label='There is no spoken dialogue in the film; the whole story is told visually.',
        choices=[[True, 'True'], [False, 'False']]
    )
    en_q6 = models.BooleanField(
        label='At one point, the pilot uses a harpoon-like device to pull water out of the air.',
        choices=[[True, 'True'], [False, 'False']]
    )
    en_q7 = models.BooleanField(
        label='The final cloud in the movie is turned into rain by a special device attached to the pilot\'s plane.',
        choices=[[True, 'True'], [False, 'False']]
    )
    en_q8 = models.BooleanField(
        label='A large desert storm destroys the pilot\'s plane before she can reach the cloud.',
        choices=[[True, 'True'], [False, 'False']]
    )
    en_q9 = models.BooleanField(
        label='The pirates use grappling hooks and nets to capture clouds.',
        choices=[[True, 'True'], [False, 'False']]
    )
    en_q10 = models.BooleanField(
        label='In the final scene, the young girl who appears earlier is shown witnessing the return of rainfall.',
        choices=[[True, 'True'], [False, 'False']]
    )
    sn_q1 = models.BooleanField(
        label='The main character is a secret agent named Walter Beckett.',
        choices=[[True, 'True'], [False, 'False']]
    )
    sn_q2 = models.BooleanField(
        label='The short begins with the protagonist accidentally dropping his high-tech briefcase in a pigeon\'s nest.',
        choices=[[True, 'True'], [False, 'False']]
    )
    sn_q3 = models.BooleanField(
        label='The briefcase contains top-secret nuclear launch codes.',
        choices=[[True, 'True'], [False, 'False']]
    )
    sn_q4 = models.BooleanField(
        label='The pigeon inside the briefcase causes chaos by pressing buttons on the equipment.',
        choices=[[True, 'True'], [False, 'False']]
    )
    sn_q5 = models.BooleanField(
        label='There is spoken dialogue in the film throughout the entire short.',
        choices=[[True, 'True'], [False, 'False']]
    )
    sn_q6 = models.BooleanField(
        label='At one point, the pigeon flies the briefcase through a restaurant, causing destruction.',
        choices=[[True, 'True'], [False, 'False']]
    )
    sn_q7 = models.BooleanField(
        label='Walter successfully recovers the briefcase without any damage by the end of the film.',
        choices=[[True, 'True'], [False, 'False']]
    )
    sn_q8 = models.BooleanField(
        label='The film is roughly 6 minutes long.',
        choices=[[True, 'True'], [False, 'False']]
    )
    sn_q9 = models.BooleanField(
        label='The pigeon is caught in the briefcase accidentally while trying to steal food.',
        choices=[[True, 'True'], [False, 'False']]
    )
    sn_q10 = models.BooleanField(
        label='The short ends with the pigeon safely flying away after the chaos.',
        choices=[[True, 'True'], [False, 'False']]
    )
    transport_1 = models.IntegerField(
        label='While I was watching the video, I could easily picture the events in it taking place.',
        max=7, min=1,
        widget=widgets.RadioSelectHorizontal, choices=[1, 2, 3, 4, 5, 6, 7]
    )
    transport_2 = models.IntegerField(
        label='I could imagine myself in the scene of the events described in the video.',
        max=7, min=1,
        widget=widgets.RadioSelectHorizontal, choices=[1, 2, 3, 4, 5, 6, 7]
    )
    transport_3 = models.IntegerField(
        label='I was mentally involved in the video while watching it.',
        max=7, min=1,
        widget=widgets.RadioSelectHorizontal, choices=[1, 2, 3, 4, 5, 6, 7]
    )
    transport_4 = models.IntegerField(
        label='After finishing the video, I found it easy to put it out of my mind.',
        max=7, min=1,
        widget=widgets.RadioSelectHorizontal, choices=[1, 2, 3, 4, 5, 6, 7]
    )
    transport_5 = models.IntegerField(
        label='The video affected me emotionally.',
        max=7, min=1,
        widget=widgets.RadioSelectHorizontal, choices=[1, 2, 3, 4, 5, 6, 7]
    )
    transport_6 = models.IntegerField(
        label='The events in the video have changed my concerns about the future.',
        max=7, min=1,
        widget=widgets.RadioSelectHorizontal, choices=[1, 2, 3, 4, 5, 6, 7]
    )
    char_image_vivid = models.IntegerField(
        label='While watching the video I had a vivid image of the main character.',
        max=7, min=1,
        widget=widgets.RadioSelectHorizontal, choices=[1, 2, 3, 4, 5, 6, 7]
    )
    char_good_bad = models.IntegerField(
        label='How would you rate the main character: <br> <i>Bad (1) – Good (7)</i>',
        max=7, min=1,
        widget=widgets.RadioSelectHorizontal,
        choices=[1, 2, 3, 4, 5, 6, 7]
    )
    char_pleasant_unpleasant = models.IntegerField(
        label='How would you rate the main character: <br> <i>Unpleasant (1) – Pleasant (7)</i>',
        max=7, min=1,
        widget=widgets.RadioSelectHorizontal, choices=[1, 2, 3, 4, 5, 6, 7]
    )
    char_attractive_unattractive = models.IntegerField(
        label='How would you rate the main character: <br> <i>Unattractive (1) – Attractive (7)</i>',
        max=7, min=1,
        widget=widgets.RadioSelectHorizontal, choices=[1, 2, 3, 4, 5, 6, 7]
    )
    char_responsible_irresponsible = models.IntegerField(
        label='How would you rate the main character: <br> <i>Irresponsible (1) – Responsible (7)</i>',
        max=7, min=1,
        widget=widgets.RadioSelectHorizontal, choices=[1, 2, 3, 4, 5, 6, 7]
    )
    belief_just_world = models.IntegerField(label='People usually receive the outcomes that they deserve.',
                                            widget=widgets.RadioSelectHorizontal, choices=[1, 2, 3, 4, 5, 6, 7]
                                            )
    belief_climate_relevant = models.IntegerField(
        label='Climate change is a relevant issue that must be actively addressed.',
        widget=widgets.RadioSelectHorizontal, choices=[1, 2, 3, 4, 5, 6, 7]
        )
    video_watched = models.BooleanField(initial=False)

    time_on_page = models.IntegerField(initial=0)
    pause_count = models.IntegerField(initial=0)

    ai_chat_log = models.LongStringField(blank=True)

    gender = models.StringField(
        label='What is your gender?',
        choices=['Male', 'Female', 'Other']
    )

    age = models.IntegerField(
        label='What is your age?',
        min=18,
        max=99
    )

    employment_status = models.StringField(
        label='Which of the following describes your situation the best?',
        choices=[
            'Unemployed',
            'Student',
            'Working either part-time or full-time',
            'Self-employed',
            'Other'
        ]
    )

    education = models.StringField(
        label='What is the highest level of education you achieved to date?',
        choices=[
            'No formal education',
            'Lower than a high school diploma',
            'High school diploma',
            'Bachelor degree',
            "Master's degree",
            'Doctoral degree'
        ]
    )

    income_us = models.IntegerField(
        label='Imagine an income scale where 1 is the lowest income group and 10 the highest in the United States.',
        min=1,
        max=10
    )

    us_state = models.StringField(
        label='In which US state are you residing?',
        choices=[
            'Alabama', 'Alaska', 'Arizona', 'Arkansas', 'California', 'Colorado', 'Connecticut',
            'Delaware', 'District of Columbia', 'Florida', 'Georgia', 'Hawaii', 'Idaho', 'Illinois',
            'Indiana', 'Iowa', 'Kansas', 'Kentucky', 'Louisiana', 'Maine', 'Maryland', 'Massachusetts',
            'Michigan', 'Minnesota', 'Mississippi', 'Missouri', 'Montana', 'Nebraska', 'Nevada',
            'New Hampshire', 'New Jersey', 'New Mexico', 'New York', 'North Carolina', 'North Dakota',
            'Ohio', 'Oklahoma', 'Oregon', 'Pennsylvania', 'Rhode Island', 'South Carolina',
            'South Dakota', 'Tennessee', 'Texas', 'Utah', 'Vermont', 'Virginia', 'Washington',
            'West Virginia', 'Wisconsin', 'Wyoming'
        ]
    )
    book = models.IntegerField(
        label="How many books, e-books, or audiobooks have you read in the past 12 months? Please include only books read for non-academic or non-professional purposes.",
        blank=False
    )

    book3 = models.StringField(
        choices=[
            ('0', '0'),
            ('1', '1–10'),
            ('2', '11–50'),
            ('3', '50–100'),
            ('4', 'More than 100')
        ],
        label="Approximately how many books do you own at home or on your computer/smartphone, in physical or digital format?",
        blank=False
    )

    teatro = models.IntegerField(
        choices=[(1, 'Never'), (2, '1–3 times'), (3, '4–6 times'), (4, '7–12 times'), (5, 'More than 12 times')],
        label="Theater",
        widget=widgets.RadioSelect
    )

    cinema = models.IntegerField(
        choices=[(1, 'Never'), (2, '1–3 times'), (3, '4–6 times'), (4, '7–12 times'), (5, 'More than 12 times')],
        label="Cinema",
        widget=widgets.RadioSelect
    )

    musei_mostre = models.IntegerField(
        choices=[(1, 'Never'), (2, '1–3 times'), (3, '4–6 times'), (4, '7–12 times'), (5, 'More than 12 times')],
        label="Museums and exhibitions",
        widget=widgets.RadioSelect
    )

    classica_opera = models.IntegerField(
        choices=[(1, 'Never'), (2, '1–3 times'), (3, '4–6 times'), (4, '7–12 times'), (5, 'More than 12 times')],
        label="Classical music concerts and opera",
        widget=widgets.RadioSelect
    )

    altri_concerti = models.IntegerField(
        choices=[(1, 'Never'), (2, '1–3 times'), (3, '4–6 times'), (4, '7–12 times'), (5, 'More than 12 times')],
        label="Other music concerts",
        widget=widgets.RadioSelect
    )

    sport = models.IntegerField(
        choices=[(1, 'Never'), (2, '1–3 times'), (3, '4–6 times'), (4, '7–12 times'), (5, 'More than 12 times')],
        label="Sporting events",
        widget=widgets.RadioSelect
    )

    discoteche = models.IntegerField(
        choices=[(1, 'Never'), (2, '1–3 times'), (3, '4–6 times'), (4, '7–12 times'), (5, 'More than 12 times')],
        label="Nightclubs, dance halls, or other venues for dancing",
        widget=widgets.RadioSelect
    )

    siti_archeologici = models.IntegerField(
        choices=[(1, 'Never'), (2, '1–3 times'), (3, '4–6 times'), (4, '7–12 times'), (5, 'More than 12 times')],
        label="Archaeological sites and monuments",
        widget=widgets.RadioSelect
    )

    # ----------------- FUNCTIONS -----------------#


def creating_session(subsession: Subsession):
    players = subsession.get_players()

    # balanced assignment for narrative_treatment: EN, SN, C
    n = len(players)
    base = ['EN', 'SN', 'C']
    reps = n // 3
    remainder = n % 3
    assignments = base * reps + base[:remainder]
    random.shuffle(assignments)

    for p, t in zip(players, assignments):
        p.narrative_treatment = t

    # balanced assignment for social_context: PE, AE
    base_sc = ['PE', 'AE']
    reps_sc = n // 2
    remainder_sc = n % 2
    contexts = base_sc * reps_sc + base_sc[:remainder_sc]
    random.shuffle(contexts)

    for p, sc in zip(players, contexts):
        p.social_context = sc


# ----------------- PAGES -----------------#


class Consent(Page):
    form_model = 'player'
    form_fields = ['consent']

    @staticmethod
    def before_next_page(player: Player, timeout_happened):
        participant = player.participant
        participant.consent = player.consent


class Intro(Page):
    form_model = 'player'

    @staticmethod
    def is_displayed(player: Player):
        return player.consent


class NarrativeTreatmentPage(Page):
    form_model = 'player'

    @staticmethod
    def is_displayed(player: Player):
        return player.narrative_treatment in ['EN', 'SN'] and player.consent

    @staticmethod
    def vars_for_template(player: Player):
        if player.narrative_treatment == 'EN':
            video_url = 'https://www.youtube.com/embed/0umI5yPcHFY'
            min_watch_seconds = 600
        elif player.narrative_treatment == 'SN':
            video_url = 'https://www.youtube.com/embed/jEjUAnPc2VA'
            min_watch_seconds = 360
        else:
            video_url = None
            min_watch_seconds = 0

        return dict(
            video_url=video_url,
            treatment=player.narrative_treatment,
            min_watch_seconds=min_watch_seconds
        )


class SocialNormsPage(Page):
    form_model = 'player'
    form_fields = [
        'appropriateness_0', 'appropriateness_1', 'appropriateness_2',
        'appropriateness_3', 'appropriateness_4', 'appropriateness_5',
        'appropriateness_6', 'appropriateness_7', 'appropriateness_8',
        'appropriateness_9', 'appropriateness_10'
    ]

    @staticmethod
    def is_displayed(player: Player):
        return player.consent

    @staticmethod
    def vars_for_template(player: Player):
        if player.social_context == 'PE':
            X, Y = 5, 5
        else:  # AE
            X, Y = 10, 0

        # Prepare allocation choices with labels
        allocations = []
        for i in range(11):
            allocations.append({
                'keep': i,
                'give': 10 - i,
                'field_name': f'appropriateness_{i}'
            })

        return dict(X=X, Y=Y, context=player.social_context, allocations=allocations)

    @staticmethod
    def live_method(player: Player, data):
        if data.get("type") == "chat":
            user_message = data.get("message", "").strip()

            if not user_message:
                return {player.id_in_group: {"error": "Empty message."}}

            # Load existing chat log
            chat_log = []
            existing_log = player.field_maybe_none('ai_chat_log')

            if existing_log:
                try:
                    chat_log = json.loads(existing_log)
                except json.JSONDecodeError:
                    chat_log = []

            # Append user message
            chat_log.append({"role": "user", "content": user_message})

            # Call OpenAI — FIXED: use new openai >= 1.0 client API
            try:
                response = client.chat.completions.create(
                    model="gpt-4o-mini",
                    max_tokens=300,
                    messages=[
                        {
                            "role": "system",
                            "content": (
                                "You are assisting a participant in a behavioral experiment about social norms. "
                                "Your role is strictly limited: "
                                "- You can clarify instructions. "
                                "- You can explain what the task means. "
                                "- You can restate what different choices represent. "
                                "You must NOT: "
                                "- Suggest which option to choose. "
                                "- Recommend any behavior. "
                                "- Help maximize earnings. "
                                "- Infer what others might answer. "
                                "- Provide strategic advice of any kind. "
                                "If the user asks for advice, refuse and say: "
                                "'I cannot help you choose an answer, but I can clarify the task.' "
                                "Keep answers short and neutral."
                            )
                        }
                    ] + chat_log
                )

                # FIXED: use attribute access, not dict access
                reply = response.choices[0].message.content

            except Exception as e:
                return {player.id_in_group: {"error": f"AI service unavailable. Please try again. ({str(e)})"}}

            # Append assistant reply
            chat_log.append({"role": "assistant", "content": reply})

            # Save updated log (player.save() is NOT needed in oTree)
            player.ai_chat_log = json.dumps(chat_log)

            return {player.id_in_group: {"reply": reply}}


class ManipulationCheck(Page):
    form_model = 'player'

    @staticmethod
    def get_form_fields(player: Player):
        if player.narrative_treatment == 'EN':
            return ['en_q1', 'en_q2', 'en_q3', 'en_q4', 'en_q5', 'en_q6', 'en_q7', 'en_q8', 'en_q9', 'en_q10']
        elif player.narrative_treatment == 'SN':
            return ['sn_q1', 'sn_q2', 'sn_q3', 'sn_q4', 'sn_q5', 'sn_q6', 'sn_q7', 'sn_q8', 'sn_q9', 'sn_q10']
        else:
            return []

    @staticmethod
    def is_displayed(player: Player):
        return player.narrative_treatment in ['EN', 'SN'] and player.consent


class Transportation(Page):
    form_model = 'player'
    form_fields = [
        'transport_1', 'transport_2', 'transport_3', 'transport_4',
        'transport_5', 'transport_6', 'char_image_vivid', 'char_good_bad',
        'char_pleasant_unpleasant', 'char_attractive_unattractive',
        'char_responsible_irresponsible'
    ]

    @staticmethod
    def is_displayed(player: Player):
        return player.narrative_treatment in ['EN', 'SN'] and player.consent


class Beliefs(Page):
    form_model = 'player'
    form_fields = ['belief_just_world', 'belief_climate_relevant']

    @staticmethod
    def is_displayed(player: Player):
        return player.consent

class Culture(Page):
    form_model = 'player'
    form_fields = ['teatro', 'cinema', 'musei_mostre', 'classica_opera', 'altri_concerti', 'sport', 'discoteche',
                   'siti_archeologici']

    @staticmethod
    def is_displayed(player: Player):
        return player.consent

class Demographics(Page):
    form_model = 'player'
    form_fields = [
        'book',
        'book3',
        'gender',
        'age',
        'employment_status',
        'education',
        'income_us',
        'us_state'
    ]

    @staticmethod
    def is_displayed(player: Player):
        return player.consent


class Thanks(Page):
    @staticmethod
    def is_displayed(player: Player):
        return player.consent


class Noconsent(Page):
    @staticmethod
    def is_displayed(player: Player):
        return not player.consent


# ----------------- SEQUENCE -----------------#

page_sequence = [
    Consent,
    Intro,
    NarrativeTreatmentPage,
    ManipulationCheck,
    Transportation,
    SocialNormsPage,
    Beliefs,
    Demographics,
    Thanks,
    Noconsent
]