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

# 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_flow")

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

    async def listen(self):
        url = f"{WS_URL}/ws/{TENANT_ID}/tickets"
        logger.info(f"Connecting to WebSocket: {url}")
        try:
            async with websockets.connect(url) as websocket:
                logger.info("Connected to WebSocket.")
                while not self.stop_event.is_set():
                    try:
                        # Use a short timeout to check stop_event periodically
                        message = await asyncio.wait_for(websocket.recv(), timeout=1.0)
                        data = json.loads(message)
                        logger.info(f" [WS RECEIVED] Event: {data.get('event')} - Ticket: {data.get('data', {}).get('ticket_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 = TicketTracker()
    listen_task = asyncio.create_task(tracker.listen())
    
    # Wait for WS to be ready
    await asyncio.sleep(3)
    
    thread_id = f"verify-{uuid.uuid4().hex[:8]}"
    logger.info(f"Starting chat flow for thread: {thread_id}")
    
    async with httpx.AsyncClient(timeout=60.0) as client:
        # We'll use a single message that provides all details to trigger ticket immediately
        query = (
            f"Please create a support ticket. "
            f"Issue: WebSocket Verification Test. "
            f"Priority: High. "
            f"Email: test@galaxiq.ai. Phone: 9999999999. "
            f"Thread ID: {thread_id}"
        )
        
        logger.info(f"Sending Chat request...")
        response = await client.post(
            f"{BASE_URL}/chat",
            data={
                "tenant_id": TENANT_ID,
                "query": query,
                "thread_id": thread_id,
                "user_name": "VerifierBot",
                "role": "user"
            }
        )
        
        if response.status_code != 200:
            logger.error(f"Chat failed: {response.text}")
            tracker.stop_event.set()
            return

        chat_data = response.json()
        ai_response = chat_data.get("response", "")
        logger.info(f"AI Response: {ai_response}")
        
        # Check if ticket ID is in AI response
        import re
        ticket_match = re.search(r"TICK-[A-Z0-9]+", ai_response)
        if ticket_match:
            ticket_id = ticket_match.group(0)
            logger.info(f"Confirmed ticket created by AI: {ticket_id}")
        else:
            logger.warning("No ticket ID found in AI response. Ticket might not have been created.")
            ticket_id = None

    # 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("FINAL TEST RESULTS")
    print("="*50)
    
    if ticket_id:
        # Look for this ticket ID in received events
        found = False
        for event in tracker.received_events:
            if event.get("data", {}).get("ticket_id") == ticket_id:
                found = True
                break
        
        if found:
            print(f"RESULT: SUCCESS - Event for {ticket_id} received via WebSocket!")
        else:
            print(f"RESULT: FAILURE - Ticket {ticket_id} created but NO WebSocket event received.")
    else:
        print("RESULT: SKIPPED - Ticket was not created via Chat.")
    
    print("="*50)

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