#!/usr/bin/env python3
"""
Agent Config Test Utility - Run agents with dynamic JSON configs
"""
import asyncio
import json
import importlib
from typing import Dict, Any, Optional, List
# Lazy import to avoid circular dependency - imported inside functions
from ..unified_chat import UnifiedChat, ProviderEnum
from ..memory.history import History
from ..baseheaventool import ToolResult, ToolError
from .find_tool_use_in_history import (
    find_tool_use_in_history, 
    count_tool_calls_in_history
)

async def agent_config_test(
    test_prompt: str,
    system_prompt: str,
    iterations: int = 1,
    agent_mode: bool = True,
    name: str = "DynamicTestAgent",
    tools: Optional[List[str]] = None,
    provider: str = "openai",
    model: Optional[str] = None,
    temperature: float = 0.7,
    max_tokens: int = 8000,
    thinking_budget: Optional[int] = None,
    additional_kws: Optional[List[str]] = None,
    additional_kw_instructions: str = "",
    known_config_paths: Optional[List[str]] = None,
    prompt_suffix_blocks: Optional[List[str]] = None,
    # BaseHeavenAgent initialization parameters
    max_tool_calls: int = 10,
    orchestrator: bool = False,
    history_id: Optional[str] = None,
    system_prompt_suffix: Optional[str] = None,
    adk: bool = False,
    duo_enabled: bool = False,
    run_on_langchain: bool = False,
    # Assertion parameters for validation
    assert_tool_used: Optional[str] = None,
    assert_no_errors: bool = False,
    assert_goal_accomplished: bool = False,
    assert_extracted_keys: Optional[List[str]] = None,
    assert_extracted_contains: Optional[Dict[str, Any]] = None,
    assert_min_tool_calls: Optional[int] = None,
    assert_output_contains: Optional[str] = None
) -> Dict[str, Any]:
    """
    Test an agent configuration by running it with a test prompt.
    
    Returns:
        Dictionary with test results including final output, execution info, and assertion results
    """
    try:
        # Lazy import to avoid circular dependency
        from ..baseheavenagent import BaseHeavenAgent, HeavenAgentConfig
        
        # Convert tool names to classes if provided (using OmniTool's approach)
        tool_classes = []
        if tools:
            from computer_use_demo.tools.base.utils.agent_and_tool_lists import get_tool_modules
            
            # Get available tool class names (PascalCase)
            available = get_tool_modules().split(", ")
            
            for tool_name in tools:
                try:
                    # Normalize to class name (same logic as OmniTool)
                    if tool_name not in available:
                        alt = ''.join(part.capitalize() for part in tool_name.split('_'))
                        if alt in available:
                            tool_name = alt
                        else:
                            print(f"Warning: Tool '{tool_name}' not found among available tools")
                            continue
                    
                    # Derive module name from class name (same logic as OmniTool)
                    module_name = ''.join(
                        ('_' + c.lower() if c.isupper() else c) for c in tool_name
                    ).lstrip('_')
                    full_module = f'computer_use_demo.tools.base.tools.{module_name}'
                    
                    # Import the tool class
                    module = importlib.import_module(full_module)
                    tool_class = getattr(module, tool_name)
                    tool_classes.append(tool_class)
                    
                except Exception as e:
                    print(f"Warning: Could not import tool {tool_name}: {e}")
        
        # Convert provider string to enum if needed
        if isinstance(provider, str):
            provider_map = {
                'anthropic': ProviderEnum.ANTHROPIC,
                'openai': ProviderEnum.OPENAI,
                'google': ProviderEnum.GOOGLE
            }
            provider_enum = provider_map.get(provider.lower(), ProviderEnum.ANTHROPIC)
        else:
            provider_enum = provider
        
        # Create HeavenAgentConfig from parameters
        agent_config = HeavenAgentConfig(
            name=name,
            system_prompt=system_prompt,
            tools=tool_classes,
            provider=provider_enum,
            model=model,
            temperature=temperature,
            max_tokens=max_tokens,
            thinking_budget=thinking_budget,
            additional_kws=additional_kws or [],
            additional_kw_instructions=additional_kw_instructions,
            known_config_paths=known_config_paths,
            prompt_suffix_blocks=prompt_suffix_blocks
        )
        
        # Create agent instance with all parameters
        agent = BaseHeavenAgent(
            config=agent_config,
            unified_chat=UnifiedChat(),
            max_tool_calls=max_tool_calls,
            orchestrator=orchestrator,
            history=History(messages=[]) if not history_id else None,
            history_id=history_id,
            system_prompt_suffix=system_prompt_suffix,
            adk=adk,
            duo_enabled=duo_enabled,
            run_on_langchain=run_on_langchain
        )
        
        # Format prompt based on mode
        if agent_mode:
            formatted_prompt = f"agent goal={test_prompt}, iterations={iterations}"
        else:
            formatted_prompt = test_prompt
        
        # Run the agent
        print("DEBUG: About to call agent.run()")
        result = await agent.run(prompt=formatted_prompt)
        print("DEBUG: agent.run() completed successfully")
        
        print(f"DEBUG: agent.run() returned type: {type(result)}")
        print(f"DEBUG: result is dict: {isinstance(result, dict)}")
        if isinstance(result, dict):
            print(f"DEBUG: result keys: {list(result.keys())}")
            print(f"DEBUG: 'history' in result: {'history' in result}")
        else:
            print(f"DEBUG: result value: {result}")
        
        # Extract key information
        final_output = ""
        message_count = 0
        tool_calls = 0
        
        if isinstance(result, dict) and "history" in result:
            history = result["history"]
            message_count = len(history.messages)
            print(f"DEBUG: About to call count_tool_calls_in_history with {len(history.messages)} messages")
            tool_calls = count_tool_calls_in_history(history)
            print(f"DEBUG: count_tool_calls_in_history returned: {tool_calls}")
            
            # # Get final output from the last AI message
            # for msg in reversed(history.messages):
            #     if hasattr(msg, 'content'):
            #         if isinstance(msg.content, list):
            #             # Get text blocks for final output
            #             text_blocks = [block.get('text', '') for block in msg.content if isinstance(block, dict) and block.get('type') == 'text']
            #             if text_blocks:
            #                 final_output = '\n'.join(text_blocks)
            #                 break
            #         elif isinstance(msg.content, str):
            #             if msg.__class__.__name__ == 'AIMessage':
            #                 final_output = msg.content
            #                 break
        
        # Get extracted content if available
        extracted_content = {}
        if (isinstance(result, dict) and 
            result.get('agent_status') and 
            result['agent_status'].extracted_content):
            extracted_content = result['agent_status'].extracted_content
        
        # Run assertions if provided
        assertion_results = {}
        all_assertions_passed = True
        
        print("DEBUG: Starting assertion processing")
        
        if assert_tool_used:
            print(f"DEBUG: assert_tool_used = {assert_tool_used}")
            tool_was_used = False
            if isinstance(result, dict) and "history" in result:
                print("DEBUG: About to call find_tool_use_in_history")
                tool_was_used = find_tool_use_in_history(result["history"], assert_tool_used)
                print(f"DEBUG: find_tool_use_in_history returned: {tool_was_used}")
            else:
                print("DEBUG: Skipping find_tool_use_in_history - no history in result")
            assertion_results["tool_used"] = {
                "passed": tool_was_used,
                "expected": assert_tool_used,
                "message": f"Tool '{assert_tool_used}' {'was' if tool_was_used else 'was NOT'} used"
            }
            if not tool_was_used:
                all_assertions_passed = False

        if assert_min_tool_calls is not None:
            enough_calls = tool_calls >= assert_min_tool_calls
            assertion_results["min_tool_calls"] = {
                "passed": enough_calls,
                "expected": f">= {assert_min_tool_calls}",
                "actual": tool_calls,
                "message": f"Made {tool_calls} tool calls, {'enough' if enough_calls else 'not enough'}"
            }
            if not enough_calls:
                all_assertions_passed = False
        
        return {
            "success": True,
            "final_output": history.messages,
            "message_count": message_count,
            "tool_calls": tool_calls,
            "history_id": result.get("history_id") if isinstance(result, dict) else None,
            "agent_status": result.get("agent_status") if isinstance(result, dict) else None,
            "extracted_content": extracted_content,
            "config_used": {
                "name": agent_config.name,
                "provider": agent_config.provider.value,
                "model": agent_config.model,
                "temperature": agent_config.temperature,
                "max_tokens": agent_config.max_tokens,
                "tools": [tool.__name__ for tool in tool_classes],
                "max_tool_calls": max_tool_calls,
                "orchestrator": orchestrator,
                "adk": adk,
                "duo_enabled": duo_enabled,
                "run_on_langchain": run_on_langchain
            },
            "assertions": {
                "all_passed": all_assertions_passed,
                "results": assertion_results,
                "count": len(assertion_results)
            }
        }
        
    except Exception as e:
        print(f"DEBUG: Exception occurred: {e}")
        print(f"DEBUG: Exception type: {type(e)}")
        return {
            "success": False,
            "error": str(e),
            "error_type": type(e).__name__,
            "final_output": "",
            "message_count": 0,
            "tool_calls": 0,
            "extracted_content": {},
            "assertions": {
                "all_passed": False,
                "results": {},
                "count": 0
            }
        }




