import argparse
import os
from typing import List, Tuple
from maihem import Simulator

# Set the Maihem API key as an environment variable
# os.environ['MAIHEM_API_KEY'] = '<YOUR_API_KEY>'

parser = argparse.ArgumentParser()

# Choose identifier of the target agent, to be defined in config.yaml
parser.add_argument("--mode", default="dev", choices=["dev", "test"], type=str)

# Choose identifier of the target agent, to be defined in config.yaml
parser.add_argument("--target_agent_identifier", default="tg_ag_therapy_1", type=str)

# Choose identifier of the maihem agent, to be defined in config.yaml
parser.add_argument("--maihem_agent_identifier", default="mh_ag_venting_1", type=str)

# Choose identifier of the test, to be defined in config.yaml
parser.add_argument("--test_identifier", default="test_a", type=str)

# Path to config.yaml
parser.add_argument("--config_path", default="./src/config.yaml", type=str)


# Example of chat function
def chat_function(conversation_id: str, agent_maihem_message: str) -> Tuple[str, List[str]]:
    """
    Callable chat function to wrap your target agent to be tested.

    params:
    - conversation_id: unique conversation ID generated by Maihem, use it to keep track of different conversations
    - agent_maihem_msg: message from the Maihem Agent

    return:
    - agent_target_message [str] - message from your Target Agent in response to the Maihem Agent message
    - contexts [List[str]] - list of contexts (required for RAG evaluations), pass empty list if not needed
    """
    
    # The message from your Target Agent. 
    # Connect your agent here to respond to agent_maihem_message. 
    # Keep track of same or new conversations with conversation_id
    agent_target_message = "Hi, how can I help you today?"
    
    # List of contexts for RAG evaluations, pass empty list if not needed
    contexts = ["Context 1", "Context 2"] 
    # contexts = ["<Context 1>", "<Context 2>"] 
    
    # Delete when you have implemented your chat function
    raise NotImplementedError("Chat function has not been implemented")
    
    return agent_target_message, contexts


if __name__ == "__main__":
    
    if parser.parse_args().mode == "dev":
        conversation = Simulator.conversation(
            chat_function, 
            target_agent_identifier=parser.parse_args().target_agent_identifier, 
            maihem_agent_identifier=parser.parse_args().maihem_agent_identifier, 
            config_path=parser.parse_args().config_path
        )
        
        print(conversation.messages)
        print(conversation.evaluation)
        
    elif parser.parse_args().mode == "test":  
        conversations = Simulator.test(
            chat_function, 
            target_agent_identifier=parser.parse_args().target_agent_identifier, 
            test_identifier=parser.parse_args().test_identifier, 
            config_path=parser.parse_args().config_path
        )
        
        print(conversations[0].messages)
        print(conversations[0].evaluation)
        
    else:
        raise ValueError("Invalid mode. Choose 'dev' or 'test'")
