import asyncio
import logging
from dotenv import load_dotenv

# Assuming you have mcp-use or a similar library
from mcp_use import MCPAgent, MCPClient 
from langchain_openai import ChatOpenAI # Or your preferred LLM provider

import hud
from hud import gym, Task
from hud.adapters import ResponseAction # For submitting final text to HUD env

# Load environment variables (e.g., API keys)
load_dotenv()

# Configure logging
logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO)
logging.getLogger("hud").setLevel(logging.DEBUG) # See HUD SDK logs

# Example MCP server configuration (replace with your actual server setup)
# This simulates an MCP server for Airbnb search that can be run via npx.
mcp_server_config = {
    "mcpServers": {
        "airbnb_search_tool": {
            "command": "npx",
            "args": ["-y", "@openbnb/mcp-server-airbnb", "--ignore-robots-txt"], # Example package
            "env": {"IGNORE_ROBOTS_TXT": "true"} 
        }
    }
}

# Initialize MCPClient and MCPAgent
mcp_client = MCPClient.from_dict(mcp_server_config)
llm = ChatOpenAI(model="gpt-4o") # Configure your LLM

# This agent will use tools available via MCP servers defined in mcp_client
mcp_agent = MCPAgent(
    llm=llm, 
    client=mcp_client, 
    max_steps=10, # Max MCP tool calls per .run()
    verbose=True # Set to False for less mcp-use library logging
)

# Task for the HUD Browser environment
browser_task = Task(
    prompt=(
        "Find a pet-friendly Airbnb in Barcelona for 2 adults for a week in August "
        "with good reviews. Summarize the top 3 options."
    ),
    gym="hud-browser", # Agent will use MCP tools, then might interact with browser
    setup=("goto", "https://www.google.com/travel/hotels/Barcelona"), # Initial browser page
    evaluate=("response_includes", ["Barcelona", "pool", "pet-friendly"]) # Check final response
)

qa_task = Task(
    prompt=(
        "Using available tools, find a pet-friendly Airbnb in Barcelona for 2 adults for a week "
        "in August with good reviews. Provide a summary of the top option."
    ),
    gym="qa", # No interactive browser needed for the agent's core work
    evaluate=("response_includes", ["Barcelona", "pet-friendly"])
)

# (Continuing from previous setup)

async def run_mcp_agent_with_tracing(task_definition: Task):
    logger.info(f"Starting task: {task_definition.prompt}")

    with hud.trace(name="airbnb_search_mcp_trace", attributes={"city": "Barcelona", "agent_type": "MCPAgent"}):
        # Initialize HUD environment if the task uses one
        env = await gym.make(task_definition)
        await env.reset() # Runs task_definition.setup

        logger.info("Running MCP agent...")
        # The MCPAgent.run() method will make one or more MCP calls to its configured servers.
        # These calls (requests, responses, errors, timings) will be captured by hud.trace.
        try:
            agent_final_output = await mcp_agent.run(
                task_definition.prompt # Agent receives the main prompt
            )
            logger.info(f"MCPAgent final output: {agent_final_output}")
            
            # Submit the agent's final textual output to the HUD environment for evaluation
            # This is crucial if your Task.evaluate depends on ResponseAction
            if agent_final_output and isinstance(agent_final_output, str):
                await env.step([ResponseAction(text=agent_final_output)])
            
        except Exception as e:
            logger.error(f"MCPAgent run failed: {e}")
            agent_final_output = f"Agent error: {e}" # Log error as part of the output
            # Optionally submit error as a response if needed for evaluation logic
            await env.step([ResponseAction(text=agent_final_output)])

        # Evaluate the task based on its definition (e.g., checking browser state or agent_final_output)
        evaluation_result = await env.evaluate()
        logger.info(f"HUD Evaluation Result: {evaluation_result}")
        
        await env.close()
        logger.info("Environment closed.")
        
    # The trace is automatically submitted at the end of the `with` block.
    # View it on app.hud.so/jobs/traces/{trace_id} (trace_id logged by HUD)

# Example of how to run it:
# async def main():
#     # hud.init_telemetry() # Ensure telemetry is initialized
#     await run_mcp_agent_with_tracing(browser_task) # or qa_task
#     await asyncio.sleep(10) # Allow time for background telemetry export
#     hud.flush() # Ensure all telemetry is sent before exiting

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

@hud.register_trace(name="find_accommodation_mcp")
async def find_accommodation(prompt_query: str, agent_instance: MCPAgent):
    logger.info(f"Agent received query: {prompt_query}")
    # All MCP calls made by agent_instance.run() here will be traced.
    result = await agent_instance.run(prompt_query)
    return result

# Later in your code:
# accommodation_summary = await find_accommodation("Book a hotel in Paris", mcp_agent)