"""
External Tool Execution Usage Examples

Demonstrates how to use external tool execution with HITL (Human-in-the-Loop).
External tools pause execution waiting for external completion before resuming.

Note: HITL continuation (continue_run_async) only supports direct call mode.
Streaming mode is not supported for continuation.
"""

import pytest
import asyncio
import os
from upsonic import Agent, Task
from upsonic.tools import tool
from upsonic.db.database import SqliteDatabase

pytestmark = pytest.mark.timeout(300)


def cleanup_db():
    """Clean up test database files."""
    if os.path.exists("external.db"):
        os.remove("external.db")


# ============================================================================
# EXTERNAL TOOLS DEFINITION
# ============================================================================

@tool(external_execution=True)
def send_email(to: str, subject: str, body: str) -> str:
    """
    Send an email - requires external execution.
    
    Args:
        to: Email recipient
        subject: Email subject
        body: Email body content
        
    Returns:
        Confirmation message
    """
    print("send_email called")
    # In a real implementation, this would call an email service
    return f"Email sent successfully to {to} with subject '{subject}'"


@tool(external_execution=True)
def execute_database_query(query: str) -> str:
    """
    Execute a database query - requires external execution.
    
    Args:
        query: SQL query to execute
        
    Returns:
        Query results
    """
    print("execute_database_query called")
    # In a real implementation, this would execute the query
    return f"Query executed: {query} | Results: 10 rows returned"


@tool(external_execution=True)
def call_external_api(endpoint: str, payload: dict = None) -> dict:
    """
    Call an external API - requires external execution.
    
    Args:
        endpoint: API endpoint URL
        payload: Request payload
        
    Returns:
        API response
    """
    print("call_external_api called")
    # In a real implementation, this would call the API
    return {"status": "success", "data": {"message": f"API called at {endpoint}"}}


# ============================================================================
# HITL HANDLER
# ============================================================================

def _execute_tool(requirement) -> str:
    """Execute an external tool and return the result string."""
    tool_exec = requirement.tool_execution
    tool_name = tool_exec.tool_name
    tool_args = tool_exec.tool_args
    
    if tool_name == "send_email":
        return send_email(**tool_args)
    elif tool_name == "execute_database_query":
        return execute_database_query(**tool_args)
    elif tool_name == "call_external_api":
        result = call_external_api(**tool_args)
        return str(result) if isinstance(result, dict) else result
    else:
        raise ValueError(f"Unknown tool: {tool_name}")


def hitl_handler(requirement) -> None:
    """
    Unified HITL handler for all pause types.
    
    Mutates the requirement in-place based on its type:
    - External execution: sets tool_execution.result
    - Confirmation: calls confirm() or reject()
    - User input: fills user_input_schema field values
    """
    if requirement.needs_external_execution:
        result = _execute_tool(requirement)
        requirement.tool_execution.result = result


# ============================================================================
# VARIANT 1: Direct Call with run_id - Same Agent
# ============================================================================

async def external_direct_call_with_run_id_same_agent():
    """
    Direct call mode with external tool using run_id and same agent instance.
    
    The run_id is obtained from the output of print_do_async.
    """
    agent = Agent("openai/gpt-4o-mini", name="external_tool_agent")
    task = Task(
        description="Send an email to test@example.com with subject 'Hello' and body 'Test message'.",
        tools=[send_email]
    )
    
    output = await agent.print_do_async(task, return_output=True)
    
    for requirement in output.active_requirements:
        if requirement.is_external_tool_execution:
            result = _execute_tool(requirement)
            requirement.tool_execution.result = result
    
    result = await agent.continue_run_async(run_id=output.run_id, return_output=True)
    return result


async def external_direct_call_with_task_same_agent():
    """
    Direct call mode with external tool using task and same agent instance.
    
    Uses in-memory context for continuation.
    """
    agent = Agent("openai/gpt-4o-mini", name="external_tool_agent")
    task = Task(
        description="Send an email to test@example.com with subject 'Hello' and body 'Test message'.",
        tools=[send_email]
    )
    
    output = await agent.print_do_async(task, return_output=True)
    
    for requirement in output.active_requirements:
        if requirement.is_external_tool_execution:
            result = _execute_tool(requirement)
            requirement.tool_execution.result = result
    
    result = await agent.continue_run_async(task=task, return_output=True)
    return result


