
import sys
import os
from unittest.mock import MagicMock

# Mock LangChain messages
class AIMessage:
    def __init__(self, content):
        self.content = content

class HumanMessage:
    def __init__(self, content):
        self.content = content

# Set the path so it can import from services
sys.path.append(os.getcwd())

def test_extraction_logic():
    # Mock the response structure from agent_executor.invoke
    mock_response_1 = {
        "messages": [
            HumanMessage("Hello"),
            AIMessage([{"text": "Block 1 "}, {"text": "Block 2"}])
        ]
    }
    
    mock_response_2 = {
        "messages": [
            HumanMessage("Hello"),
            AIMessage("Simple string response")
        ]
    }
    
    mock_response_empty = {
        "messages": [
            HumanMessage("Hello"),
            AIMessage("")
        ]
    }
    
    # Simulation function based on the logic in chat_agent.py
    def extract(response):
        answer_str = ""
        for msg in reversed(response["messages"]):
            if isinstance(msg, AIMessage):
                content = msg.content
                if isinstance(content, str) and content.strip():
                    answer_str = content.strip()
                    break
                elif isinstance(content, list):
                    parts = []
                    for block in content:
                        if isinstance(block, dict):
                            parts.append(block.get("text", ""))
                        elif isinstance(block, str):
                            parts.append(block)
                    joined = "".join(parts).strip()
                    if joined:
                        answer_str = joined
                        break
        return answer_str

    print("Test 1 (List of dicts):", extract(mock_response_1))
    assert extract(mock_response_1) == "Block 1 Block 2"
    
    print("Test 2 (String):", extract(mock_response_2))
    assert extract(mock_response_2) == "Simple string response"
    
    print("Test 3 (Empty):", extract(mock_response_empty))
    assert extract(mock_response_empty) == ""

if __name__ == "__main__":
    try:
        test_extraction_logic()
        print("\nExtraction logic verification: SUCCESS")
    except AssertionError as e:
        print("\nExtraction logic verification: FAILED")
        sys.exit(1)
