from typing import Tuple, List
import os
import random

from maihem.clients import Maihem
from maihem.utils import extract_text


# Create a Maihem client by providing your API key. 
# Add API key to MAIHEM_API_KEY environment variable.
# os.environ["MAIHEM_API_KEY"] = api_key
maihem_client = Maihem()

# Create (and describe) the target agent you want to test
try:
    target_agent = maihem_client.create_target_agent(
        identifier="tg_agent_pg_essays",
        role="question answering",
        industry="entrepreneurship",
        description="A chatbot that answers questions about entrepreneurship and startups, based in the essays of Paul Graham, the co-founder of Y Combinator.",
    )
except Exception as e:
    # or get target agent if it exists
    target_agent = maihem_client.get_target_agent(identifier="tg_agent_pg_essays")

# 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>"] 
    
    return agent_target_message, contexts
    
# Set the chat function to the Target Agent, so we can automatically chat with your Target Agent
target_agent.set_chat_function(chat_function)

# For RAG testing, add a list of documents so Maihem can generate questions from them to evaluate your Target Agent
target_agent.add_documents(documents=["./src/test_data/RAG_PG_essays.pdf"])

# Select the Metrics you want to evaluate your agent on, 
# and how many conversations you want to run for each metric.
# You can find the list of supported metrics here: https://docs.maihem.ai/metrics/collection
metrics_config = {
    "qa_cx_helpfulness": 2, # Number of conversations to run for this metric
    "qa_cx_goal_completion": 2, # Number of conversations to run for this metric
    "qa_rag_hallucination": 2, # Number of conversations to run for this metric
    "qa_rag_answer_relevance": 2, # Number of conversations to run for this metric
}

# Create a Test
test = maihem_client.create_test(
    identifier="test_pg_essays_rag",
    metrics_config=metrics_config,
    maihem_agent_behavior_prompt="Ask questions about entrepreneurship and startups based on the essays of Paul Graham. Image you are a young entrepreneur looking for advice, deciding wheter to drop out of college to start a company.",
    conversation_turns_max=5
)

# # Create a Test Run from a Test and a Target Agent
test_run = maihem_client.create_test_run(
    test_identifier="test_pg_essays_rag", # The identifier of the Test you want to run
    target_agent=target_agent, # The Target Agent you want to test
    concurrent_conversations=10, # How many conversations to run concurrently
)

# Get the results of the Test Run (status and how many conversations passed or failed)
results = maihem_client.get_test_run_result(test_run_id=test_run.id)
print(results)

# Get the conversations generated in the Test Run
conversations = maihem_client.get_test_run_result_conversations(test_run_id=test_run.id)
print(conversations)