# ============================================================================
# VARIANT 2A: Direct Call with task - New Agent (Cross-process resumption)
# ============================================================================

async def external_direct_call_with_task_new_agent():
    """
    Direct call mode with external tool using task and new agent instance.
    
    Simulates cross-process resumption where a new agent instance uses
    the task object for continuation (in-memory context).
    
    Note: This uses task parameter, so it relies on in-memory context.
    For true cross-process, use run_id variant instead.
    """
    cleanup_db()
    db = SqliteDatabase(db_file="external.db", session_id="session_1", user_id="user_1")
    agent = Agent("openai/gpt-4o-mini", name="external_tool_agent", db=db)
    task = Task(
        description="Send an email to test@example.com with subject 'Hello' and body 'Test message'.",
        tools=[send_email]
    )
    
    output = await agent.print_do_async(task, return_output=True)
    
    for requirement in output.active_requirements:
        if requirement.is_external_tool_execution:
            result = _execute_tool(requirement)
            requirement.tool_execution.result = result
    
    # New agent uses task for continuation
    # Pass requirements (with results) so new agent can inject them into loaded state
    new_agent = Agent("openai/gpt-4o-mini", name="external_tool_agent", db=db)
    result = await new_agent.continue_run_async(
        task=task, 
        requirements=output.requirements,
        return_output=True
    )
    return result


# ============================================================================
# VARIANT 2B: Direct Call with run_id - New Agent (Cross-process resumption)
# ============================================================================

async def external_direct_call_with_run_id_new_agent():
    """
    Direct call mode with external tool using run_id and new agent instance.
    
    Simulates cross-process resumption where a new agent instance loads
    the paused run from storage using run_id.
    """
    cleanup_db()
    db = SqliteDatabase(db_file="external.db", session_id="session_1", user_id="user_1")
    agent = Agent("openai/gpt-4o-mini", name="external_tool_agent", db=db, debug=True)
    task = Task(
        description="Send an email to test@example.com with subject 'Hello' and body 'Test message'.",
        tools=[send_email]
    )
    
    output = await agent.print_do_async(task, return_output=True)
    run_id = output.run_id
    
    # Execute tools and set results on in-memory requirements
    for requirement in output.active_requirements:
        if requirement.is_external_tool_execution:
            result = _execute_tool(requirement)
            requirement.tool_execution.result = result
    
    # New agent loads from storage and injects results from passed requirements
    new_agent = Agent("openai/gpt-4o-mini", name="external_tool_agent", db=db)
    result = await new_agent.continue_run_async(
        run_id=run_id,
        requirements=output.requirements,
        return_output=True
    )
    return result


# ============================================================================
# VARIANT 3A: Using hitl_handler parameter with task
# ============================================================================

async def external_with_executor_callback_task():
    """
    Use the hitl_handler parameter with task parameter to handle tools automatically.
    
    When provided, if the agent pauses again with NEW HITL requirements,
    the handler is called automatically for each requirement.
    """
    agent = Agent("openai/gpt-4o-mini", name="external_tool_agent")
    task = Task(
        description="Send an email to test@example.com with subject 'Hello' and body 'Test message'.",
        tools=[send_email]
    )
    
    output = await agent.print_do_async(task, return_output=True)
    
    for requirement in output.active_requirements:
        if requirement.is_external_tool_execution:
            result = _execute_tool(requirement)
            requirement.tool_execution.result = result
    
    result = await agent.continue_run_async(
        task=task, 
        return_output=True,
        hitl_handler=hitl_handler,
    )
    return result


# ============================================================================
# VARIANT 3B: Using hitl_handler parameter with run_id
# ============================================================================

