from autogen_agentchat.agents import AssistantAgent
from autogen_ext.models.openai import OpenAIChatCompletionClient
from autogen_ext.code_executors.docker_jupyter import DockerJupyterCodeExecutor, DockerJupyterServer
from autogen_ext.tools.code_execution import PythonCodeExecutionTool
from autogen_agentchat.ui import Console
from autogen_core.code_executor import CodeBlock
from autogen_core import CancellationToken
from autogen_agentchat.teams import RoundRobinGroupChat
from autogen_agentchat.conditions import TextMessageTermination

# Download the dataset from https://www.kaggle.com/datasets/nelgiriyewithana/top-spotify-songs-2023
# and place it the coding directory as `spotify-2023.csv`.
bind_dir = "./coding"

# Use a custom docker image with the Jupyter kernel gateway and data science libraries installed.
# Custom docker image: ds-kernel-gateway:latest -- you need to build this image yourself.
# Dockerfile:
# FROM quay.io/jupyter/docker-stacks-foundation:latest
# 
# # ensure that 'mamba' and 'fix-permissions' are on the PATH
# SHELL ["/bin/bash", "-o", "pipefail", "-c"]
# 
# # Switch to the default notebook user
# USER ${NB_UID}
# 
# # Install data-science packages + kernel gateway
# RUN mamba install --quiet --yes \
#     numpy \
#     pandas \
#     scipy \
#     matplotlib \
#     scikit-learn \
#     seaborn \
#     jupyter_kernel_gateway \
#     ipykernel \
#     && mamba clean --all -f -y \
#     && fix-permissions "${CONDA_DIR}" \
#     && fix-permissions "/home/${NB_USER}"
# 
# # Allow you to set a token at runtime (or leave blank for no auth)
# ENV TOKEN=""
# 
# # Launch the Kernel Gateway, listening on all interfaces,
# # with the HTTP endpoint for listing kernels enabled
# CMD ["python", "-m", "jupyter", "kernelgateway", \
#     "--KernelGatewayApp.ip=0.0.0.0", \
#     "--KernelGatewayApp.port=8888", \
#     # "--KernelGatewayApp.auth_token=${TOKEN}", \
#     "--JupyterApp.answer_yes=true", \
#     "--JupyterWebsocketPersonality.list_kernels=true"]
# 
# EXPOSE 8888
# 
# WORKDIR "${HOME}"

async def main():
    model = OpenAIChatCompletionClient(model="gpt-4.1")
    async with DockerJupyterServer(
        custom_image_name="ds-kernel-gateway:latest", 
        bind_dir=bind_dir,
    ) as server:
        async with DockerJupyterCodeExecutor(jupyter_server=server) as code_executor:
            await code_executor.execute_code_blocks([
                CodeBlock(code="import pandas as pd\ndf = pd.read_csv('/workspace/spotify-2023.csv', encoding='latin-1')", language="python"),
            ],
                cancellation_token=CancellationToken(),
            )
            tool = PythonCodeExecutionTool(
                executor=code_executor,
            )
            assistant = AssistantAgent(
                "assistant",
                model_client=model,
                system_message="You have access to a Jupyter kernel. Do not write all code at once. Write one code block, observe the output, and then write the next code block.",
                tools=[tool],
            )
            team = RoundRobinGroupChat(
                [assistant],
                termination_condition=TextMessageTermination(source="assistant"),
            )
            task = f"Datafile has been loaded as variable `df`. First preview dataset. Then answer the following question: What is the highest streamed artist in the dataset?"
            await Console(team.run_stream(task=task))

if __name__ == "__main__":
    import asyncio
    asyncio.run(main())