# BROKEN
# #!/usr/bin/env python3
# """
# Agent Config Test Utility - Run agents with dynamic JSON configs
# """
# import asyncio
# import json
# import importlib
# from typing import Dict, Any, Optional, List
# from computer_use_demo.tools.base.baseheavenagent import BaseHeavenAgent, HeavenAgentConfig
# from computer_use_demo.tools.base.unified_chat import UnifiedChat, ProviderEnum
# from computer_use_demo.tools.base.memory.history import History
# from computer_use_demo.tools.base.baseheaventool import ToolResult, ToolError
# from computer_use_demo.tools.base.tool_utils.find_tool_use_in_history import (
#     find_tool_use_in_history, 
#     count_tool_calls_in_history
# )

# async def agent_config_test(
#     test_prompt: str,
#     system_prompt: str,
#     iterations: int = 1,
#     agent_mode: bool = True,
#     name: str = "DynamicTestAgent",
#     tools: Optional[List[str]] = None,
#     provider: str = "openai",
#     model: Optional[str] = None,
#     temperature: float = 0.7,
#     max_tokens: int = 8000,
#     thinking_budget: Optional[int] = None,
#     additional_kws: Optional[List[str]] = None,
#     additional_kw_instructions: str = "",
#     known_config_paths: Optional[List[str]] = None,
#     prompt_suffix_blocks: Optional[List[str]] = None,
#     # BaseHeavenAgent initialization parameters
#     max_tool_calls: int = 10,
#     orchestrator: bool = False,
#     history_id: Optional[str] = None,
#     system_prompt_suffix: Optional[str] = None,
#     adk: bool = False,
#     duo_enabled: bool = False,
#     run_on_langchain: bool = False,
#     # Assertion parameters for validation
#     assert_tool_used: Optional[str] = None,
#     assert_no_errors: bool = False,
#     assert_goal_accomplished: bool = False,
#     assert_extracted_keys: Optional[List[str]] = None,
#     assert_extracted_contains: Optional[Dict[str, Any]] = None,
#     assert_min_tool_calls: Optional[int] = None,
#     assert_output_contains: Optional[str] = None
# ) -> Dict[str, Any]:
#     """
#     Test an agent configuration by running it with a test prompt.
    
