import logging
import random
from typing import Annotated

# Import necessary components from autogen
from autogen import (
    ConversableAgent,
    GroupChat,
    GroupChatManager,
    LLMConfig,
    UserProxyAgent,
    config_list_from_json,
)
from autogen.tools.experimental.reliable import ReliableTool

# Configure logging
# Set level to DEBUG to see more internal details if needed
# logging.basicConfig(level=logging.DEBUG, format='%(asctime)s [%(levelname)s] %(name)s: %(message)s')
logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(name)s: %(message)s")
logger = logging.getLogger(__name__)

print("Successfully imported components from autogen.")

should_error_for_demo = True

def buy_airplane_ticket(
    from_location: Annotated[str, "The departure city."],
    to_location: Annotated[str, "The destination city."],
    date: Annotated[str, "The date of travel (e.g., 'April 12, 2025')."],
) -> str:
    """
    Simulates buying an airplane ticket.

    Args:
        from_location: The departure city.
        to_location: The destination city.
        date: The date of travel.

    Returns:
        A confirmation string with a random ticket number.
    """
    global should_error_for_demo
    if should_error_for_demo:
        should_error_for_demo = not should_error_for_demo
        return "ERROR: Something weird happened."
    logger.info(f"Attempting to buy ticket: {from_location} -> {to_location} on {date}")
    ticket_number = random.randint(1000, 9999)
    confirmation = f"""Your ticket from {from_location} to {to_location} on {date} has been booked.
Your ticket number is {ticket_number}.
Please keep this number for future reference."""
    logger.info(f"Ticket bought successfully: {ticket_number}")
    return confirmation

def cancel_airplane_ticket(ticket_number: Annotated[str, "The 4-digit ticket number to cancel."]) -> str:
    """
    Simulates cancelling an airplane ticket.

    Args:
        ticket_number: The ticket number to cancel.

    Returns:
        A cancellation confirmation string.
    """
    logger.info(f"Attempting to cancel ticket: {ticket_number}")
    if not (isinstance(ticket_number, str) and ticket_number.isdigit() and len(ticket_number) == 4):
        logger.warning(f"Invalid ticket number format for cancellation: {ticket_number}")
        return f"Error: Invalid ticket number format '{ticket_number}'. Please provide a 4-digit number."

    confirmation = f"Your ticket with ticket number {ticket_number} has been successfully canceled."
    logger.info(f"Ticket cancelled successfully: {ticket_number}")
    return confirmation

llm_config: LLMConfig | None = None
try:
    config_list = config_list_from_json(
        "OAI_CONFIG_LIST",
    )
    llm_config = LLMConfig(
        config_list=config_list,
        temperature=0.7,
    )
    print(f"Using LLM config: {config_list[0].get('base_url', 'Default Endpoint')}, Model: {config_list[0]['model']}")
except Exception as e_config:
    print(f"Error creating LLM config: {e_config}. Tools cannot be initialized.")
    llm_config = None
buy_system_message_addition_for_tool_calling = """You are an assistant that prepares function calls.
The user (the outer agent) has provided the necessary arguments: `from_location`, `to_location`, and `date`.
Your task is to invoke the internal tool `buy_airplane_ticket` using these exact arguments.
You MUST invoke the `buy_airplane_ticket` tool exactly once per response using the provided arguments.
Do not add conversational text. Output only the tool call.
"""

# Validator: Instructs the internal LLM on how to validate the result returned by buy_airplane_ticket.
buy_system_message_addition_for_result_validation = """You are a validation assistant. Your task is to validate the output received from the `buy_airplane_ticket` tool.

**Input:** You will receive the `arguments` passed to the tool and the `tool_output` string returned by it.

**Validation Criteria:**
1.  **Success Indication:** The `tool_output` should contain the phrase "has been booked" and "Your ticket number is".
2.  **Content Presence:** The `tool_output` string should not be empty or indicate an obvious error (unless an error was expected, which is not the case here).
3.  **Argument Reflection:** The output should ideally reflect the `from_location`, `to_location`, and `date` from the input `arguments`. (This might be hard for the LLM to perfectly verify from just the args and output string, focus on criteria 1 & 2).

Analyze the `tool_output` based on the `arguments`. Respond ONLY with your validation assessment in the required JSON format:
{"validation_result": boolean, "justification": "string"}
- If criteria 1 and 2 are met, set `validation_result` to `true`.
- If any criterion is not met, set `validation_result` to `false` and state the reason (e.g., "Confirmation phrase missing", "Output is empty", "Output indicates an error").
Do not add any other text before or after the JSON.
"""

reliable_buy_tool: ReliableTool | None = None
if llm_config:
    try:
        reliable_buy_tool = ReliableTool(
            name="ReliableBuyTicket",
            func_or_tool=buy_airplane_ticket,  # Pass the function itself
            description="Reliably buys an airplane ticket using provided location, destination, and date, with validation.",
            runner_llm_config=llm_config,
            validator_llm_config=llm_config,
            system_message_addition_for_tool_calling=buy_system_message_addition_for_tool_calling,
            system_message_addition_for_result_validation=buy_system_message_addition_for_result_validation,
            ground_truth=["####Test ground truth entry####"],
        )
        logger.info("ReliableTool for buying tickets created successfully.")
    except Exception as e_reliable_init:
        logger.error(f"Error creating ReliableTool for buying tickets: {e_reliable_init}", exc_info=True)