async def external_with_executor_callback():
    """
    Use the hitl_handler parameter to handle tools automatically.
    
    When provided, if the agent pauses again with NEW HITL requirements,
    the handler is called automatically for each requirement.
    """
    agent = Agent("openai/gpt-4o-mini", name="external_tool_agent")
    task = Task(
        description="Send an email to test@example.com with subject 'Hello' and body 'Test message'.",
        tools=[send_email]
    )
    
    output = await agent.print_do_async(task, return_output=True)
    
    for requirement in output.active_requirements:
        if requirement.is_external_tool_execution:
            result = _execute_tool(requirement)
            requirement.tool_execution.result = result
    
    result = await agent.continue_run_async(
        run_id=output.run_id, 
        return_output=True,
        hitl_handler=hitl_handler
    )
    return result


# ============================================================================
# VARIANT 4A: Multiple External Tools with task parameter
# ============================================================================

async def external_multiple_tools_direct_call_task():
    """
    Multiple external tools executed in a single run with task parameter.
    
    Uses a loop to handle sequential external tool pauses.
    """
    agent = Agent("openai/gpt-4o-mini", name="external_tool_agent")
    task = Task(
        description=(
            "First, send an email to admin@example.com with subject 'Report' and body 'Monthly report'. "
            "Then query the database with 'SELECT * FROM users'. "
            "Finally, call the external API at https://api.example.com/data."
        ),
        tools=[send_email, execute_database_query, call_external_api]
    )
    
    output = await agent.print_do_async(task, return_output=True)
    
    while output.active_requirements:
        for requirement in output.active_requirements:
            if requirement.is_external_tool_execution:
                result = _execute_tool(requirement)
                requirement.tool_execution.result = result
        
        output = await agent.continue_run_async(
            task=task, 
            return_output=True
        )
    
    return output


async def external_multiple_tools_with_executor_task():
    """
    Multiple external tools with automatic executor callback using task parameter.
    
    The hitl_handler handles all subsequent tool pauses automatically.
    """
    agent = Agent("openai/gpt-4o-mini", name="external_tool_agent")
    task = Task(
        description=(
            "First, send an email to admin@example.com with subject 'Report' and body 'Monthly report'. "
            "Then query the database with 'SELECT * FROM users'. "
            "Finally, call the external API at https://api.example.com/data."
        ),
        tools=[send_email, execute_database_query, call_external_api]
    )
    
    output = await agent.print_do_async(task, return_output=True)
    
    for requirement in output.active_requirements:
        if requirement.is_external_tool_execution:
            result = _execute_tool(requirement)
            requirement.tool_execution.result = result
    
    result = await agent.continue_run_async(
        task=task, 
        return_output=True,
        hitl_handler=hitl_handler
    )
    
    return result


# ============================================================================
# VARIANT 4B: Multiple External Tools in Single Run with run_id
# ============================================================================

async def external_multiple_tools_direct_call():
    """
    Multiple external tools executed in a single run with direct call mode.
    
    Uses a loop to handle sequential external tool pauses.
    """
    agent = Agent("openai/gpt-4o-mini", name="external_tool_agent")
    task = Task(
        description=(
            "First, send an email to admin@example.com with subject 'Report' and body 'Monthly report'. "
            "Then query the database with 'SELECT * FROM users'. "
            "Finally, call the external API at https://api.example.com/data."
        ),
        tools=[send_email, execute_database_query, call_external_api]
    )
    
    output = await agent.print_do_async(task, return_output=True)
    
    while output.active_requirements:
        for requirement in output.active_requirements:
            if requirement.is_external_tool_execution:
                result = _execute_tool(requirement)
                requirement.tool_execution.result = result
        
        output = await agent.continue_run_async(
            run_id=output.run_id, 
            return_output=True
        )
    
    return output


async def external_multiple_tools_with_executor():
    """
    Multiple external tools with automatic executor callback.
    
    The hitl_handler handles all subsequent tool pauses automatically.
    """
    agent = Agent("openai/gpt-4o-mini", name="external_tool_agent")
    task = Task(
        description=(
            "First, send an email to admin@example.com with subject 'Report' and body 'Monthly report'. "
            "Then query the database with 'SELECT * FROM users'. "
            "Finally, call the external API at https://api.example.com/data."
        ),
        tools=[send_email, execute_database_query, call_external_api]
    )
    
    output = await agent.print_do_async(task, return_output=True)
    
    for requirement in output.active_requirements:
        if requirement.is_external_tool_execution:
            result = _execute_tool(requirement)
            requirement.tool_execution.result = result
    
    result = await agent.continue_run_async(
        run_id=output.run_id, 
        return_output=True,
        hitl_handler=hitl_handler
    )
    
    return result