#     Args:
#         test_prompt: The prompt to test the agent with
#         system_prompt: System prompt for the agent
#         iterations: Number of iterations (only used in agent mode)
#         agent_mode: Whether to run in agent mode or direct mode
#         name: Name for the agent configuration
#         tools: List of tool names to include
#         provider: AI provider (anthropic, openai, google)
#         model: Model name
#         temperature: Temperature for AI generation
#         max_tokens: Maximum tokens for AI response
#         thinking_budget: Thinking budget for reasoning
#         additional_kws: Additional keywords for extraction
#         additional_kw_instructions: Instructions for additional keyword extraction
#         known_config_paths: List of Hermes config paths for orchestrator mode
#         prompt_suffix_blocks: List of prompt suffix block names to append
#         max_tool_calls: Maximum number of tool calls allowed
#         orchestrator: Whether to run in orchestrator mode with Hermes Switchboard
#         history_id: Existing history ID to continue from
#         system_prompt_suffix: Additional text to append to the system prompt
#         adk: Whether to use Google ADK instead of LangChain
#         duo_enabled: Whether to enable DUO prompt injection system
#         run_on_langchain: Force LangChain usage even with Google provider
#         assert_tool_used: Assert that this specific tool was used
#         assert_no_errors: Assert that no tool errors occurred during execution
#         assert_goal_accomplished: Assert that the agent marked the goal as accomplished
#         assert_extracted_keys: Assert that these keys exist in extracted_content
#         assert_extracted_contains: Assert that extracted_content contains specific key-value pairs
#         assert_min_tool_calls: Assert minimum number of tool calls made
#         assert_output_contains: Assert that final output contains this substring
        
