import asyncio
import httpx
import uuid
import json
import logging
import websockets

# Configuration
BASE_URL = "https://dev-strategistai.galaxiq.ai"
WS_URL = "wss://dev-strategistai.galaxiq.ai"
TENANT_ID = "org_d1e1c1ea-6c90-431f-9337-3375dcdfef5c"

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger("verify_conversation")

class ConversationTracker:
    def __init__(self):
        self.received_events = []
        self.stop_event = asyncio.Event()

    async def listen(self):
        url = f"{WS_URL}/ws/{TENANT_ID}/conversation_list"
        logger.info(f"Connecting to WebSocket: {url}")
        try:
            async with websockets.connect(url) as websocket:
                logger.info("Connected to Conversation List WebSocket.")
                while not self.stop_event.is_set():
                    try:
                        message = await asyncio.wait_for(websocket.recv(), timeout=1.0)
                        data = json.loads(message)
                        event_name = data.get('event')
                        # Handle both flat and nested structures
                        thread_id = data.get('thread_id') or data.get('data', {}).get('thread_id')
                        logger.info(f" [WS RECEIVED] Event: {event_name} - Thread: {thread_id}")
                        self.received_events.append(data)
                    except asyncio.TimeoutError:
                        continue
                    except Exception as e:
                        logger.error(f"WebSocket receive error: {e}")
                        break
        except Exception as e:
            logger.error(f"WebSocket connection error: {e}")

async def run_test():
    tracker = ConversationTracker()
    listen_task = asyncio.create_task(tracker.listen())
    
    # Wait for WS to be ready
    await asyncio.sleep(3)
    
    # Create a completely new thread ID
    thread_id = f"thread-{uuid.uuid4().hex[:12]}"
    logger.info(f"Starting new conversation with thread: {thread_id}")
    
    async with httpx.AsyncClient(timeout=60.0) as client:
        logger.info(f"Sending first chat message to start the conversation...")
        response = await client.post(
            f"{BASE_URL}/chat",
            data={
                "tenant_id": TENANT_ID,
                "query": "Hello! I am starting a new conversation to test the list update.",
                "thread_id": thread_id,
                "user_name": "ListTester",
                "role": "user"
            }
        )
        
        if response.status_code != 200:
            logger.error(f"Chat failed: {response.text}")
            tracker.stop_event.set()
            return

        logger.info(f"Chat response received. AI: {response.json().get('response')[:50]}...")

    # Wait for WS event
    logger.info("Waiting 10 seconds for WebSocket event...")
    await asyncio.sleep(10)
    
    # Analysis
    tracker.stop_event.set()
    await listen_task
    
    print("\n" + "="*50)
    print("CONVERSATION LIST TEST RESULTS")
    print("="*50)
    
    # Look for this thread ID in received events
    found = False
    for event in tracker.received_events:
        received_tid = event.get('thread_id') or event.get('data', {}).get('thread_id')
        if received_tid == thread_id:
            found = True
            break
    
    if found:
        print(f"RESULT: SUCCESS - Event for thread {thread_id} received via WebSocket!")
    else:
        print(f"RESULT: FAILURE - Thread {thread_id} started but NO WebSocket event received.")
    
    print("="*50)

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