# ============================================================================
# VARIANT 5A: Cross-process External Tool Handling with task
# ============================================================================

async def external_cross_process_handling_task():
    """
    Pattern for handling external tools across process restarts using task parameter.
    
    Demonstrates the flow:
    1. Process A: print_do_async returns with paused status and requirements
    2. External system executes tools and sets results on in-memory requirements
    3. Process B: Creates new agent, uses task object for continuation
    
    Note: This uses task parameter, so it relies on in-memory context.
    For true cross-process persistence, use run_id variant instead.
    """
    cleanup_db()
    
    # STEP 1: Initial run pauses for external tool
    db = SqliteDatabase(db_file="external.db", session_id="session_1", user_id="user_1")
    agent = Agent("openai/gpt-4o-mini", name="external_tool_agent", db=db)
    task = Task(
        description="Send an email to test@example.com with subject 'Test' and body 'Hello'.",
        tools=[send_email]
    )
    
    output = await agent.print_do_async(task, return_output=True)
    run_id = output.run_id
    
    if output.is_paused and output.active_requirements:
        print(f"Run {run_id} paused for external tools:")
        for req in output.active_requirements:
            if req.tool_execution:
                print(f"  - Tool: {req.tool_execution.tool_name}")
                print(f"    Call ID: {req.tool_execution.tool_call_id}")
                print(f"    Args: {req.tool_execution.tool_args}")
    
    # STEP 2: Execute tools and set results on in-memory requirements
    print("\nExecuting external tools and setting results...")
    for req in output.active_requirements:
        if req.is_external_tool_execution and not req.is_resolved:
            tool_result = _execute_tool(req)
            req.tool_execution.result = tool_result
            print(f"  Set result for {req.tool_execution.tool_name}: {tool_result}")
    
    # STEP 3: New agent instance uses task for continuation (simulates Process B)
    # Pass requirements (with results) so new agent can inject them into loaded state
    print(f"\nCreating new agent to resume run {run_id} with task...")
    new_db = SqliteDatabase(db_file="external.db", session_id="session_1", user_id="user_1")
    new_agent = Agent("openai/gpt-4o-mini", name="external_tool_agent", db=new_db)
    
    result = await new_agent.continue_run_async(
        task=task, 
        requirements=output.requirements,
        return_output=True
    )
    print(f"Final result: {result.output}")
    return result


# ============================================================================
# VARIANT 5B: Cross-process External Tool Handling with run_id
# ============================================================================

async def external_cross_process_handling():
    """
    Pattern for handling external tools across process restarts.
    
    Demonstrates the full flow:
    1. Process A: print_do_async returns with paused status and requirements
    2. External system executes tools and sets results on in-memory requirements
    3. Process B: Creates new agent, loads run by run_id, and continues with requirements
    """
    cleanup_db()
    
    # STEP 1: Initial run pauses for external tool
    db = SqliteDatabase(db_file="external.db", session_id="session_1", user_id="user_1")
    agent = Agent("openai/gpt-4o-mini", name="external_tool_agent", db=db)
    task = Task(
        description="Send an email to test@example.com with subject 'Test' and body 'Hello'.",
        tools=[send_email]
    )
    
    output = await agent.print_do_async(task, return_output=True)
    run_id = output.run_id
    
    if output.is_paused and output.active_requirements:
        print(f"Run {run_id} paused for external tools:")
        for req in output.active_requirements:
            if req.tool_execution:
                print(f"  - Tool: {req.tool_execution.tool_name}")
                print(f"    Call ID: {req.tool_execution.tool_call_id}")
                print(f"    Args: {req.tool_execution.tool_args}")
    
    # STEP 2: Execute tools and set results on in-memory requirements
    print("\nExecuting external tools and setting results...")
    for req in output.active_requirements:
        if req.is_external_tool_execution and not req.is_resolved:
            tool_result = _execute_tool(req)
            req.tool_execution.result = tool_result
            print(f"  Set result for {req.tool_execution.tool_name}: {tool_result}")
    
    # STEP 3: New agent instance resumes from storage (simulates Process B)
    # Pass requirements (with results) so new agent can inject them into loaded state
    print(f"\nCreating new agent to resume run {run_id}...")
    new_db = SqliteDatabase(db_file="external.db", session_id="session_1", user_id="user_1")
    new_agent = Agent("openai/gpt-4o-mini", name="external_tool_agent", db=new_db)
    
    result = await new_agent.continue_run_async(
        run_id=run_id,
        requirements=output.requirements,
        return_output=True
    )
    print(f"Final result: {result.output}")
    return result