#     Returns:
#         Dictionary with test results including final output, execution info, and assertion results
#     """
#     try:
#         # Convert tool names to classes if provided (using OmniTool's approach)
#         tool_classes = []
#         if tools:
#             from computer_use_demo.tools.base.utils.agent_and_tool_lists import get_tool_modules
            
#             # Get available tool class names (PascalCase)
#             available = get_tool_modules().split(", ")
            
#             for tool_name in tools:
#                 try:
#                     # Normalize to class name (same logic as OmniTool)
#                     if tool_name not in available:
#                         alt = ''.join(part.capitalize() for part in tool_name.split('_'))
#                         if alt in available:
#                             tool_name = alt
#                         else:
#                             print(f"Warning: Tool '{tool_name}' not found among available tools")
#                             continue
                    
#                     # Derive module name from class name (same logic as OmniTool)
#                     module_name = ''.join(
#                         ('_' + c.lower() if c.isupper() else c) for c in tool_name
#                     ).lstrip('_')
#                     full_module = f'computer_use_demo.tools.base.tools.{module_name}'
                    
#                     # Import the tool class
#                     module = importlib.import_module(full_module)
#                     tool_class = getattr(module, tool_name)
#                     tool_classes.append(tool_class)
                    
#                 except Exception as e:
#                     print(f"Warning: Could not import tool {tool_name}: {e}")
        
#         # Convert provider string to enum if needed
#         if isinstance(provider, str):
#             provider_map = {
#                 'anthropic': ProviderEnum.ANTHROPIC,
#                 'openai': ProviderEnum.OPENAI,
#                 'google': ProviderEnum.GOOGLE
#             }
#             provider_enum = provider_map.get(provider.lower(), ProviderEnum.ANTHROPIC)
#         else:
#             provider_enum = provider
        
#         # Create HeavenAgentConfig from parameters
#         agent_config = HeavenAgentConfig(
#             name=name,
#             system_prompt=system_prompt,
#             tools=tool_classes,
#             provider=provider_enum,
#             model=model,
#             temperature=temperature,
#             max_tokens=max_tokens,
#             thinking_budget=thinking_budget,
#             additional_kws=additional_kws or [],
#             additional_kw_instructions=additional_kw_instructions,
#             known_config_paths=known_config_paths,
#             prompt_suffix_blocks=prompt_suffix_blocks
#         )
        
#         # Create agent instance with all parameters
#         agent = BaseHeavenAgent(
#             config=agent_config,
#             unified_chat=UnifiedChat(),
#             max_tool_calls=max_tool_calls,
#             orchestrator=orchestrator,
#             history=History(messages=[]) if not history_id else None,
#             history_id=history_id,
#             system_prompt_suffix=system_prompt_suffix,
#             adk=adk,
#             duo_enabled=duo_enabled,
#             run_on_langchain=run_on_langchain
#         )
        
#         # Format prompt based on mode
#         if agent_mode:
#             formatted_prompt = f"agent goal={test_prompt}, iterations={iterations}"
#         else:
#             formatted_prompt = test_prompt
        