else:
    logger.warning("LLM Configuration not ready. Cannot create ReliableTool for buying.")

# --- ReliableTool for Cancelling Tickets ---

# Runner: Similar to the buy runner, it confirms the call to the internal function.
cancel_system_message_addition_for_tool_calling = """You are an assistant that prepares function calls.
The user (the outer agent) has provided the necessary argument: `ticket_number`.
Your task is to invoke the internal tool `cancel_airplane_ticket` using this exact argument.
You MUST invoke the `cancel_airplane_ticket` tool exactly once per response using the provided argument.
Do not add conversational text. Output only the tool call.
"""

# Validator: Validates the result of cancel_airplane_ticket.
cancel_system_message_addition_for_result_validation = """You are a validation assistant. Your task is to validate the output received from the `cancel_airplane_ticket` tool.

**Input:** You will receive the `arguments` passed to the tool (containing `ticket_number`) and the `tool_output` string returned by it.

**Validation Criteria:**
1.  **Success Indication:** The `tool_output` should contain the phrase "has been successfully canceled".
2.  **Error Handling:** If the `tool_output` contains "Error: Invalid ticket number format", the function call technically worked but failed validation, so validation should be `false`.
3.  **Content Presence:** The `tool_output` string should not be empty.

Analyze the `tool_output` based on the `arguments`. Respond ONLY with your validation assessment in the required JSON format:
{"validation_result": boolean, "justification": "string"}
- If criteria 1 and 3 are met, and criterion 2 is NOT met, set `validation_result` to `true`.
- Otherwise, set `validation_result` to `false` and state the reason (e.g., "Cancellation confirmation missing", "Output is empty", "Output indicates invalid ticket number error").
Do not add any other text before or after the JSON.
"""

reliable_cancel_tool: ReliableTool | None = None
if llm_config:
    try:
        reliable_cancel_tool = ReliableTool(
            name="ReliableCancelTicket",
            func_or_tool=cancel_airplane_ticket,  # Pass the function itself
            description="Reliably cancels an airplane ticket using the provided ticket number, with validation.",
            runner_llm_config=llm_config,
            validator_llm_config=llm_config,
            system_message_addition_for_tool_calling=cancel_system_message_addition_for_tool_calling,
            system_message_addition_for_result_validation=cancel_system_message_addition_for_result_validation,
        )
        logger.info("ReliableTool for cancelling tickets created successfully.")
    except Exception as e_reliable_init:
        logger.error(f"Error creating ReliableTool for cancelling tickets: {e_reliable_init}", exc_info=True)
else:
    logger.warning("LLM Configuration not ready. Cannot create ReliableTool for cancelling.")

if not llm_config:
    raise ValueError("LLM Config is not set up. Cannot proceed.")

sales_agent = ConversableAgent(
    name="SalesAgent",
    system_message="You are a sales agent responsible for helping customers buy airplane tickets. "
    "Use the 'ReliableBuyTicket' tool when a user wants to buy a ticket. "
    "Extract the from_location, to_location, and date from the user's request for the tool.",
    llm_config=llm_config,
)

cancellation_agent = ConversableAgent(
    name="CancellationAgent",
    system_message="You are a cancellation agent responsible for helping customers cancel airplane tickets. "
    "Use the 'ReliableCancelTicket' tool when a user wants to cancel a ticket. "
    "Extract the ticket_number from the user's request for the tool.",
    llm_config=llm_config,
)

user_proxy = UserProxyAgent(
    name="UserProxy",
    human_input_mode="ALWAYS",  # Ask user for confirmation before executing any tool
    code_execution_config=False,  # Disable code execution for safety
)
if reliable_buy_tool and reliable_cancel_tool and llm_config:
    reliable_buy_tool.register_tool(sales_agent)
    reliable_cancel_tool.register_tool(cancellation_agent)
    logger.info("Registered ReliableTools using agent.register_tools method.")
else:
    logger.error("One or more ReliableTools are not available for registration.")

if not reliable_buy_tool or not reliable_cancel_tool:
    print("\nSkipping chat initiation because one or more ReliableTools failed to initialize.")
else:
    groupchat = GroupChat(
        agents=[user_proxy, cancellation_agent, sales_agent],
        messages=[],
        speaker_selection_method="auto",  # LLM decides who speaks next
        allow_repeat_speaker=False,  # Prevent back-and-forth loops between two agents
    )

    manager = GroupChatManager(
        name="GroupManager",
        groupchat=groupchat,
        llm_config=llm_config,
        # My local qwen3 needs this so it doesn't get chatty when figuring out who to pick
        system_message="/no_think You are the manager of a chat between a user, a SalesAgent, and a CancellationAgent. "
        "Coordinate the conversation. If the user wants to buy a ticket, ensure the SalesAgent responds. "
        "If the user wants to cancel a ticket, ensure the CancellationAgent responds. "
        "The UserProxy executes tools after confirmation.",
    )
    print("\n\n--- Initiating Chat: Buy Ticket ---")
    user_proxy.initiate_chat(
        recipient=manager,
        message="Hi, I need to buy a plane ticket from London Heathrow to Tokyo Haneda for May 1st, 2026.",
        clear_history=True,
    )