# ============================================================================
# PYTEST TESTS
# ============================================================================

@pytest.mark.asyncio
async def test_external_direct_call_with_run_id_same_agent():
    """Test: Direct Call with run_id - Same Agent"""
    result = await external_direct_call_with_run_id_same_agent()
    assert result.is_complete, f"Expected complete, got {result.status}"


@pytest.mark.asyncio
async def test_external_direct_call_with_task_same_agent():
    """Test: Direct Call with task - Same Agent"""
    result = await external_direct_call_with_task_same_agent()
    assert result.is_complete, f"Expected complete, got {result.status}"


@pytest.mark.asyncio
async def test_external_direct_call_with_run_id_new_agent():
    """Test: Direct Call with run_id - New Agent (Cross-process)"""
    result = await external_direct_call_with_run_id_new_agent()
    assert result.is_complete, f"Expected complete, got {result.status}"


@pytest.mark.asyncio
async def test_external_direct_call_with_task_new_agent():
    """Test: Direct Call with task - New Agent (Cross-process)"""
    result = await external_direct_call_with_task_new_agent()
    assert result.is_complete, f"Expected complete, got {result.status}"


@pytest.mark.asyncio
async def test_external_with_executor_callback():
    """Test: Using hitl_handler parameter with run_id"""
    result = await external_with_executor_callback()
    assert result.is_complete, f"Expected complete, got {result.status}"


@pytest.mark.asyncio
async def test_external_with_executor_callback_task():
    """Test: Using hitl_handler parameter with task"""
    result = await external_with_executor_callback_task()
    assert result.is_complete, f"Expected complete, got {result.status}"


@pytest.mark.asyncio
async def test_external_multiple_tools_direct_call():
    """Test: Multiple External Tools - Direct Call with run_id"""
    result = await external_multiple_tools_direct_call()
    assert result.is_complete, f"Expected complete, got {result.status}"


@pytest.mark.asyncio
async def test_external_multiple_tools_direct_call_task():
    """Test: Multiple External Tools - Direct Call with task"""
    result = await external_multiple_tools_direct_call_task()
    assert result.is_complete, f"Expected complete, got {result.status}"


@pytest.mark.asyncio
async def test_external_multiple_tools_with_executor():
    """Test: Multiple External Tools - With Executor (run_id)"""
    result = await external_multiple_tools_with_executor()
    assert result.is_complete, f"Expected complete, got {result.status}"


@pytest.mark.asyncio
async def test_external_multiple_tools_with_executor_task():
    """Test: Multiple External Tools - With Executor (task)"""
    result = await external_multiple_tools_with_executor_task()
    assert result.is_complete, f"Expected complete, got {result.status}"


@pytest.mark.asyncio
async def test_external_cross_process_handling():
    """Test: Cross-process External Tool Handling (run_id)"""
    result = await external_cross_process_handling()
    assert result.is_complete, f"Expected complete, got {result.status}"


@pytest.mark.asyncio
async def test_external_cross_process_handling_task():
    """Test: Cross-process External Tool Handling (task)"""
    result = await external_cross_process_handling_task()
    assert result.is_complete, f"Expected complete, got {result.status}"


