from collections.abc import Generator
from typing import Any, Optional


def process_and_print_streaming_response(response: Generator[dict[str, Any], None, None]) -> Optional[Any]:
    content = ''
    last_sender = ''

    for chunk in response:
        if 'sender' in chunk:
            last_sender = chunk['sender']

        if 'content' in chunk and chunk['content'] is not None:
            if not content and last_sender:
                print(f'\033[94m{last_sender}:\033[0m', end=' ', flush=True)
                last_sender = ''
            print(chunk['content'], end='', flush=True)
            content += chunk['content']

        if 'tool_calls' in chunk and chunk['tool_calls'] is not None:
            for tool_call in chunk['tool_calls']:
                f = tool_call['function']
                name = f['name']
                if not name:
                    continue
                print(f'\033[94m{last_sender}: \033[95m{name}\033[0m()')

        if 'delim' in chunk and chunk['delim'] == 'end' and content:
            print()  # End of response message
            content = ''

        if 'response' in chunk:
            return chunk['response']

    return None


def run_demo_loop(
    client, starting_agent, context_variables: Optional[dict[str, Any]] = None, debug: bool = False
) -> None:
    print('Starting Swarm CLI 🐝')

    messages = []
    agent = starting_agent

    while True:
        user_input = input('\033[90mUser\033[0m: ')
        messages.append({'role': 'user', 'content': user_input})

        response = client.run(
            agent=agent,
            messages=messages,
            context_variables=context_variables or {},
            stream=True,
            debug=debug,
        )

        response = process_and_print_streaming_response(response)

        messages.extend(response.messages)
        agent = response.agent
