#!/usr/bin/env python3
"""
Test client for hwGpt WebSocket chatbot.
This script demonstrates how to connect and interact with the chatbot.
"""

import asyncio
import websockets
import json
import logging
from datetime import datetime

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class ChatbotTestClient:
    """Test client for the hwGpt WebSocket chatbot."""
    
    def __init__(self, websocket_url: str, user_id: str, store_id: str, branch_id: str = None):
        self.websocket_url = websocket_url
        self.user_id = user_id
        self.store_id = store_id
        self.branch_id = branch_id
        self.websocket = None
        self.thread_id = None
    
    async def connect(self):
        """Connect to the WebSocket chatbot."""
        try:
            if self.branch_id:
                url = f"{self.websocket_url}/chat/{self.user_id}/{self.store_id}/{self.branch_id}"
            else:
                url = f"{self.websocket_url}/chat/{self.user_id}/{self.store_id}"
            
            logger.info(f"Connecting to: {url}")
            self.websocket = await websockets.connect(url)
            
            # Wait for connection message
            response = await self.websocket.recv()
            data = json.loads(response)
            
            if data.get("type") == "connection_established":
                logger.info("✅ Connected to chatbot successfully!")
                logger.info(f"Connection details: {data}")
                return True
            else:
                logger.error(f"Unexpected connection response: {data}")
                return False
                
        except Exception as e:
            logger.error(f"Failed to connect: {str(e)}")
            return False
    
    async def send_message(self, message: str, thread_id: int = None):
        """Send a chat message to the chatbot."""
        if not self.websocket:
            logger.error("Not connected to chatbot")
            return
        
        try:
            # Prepare message data
            message_data = {
                "type": "chat",
                "message": message
            }
            
            if thread_id:
                message_data["thread_id"] = thread_id
                self.thread_id = thread_id
            
            # Send message
            await self.websocket.send(json.dumps(message_data))
            logger.info(f"📤 Sent message: {message}")
            
            # Wait for responses
            await self.wait_for_responses()
            
        except Exception as e:
            logger.error(f"Error sending message: {str(e)}")
    
    async def create_thread(self, title: str, description: str = "", thread_type: str = "chat"):
        """Create a new chat thread."""
        if not self.websocket:
            logger.error("Not connected to chatbot")
            return
        
        try:
            thread_data = {
                "type": "thread_operation",
                "operation": "create_thread",
                "data": {
                    "title": title,
                    "description": description,
                    "thread_type": thread_type
                }
            }
            
            await self.websocket.send(json.dumps(thread_data))
            logger.info(f"📤 Creating thread: {title}")
            
            # Wait for thread creation response
            await self.wait_for_responses()
            
        except Exception as e:
            logger.error(f"Error creating thread: {str(e)}")
    
    async def list_threads(self, page: int = 1, page_size: int = 20):
        """List user threads."""
        if not self.websocket:
            logger.error("Not connected to chatbot")
            return
        
        try:
            thread_data = {
                "type": "thread_operation",
                "operation": "list_threads",
                "data": {
                    "page": page,
                    "page_size": page_size
                }
            }
            
            await self.websocket.send(json.dumps(thread_data))
            logger.info("📤 Requesting thread list")
            
            # Wait for thread list response
            await self.wait_for_responses()
            
        except Exception as e:
            logger.error(f"Error listing threads: {str(e)}")
    
    async def switch_thread(self, thread_id: int):
        """Switch to a different thread."""
        if not self.websocket:
            logger.error("Not connected to chatbot")
            return
        
        try:
            thread_data = {
                "type": "thread_operation",
                "operation": "switch_thread",
                "data": {
                    "thread_id": thread_id
                }
            }
            
            await self.websocket.send(json.dumps(thread_data))
            logger.info(f"📤 Switching to thread: {thread_id}")
            
            # Wait for thread switch response
            await self.wait_for_responses()
            
        except Exception as e:
            logger.error(f"Error switching thread: {str(e)}")
    
    async def get_suggestions(self, context: str = "general"):
        """Get chat suggestions."""
        if not self.websocket:
            logger.error("Not connected to chatbot")
            return
        
        try:
            thread_data = {
                "type": "thread_operation",
                "operation": "get_suggestions",
                "data": {
                    "context": context
                }
            }
            
            await self.websocket.send(json.dumps(thread_data))
            logger.info(f"📤 Requesting suggestions for context: {context}")
            
            # Wait for suggestions response
            await self.wait_for_responses()
            
        except Exception as e:
            logger.error(f"Error getting suggestions: {str(e)}")
    
    async def ping(self):
        """Send ping to check connection health."""
        if not self.websocket:
            logger.error("Not connected to chatbot")
            return
        
        try:
            ping_data = {"type": "ping"}
            await self.websocket.send(json.dumps(ping_data))
            logger.info("📤 Sent ping")
            
            # Wait for pong response
            await self.wait_for_responses()
            
        except Exception as e:
            logger.error(f"Error sending ping: {str(e)}")
    
    async def wait_for_responses(self, timeout: float = 10.0):
        """Wait for responses from the chatbot."""
        try:
            start_time = asyncio.get_event_loop().time()
            
            while True:
                # Check timeout
                if asyncio.get_event_loop().time() - start_time > timeout:
                    logger.warning("⏰ Response timeout reached")
                    break
                
                # Try to receive response with timeout
                try:
                    response = await asyncio.wait_for(self.websocket.recv(), timeout=1.0)
                    data = json.loads(response)
                    
                    # Handle different response types
                    await self.handle_response(data)
                    
                    # If we got a response, reset timeout
                    start_time = asyncio.get_event_loop().time()
                    
                except asyncio.TimeoutError:
                    # No response in this iteration, continue
                    continue
                except Exception as e:
                    logger.error(f"Error receiving response: {str(e)}")
                    break
                    
        except Exception as e:
            logger.error(f"Error waiting for responses: {str(e)}")
    
    async def handle_response(self, data: dict):
        """Handle different types of responses from the chatbot."""
        response_type = data.get("type", "unknown")
        
        if response_type == "user_message_received":
            logger.info(f"✅ User message received: {data.get('content', '')[:50]}...")
            logger.info(f"   Message ID: {data.get('message_id')}")
            logger.info(f"   Timestamp: {data.get('timestamp')}")
            
        elif response_type == "ai_response":
            logger.info(f"🤖 AI Response: {data.get('content', '')[:100]}...")
            logger.info(f"   Message ID: {data.get('message_id')}")
            logger.info(f"   Thread ID: {data.get('thread_id')}")
            logger.info(f"   Tokens Used: {data.get('tokens_used')}")
            logger.info(f"   Model: {data.get('model_used')}")
            
            # Update thread ID if not set
            if not self.thread_id:
                self.thread_id = data.get('thread_id')
            
        elif response_type == "thread_created":
            logger.info(f"📝 Thread created: {data.get('title', '')}")
            logger.info(f"   Thread ID: {data.get('thread_id')}")
            
            # Update thread ID
            self.thread_id = data.get('thread_id')
            
        elif response_type == "threads_listed":
            threads = data.get('threads', [])
            total = data.get('total', 0)
            logger.info(f"📋 Threads listed: {len(threads)} of {total} total")
            for thread in threads[:3]:  # Show first 3
                logger.info(f"   - {thread.get('title', 'No title')} (ID: {thread.get('id')})")
            
        elif response_type == "thread_switched":
            thread_id = data.get('thread_id')
            messages = data.get('messages', [])
            logger.info(f"🔄 Switched to thread: {thread_id}")
            logger.info(f"   Loaded {len(messages)} messages")
            
            # Update thread ID
            self.thread_id = thread_id
            
        elif response_type == "suggestions":
            suggestions = data.get('suggestions', [])
            logger.info(f"💡 Chat suggestions ({len(suggestions)}):")
            for i, suggestion in enumerate(suggestions, 1):
                logger.info(f"   {i}. {suggestion}")
                
        elif response_type == "pong":
            logger.info("🏓 Pong received - connection healthy")
            
        elif response_type == "error":
            logger.error(f"❌ Error: {data.get('message', 'Unknown error')}")
            if data.get('details'):
                logger.error(f"   Details: {data.get('details')}")
                
        else:
            logger.info(f"📨 Unknown response type: {response_type}")
            logger.info(f"   Data: {data}")
    
    async def disconnect(self):
        """Disconnect from the chatbot."""
        if self.websocket:
            await self.websocket.close()
            logger.info("🔌 Disconnected from chatbot")