# ============================================================================
# TEST RUNNER (for manual execution)
# ============================================================================

async def run_all_tests():
    """Run all external tool execution test variants."""
    
    print("\n" + "="*80)
    print("TEST 1: Direct Call with run_id - Same Agent")
    print("="*80)
    result = await external_direct_call_with_run_id_same_agent()
    assert result.is_complete, f"TEST 1 FAILED: Expected complete, got {result.status}"
    print("TEST 1 PASSED")
    
    print("\n" + "="*80)
    print("TEST 2: Direct Call with task - Same Agent")
    print("="*80)
    result = await external_direct_call_with_task_same_agent()
    assert result.is_complete, f"TEST 2 FAILED: Expected complete, got {result.status}"
    print("TEST 2 PASSED")
    
    print("\n" + "="*80)
    print("TEST 3: Direct Call with run_id - New Agent (Cross-process)")
    print("="*80)
    result = await external_direct_call_with_run_id_new_agent()
    assert result.is_complete, f"TEST 3 FAILED: Expected complete, got {result.status}"
    print("TEST 3 PASSED")
    
    print("\n" + "="*80)
    print("TEST 3A: Direct Call with task - New Agent (Cross-process)")
    print("="*80)
    result = await external_direct_call_with_task_new_agent()
    assert result.is_complete, f"TEST 3A FAILED: Expected complete, got {result.status}"
    print("TEST 3A PASSED")
    
    print("\n" + "="*80)
    print("TEST 4: Using hitl_handler parameter with run_id")
    print("="*80)
    result = await external_with_executor_callback()
    assert result.is_complete, f"TEST 4 FAILED: Expected complete, got {result.status}"
    print("TEST 4 PASSED")
    
    print("\n" + "="*80)
    print("TEST 4A: Using hitl_handler parameter with task")
    print("="*80)
    result = await external_with_executor_callback_task()
    assert result.is_complete, f"TEST 4A FAILED: Expected complete, got {result.status}"
    print("TEST 4A PASSED")
    
    print("\n" + "="*80)
    print("TEST 5: Multiple External Tools - Direct Call with run_id")
    print("="*80)
    result = await external_multiple_tools_direct_call()
    assert result.is_complete, f"TEST 5 FAILED: Expected complete, got {result.status}"
    print("TEST 5 PASSED")
    
    print("\n" + "="*80)
    print("TEST 5A: Multiple External Tools - Direct Call with task")
    print("="*80)
    result = await external_multiple_tools_direct_call_task()
    assert result.is_complete, f"TEST 5A FAILED: Expected complete, got {result.status}"
    print("TEST 5A PASSED")
    
    print("\n" + "="*80)
    print("TEST 6: Multiple External Tools - With Executor (run_id)")
    print("="*80)
    result = await external_multiple_tools_with_executor()
    assert result.is_complete, f"TEST 6 FAILED: Expected complete, got {result.status}"
    print("TEST 6 PASSED")
    
    print("\n" + "="*80)
    print("TEST 6A: Multiple External Tools - With Executor (task)")
    print("="*80)
    result = await external_multiple_tools_with_executor_task()
    assert result.is_complete, f"TEST 6A FAILED: Expected complete, got {result.status}"
    print("TEST 6A PASSED")
    
    print("\n" + "="*80)
    print("TEST 7: Cross-process External Tool Handling (run_id)")
    print("="*80)
    result = await external_cross_process_handling()
    assert result.is_complete, f"TEST 7 FAILED: Expected complete, got {result.status}"
    print("TEST 7 PASSED")
    
    print("\n" + "="*80)
    print("TEST 7A: Cross-process External Tool Handling (task)")
    print("="*80)
    result = await external_cross_process_handling_task()
    assert result.is_complete, f"TEST 7A FAILED: Expected complete, got {result.status}"
    print("TEST 7A PASSED")
    
    # Cleanup
    cleanup_db()
    
    print("\n" + "="*80)
    print("ALL TESTS PASSED!")
    print("="*80)


if __name__ == "__main__":
    asyncio.run(run_all_tests())