#         # Run the agent
#         print("DEBUG: About to call agent.run()")
#         try:
#             result = await agent.run(prompt=formatted_prompt)
#             print("DEBUG: agent.run() completed successfully")
#         except Exception as e:
#             print(f"DEBUG: agent.run() failed with exception: {e}")
#             print(f"DEBUG: Exception type: {type(e)}")
#             raise
        
#         print(f"DEBUG: agent.run() returned type: {type(result)}")
#         print(f"DEBUG: result is dict: {isinstance(result, dict)}")
#         if isinstance(result, dict):
#             print(f"DEBUG: result keys: {list(result.keys())}")
#             print(f"DEBUG: 'history' in result: {'history' in result}")
#         else:
#             print(f"DEBUG: result value: {result}")
        
#         # Extract key information
#         final_output = ""
#         message_count = 0
#         tool_calls = 0
        
#         if isinstance(result, dict) and "history" in result:
#             history = result["history"]
#             message_count = len(history.messages)
#             tool_calls = count_tool_calls_in_history(history)
            
#             # Get final output from the last AI message
#             for msg in reversed(history.messages):
#                 if hasattr(msg, 'content'):
#                     if isinstance(msg.content, list):
#                         # Get text blocks for final output
#                         text_blocks = [block.get('text', '') for block in msg.content if isinstance(block, dict) and block.get('type') == 'text']
#                         if text_blocks:
#                             final_output = '\n'.join(text_blocks)
#                             break
#                     elif isinstance(msg.content, str):
#                         if msg.__class__.__name__ == 'AIMessage':
#                             final_output = msg.content
#                             break
        
#         # Get extracted content if available
#         extracted_content = {}
#         if (isinstance(result, dict) and 
#             result.get('agent_status') and 
#             result['agent_status'].extracted_content):
#             extracted_content = result['agent_status'].extracted_content
        
#         # Run assertions if provided
#         assertion_results = {}
#         all_assertions_passed = True
        
#         print("DEBUG: Starting assertion processing")
#         try:
#             if assert_tool_used:
#                 print(f"DEBUG: assert_tool_used = {assert_tool_used}")
#                 tool_was_used = False
#                 if isinstance(result, dict) and "history" in result:
#                     print("DEBUG: About to call find_tool_use_in_history")
#                     tool_was_used = find_tool_use_in_history(result["history"], assert_tool_used)
#                     print(f"DEBUG: find_tool_use_in_history returned: {tool_was_used}")
#                 else:
#                     print("DEBUG: Skipping find_tool_use_in_history - no history in result")
#                 assertion_results["tool_used"] = {
#                     "passed": tool_was_used,
#                     "expected": assert_tool_used,
#                     "message": f"Tool '{assert_tool_used}' {'was' if tool_was_used else 'was NOT'} used"
#                 }
#                 if not tool_was_used:
#                     all_assertions_passed = False
        
#             if assert_no_errors:
#                 has_errors = False
#             error_details = []
#             if isinstance(result, dict) and "history" in result:
#                 for msg in result["history"].messages:
#                     if (hasattr(msg, 'content') and 
#                         isinstance(msg.content, str) and 
#                         'error' in msg.content.lower()):
#                         has_errors = True
#                         error_details.append(msg.content[:100] + "..." if len(msg.content) > 100 else msg.content)
#             assertion_results["no_errors"] = {
#                 "passed": not has_errors,
#                 "expected": "no errors",
#                 "message": f"{'No errors found' if not has_errors else f'Errors found: {error_details}'}"
#             }
#             if has_errors:
#                 all_assertions_passed = False
        
#         if assert_goal_accomplished:
#             goal_accomplished = (isinstance(result, dict) and 
#                                result.get('agent_status') and 
#                                result['agent_status'].completed)
#             assertion_results["goal_accomplished"] = {
#                 "passed": goal_accomplished,
#                 "expected": True,
#                 "message": f"Goal {'was' if goal_accomplished else 'was NOT'} marked as accomplished"
#             }
#             if not goal_accomplished:
#                 all_assertions_passed = False
        
