import deliberate_lab as dl

client = dl.Client(env="prod")  # Use env="dev" for local emulator

# Create experiment with just a name
result = client.create_experiment(name="My First Experiment")
print(f"Created experiment: {result['experiment']['id']}")

import deliberate_lab as dl

client = dl.Client(env="prod")

# Define stage configurations
stages = [
    # Stage 1: Profile setup (anonymous animal names)
    dl.ProfileStageConfig(
        id="profile",
        kind="profile",
        name="Profile Setup",
        descriptions={"primaryText": "Set up your profile", "infoText": "", "helpText": ""},
        progress={"minParticipants": 1, "waitForAllParticipants": False, "showParticipantProgress": False},
        profileType=dl.ProfileType.ANONYMOUS_ANIMAL,
    ),
    
    # Stage 2: Pre-discussion survey
    dl.SurveyStageConfig(
        id="pre-survey",
        kind="survey",
        name="Pre-Discussion Survey",
        descriptions={"primaryText": "Answer a few questions before the discussion", "infoText": "", "helpText": ""},
        progress={"minParticipants": 1, "waitForAllParticipants": False, "showParticipantProgress": False},
        questions=[
            dl.TextSurveyQuestion(
                id="q1",
                kind="text",
                questionTitle="What topic would you like to discuss today?",
            ),
            dl.ScaleSurveyQuestion(
                id="q2",
                kind="scale",
                questionTitle="How familiar are you with group discussions?",
                lowerValue=1,
                lowerText="Not at all",
                upperValue=5,
                upperText="Very familiar",
            ),
            dl.MultipleChoiceSurveyQuestion(
                id="q3",
                kind="mc",
                questionTitle="What's your preferred discussion style?",
                options=[
                    dl.MultipleChoiceItem(id="opt1", text="Collaborative", imageId=""),
                    dl.MultipleChoiceItem(id="opt2", text="Debate-oriented", imageId=""),
                    dl.MultipleChoiceItem(id="opt3", text="Listener", imageId=""),
                ],
            ),
        ],
    ),
    
    # Stage 3: Group chat (10 minute discussion)
    dl.ChatStageConfig(
        id="discussion",
        kind="chat",
        name="Group Discussion",
        descriptions={"primaryText": "Discuss the topic with your group", "infoText": "", "helpText": ""},
        progress={"minParticipants": 2, "waitForAllParticipants": True, "showParticipantProgress": True},
        timeLimitInMinutes=10,
    ),
]

# Create the experiment
result = client.create_experiment(
    name="Demo Discussion Study",
    description="A demo experiment with profile, survey, and chat stages",
    stages=stages,
)

experiment_id = result["experiment"]["id"]
print(f"Created experiment: {experiment_id}")

import deliberate_lab as dl

client = dl.Client(env="prod")

# Define the AI agent persona and prompt
agent_mediator = dl.AgentMediatorTemplate(
    persona=dl.Persona(
        id="helper-agent",
        name="Helper",
        defaultModelSettings=dl.AgentModelSettings(
            apiType=dl.ApiKeyType.GEMINI,
            modelName="gemini-2.0-flash",
        ),
    ),
    promptMap={
        "discussion": dl.ChatPromptConfig(
            id="discussion",
            type=dl.ChatStageType.chat,
            prompt=[
                dl.TextPromptItem(
                    type="TEXT",
                    text="""You are a helpful discussion facilitator. 
Your role is to:
- Help keep the conversation on track
- Ask clarifying questions when needed  
- Summarize key points periodically
- Stay neutral and encourage all participants

Only speak when it adds value. Keep responses concise (1-2 sentences).""",
                ),
            ],
            chatSettings=dl.AgentChatSettings(
                minMessagesBeforeResponding=3,
                canSelfTriggerCalls=True,
                initialMessage="",
            ),
        ),
    },
)

# Create experiment with stages (from previous section) and agent
result = client.create_experiment(
    name="Demo with AI Facilitator",
    description="A discussion experiment with an AI facilitator",
    stages=stages,  # stages from previous section
    agent_mediators=[agent_mediator],
)

print(f"Created experiment with agent: {result['experiment']['id']}")

# Create a cohort for 3-5 participants
cohort_result = client.create_cohort(
    experiment_id=experiment_id,
    name="Cohort A",
    description="First group of participants",
    participant_config=dl.CohortParticipantConfig(
        minParticipantsPerCohort=3,
        maxParticipantsPerCohort=5,
        includeAllParticipantsInCohortCount=True,
        botProtection=True,
    ),
)

cohort_id = cohort_result["cohort"]["id"]
print(f"Created cohort: {cohort_id}")

# List all cohorts
cohorts = client.list_cohorts(experiment_id)
print(f"Total cohorts: {cohorts['total']}")

# Export full experiment data
data = client.export_experiment(experiment_id)

# Access exported structure (uses Map format)
experiment = data["experiment"]
stage_map = data["stageMap"]
cohort_map = data["cohortMap"]
participant_map = data["participantMap"]

print(f"Experiment: {experiment['metadata']['name']}")
print(f"Stages: {len(stage_map)}")
print(f"Cohorts: {len(cohort_map)}")
print(f"Total participants: {len(participant_map)}")

# Save to file
import json
with open("experiment_export.json", "w") as f:
    json.dump(data, f, indent=2)