async def run_chatbot_demo():
    """Run a demo of the chatbot functionality."""
    # Configuration
    websocket_url = "ws://localhost:8000/hw_gpt"  # Adjust port if needed
    user_id = "1"
    store_id = "35"
    branch_id = None  # Set to specific branch ID if needed
    
    # Create test client
    client = ChatbotTestClient(websocket_url, user_id, store_id, branch_id)
    
    try:
        # Connect to chatbot
        if not await client.connect():
            logger.error("Failed to connect to chatbot")
            return
        
        logger.info("🚀 Starting chatbot demo...")
        logger.info("=" * 50)
        
        # Test 1: Get chat suggestions
        logger.info("\n1️⃣ Getting chat suggestions...")
        await client.get_suggestions("general")
        await asyncio.sleep(2)
        
        # Test 2: Create a new thread
        logger.info("\n2️⃣ Creating new marketing thread...")
        await client.create_thread(
            title="Social Media Marketing Strategy",
            description="Planning social media campaigns and content",
            thread_type="marketing"
        )
        await asyncio.sleep(2)
        
        # Test 3: Send a message
        logger.info("\n3️⃣ Sending first message...")
        await client.send_message("Hello! I need help with my social media marketing strategy.")
        await asyncio.sleep(3)
        
        # Test 4: Send follow-up message
        logger.info("\n4️⃣ Sending follow-up message...")
        await client.send_message("What type of content should I create for my target audience?")
        await asyncio.sleep(3)
        
        # Test 5: List threads
        logger.info("\n5️⃣ Listing user threads...")
        await client.list_threads()
        await asyncio.sleep(2)
        
        # Test 6: Ping connection
        logger.info("\n6️⃣ Testing connection health...")
        await client.ping()
        await asyncio.sleep(1)
        
        # Test 7: Send another message
        logger.info("\n7️⃣ Sending another message...")
        await client.send_message("How can I measure the ROI of my marketing campaigns?")
        await asyncio.sleep(3)
        
        logger.info("\n🎉 Demo completed successfully!")
        
    except Exception as e:
        logger.error(f"Demo error: {str(e)}")
        
    finally:
        # Disconnect
        await client.disconnect()


if __name__ == "__main__":
    print("🚀 hwGpt WebSocket Chatbot Test Client")
    print("=" * 50)
    print("Make sure your FastAPI server is running on localhost:8000")
    print("This client will connect to the WebSocket endpoint and test the chatbot")
    print()
    
    # Run the demo
    asyncio.run(run_chatbot_demo())