#         if assert_extracted_keys:
#             missing_keys = [key for key in assert_extracted_keys if key not in extracted_content]
#             assertion_results["extracted_keys"] = {
#                 "passed": len(missing_keys) == 0,
#                 "expected": assert_extracted_keys,
#                 "actual": list(extracted_content.keys()),
#                 "message": f"Missing keys: {missing_keys}" if missing_keys else "All expected keys found"
#             }
#             if missing_keys:
#                 all_assertions_passed = False
        
#         if assert_extracted_contains:
#             mismatches = []
#             for key, expected_value in assert_extracted_contains.items():
#                 if key not in extracted_content:
#                     mismatches.append(f"Key '{key}' not found")
#                 elif extracted_content[key] != expected_value:
#                     mismatches.append(f"Key '{key}': expected {expected_value}, got {extracted_content[key]}")
#             assertion_results["extracted_contains"] = {
#                 "passed": len(mismatches) == 0,
#                 "expected": assert_extracted_contains,
#                 "actual": extracted_content,
#                 "message": f"Mismatches: {mismatches}" if mismatches else "All key-value pairs match"
#             }
#             if mismatches:
#                 all_assertions_passed = False
        
#         if assert_min_tool_calls is not None:
#             enough_calls = tool_calls >= assert_min_tool_calls
#             assertion_results["min_tool_calls"] = {
#                 "passed": enough_calls,
#                 "expected": f">= {assert_min_tool_calls}",
#                 "actual": tool_calls,
#                 "message": f"Made {tool_calls} tool calls, {'enough' if enough_calls else 'not enough'}"
#             }
#             if not enough_calls:
#                 all_assertions_passed = False
        
#         if assert_output_contains:
#             contains_text = assert_output_contains in final_output
#             assertion_results["output_contains"] = {
#                 "passed": contains_text,
#                 "expected": assert_output_contains,
#                 "actual": final_output[:200] + "..." if len(final_output) > 200 else final_output,
#                 "message": f"Output {'contains' if contains_text else 'does NOT contain'} expected text"
#             }
#             if not contains_text:
#                 all_assertions_passed = False
        
#         except Exception as e:
#             print(f"DEBUG: Exception in assertion processing: {e}")
#             print(f"DEBUG: Exception type: {type(e)}")
#             raise
        
#         return {
#             "success": True,
#             "final_output": final_output,
#             "message_count": message_count,
#             "tool_calls": tool_calls,
#             "history_id": result.get("history_id") if isinstance(result, dict) else None,
#             "agent_status": result.get("agent_status") if isinstance(result, dict) else None,
#             "extracted_content": extracted_content,
#             "config_used": {
#                 "name": agent_config.name,
#                 "provider": agent_config.provider.value,
#                 "model": agent_config.model,
#                 "temperature": agent_config.temperature,
#                 "max_tokens": agent_config.max_tokens,
#                 "tools": [tool.__name__ for tool in tool_classes],
#                 "max_tool_calls": max_tool_calls,
#                 "orchestrator": orchestrator,
#                 "adk": adk,
#                 "duo_enabled": duo_enabled,
#                 "run_on_langchain": run_on_langchain
#             },
#             "assertions": {
#                 "all_passed": all_assertions_passed,
#                 "results": assertion_results,
#                 "count": len(assertion_results)
#             }
#         }
        
#     except Exception as e:
#         return {
#             "success": False,
#             "error": str(e),
#             "error_type": type(e).__name__,
#             "final_output": "",
#             "message_count": 0,
#             "tool_calls": 0,
#             "extracted_content": {},
#             "assertions": {
#                 "all_passed": False,
#                 "results": {},
#                 "count": 0
#             }
#         }

# if __name__ == "__main__":
#     # Example usage
#     async def main():
#         result = await agent_config_test(
#             test_prompt="Analyze the file /tmp/test.py",
#             system_prompt="You are a helpful AI assistant that can analyze code.",
#             name="TestAgent",
#             tools=["SafeCodeReaderTool"],
#             provider="anthropic",
#             model="claude-3-5-sonnet-20241022",
#             temperature=0.3,
#             iterations=1
#         )
        
#         print(json.dumps(result, indent=2))
    
#     asyncio.run(main())