#!/usr/bin/env python3
"""
Test script for Archie conversation - automated test
"""

import asyncio
import websockets
import json
import sys
from pathlib import Path

WS_URL = "ws://localhost:3050/ws/orchestrator"
TOKEN_FILE = Path("/tmp/token.txt")


async def test_conversation():
    """Test a simple conversation with Archie"""

    # Load token if available
    token = None
    if TOKEN_FILE.exists():
        with open(TOKEN_FILE) as f:
            token = f.read().strip()
        print(f"✓ Using JWT token")
    else:
        print("⚠ No JWT token, using anonymous")

    url = f"{WS_URL}?token={token}" if token else WS_URL

    print(f"\n🔌 Connecting to {WS_URL}...")

    try:
        async with websockets.connect(url, ping_interval=None) as websocket:
            # Wait for connection confirmation
            conn_msg = await asyncio.wait_for(websocket.recv(), timeout=5.0)
            conn_data = json.loads(conn_msg)

            if conn_data.get("type") == "connected":
                print(f"✓ {conn_data.get('message')}")
                print(f"  User: {conn_data.get('user_id')}\n")

            # Test query
            test_query = "What is the capital of France?"
            print(f"📤 Sending: {test_query}\n")

            query_msg = {
                "type": "query",
                "query": test_query,
                "conversation_id": "test-conv-001",
                "conversation_history": []
            }

            await websocket.send(json.dumps(query_msg))
            print("📥 Receiving response:\n")
            print("Archie: ", end="", flush=True)

            full_response = ""
            completed = False

            while not completed:
                try:
                    response = await asyncio.wait_for(websocket.recv(), timeout=30.0)
                    data = json.loads(response)
                    msg_type = data.get("type")

                    if msg_type == "chunk":
                        content = data.get("content", "")
                        print(content, end="", flush=True)
                        full_response += content

                    elif msg_type == "completed":
                        print("\n")
                        confidence = data.get("confidence")
                        if confidence:
                            print(f"\n✓ Completed [Confidence: {confidence:.2%}]")
                        else:
                            print("\n✓ Completed")
                        completed = True

                    elif msg_type == "error":
                        error_msg = data.get("error", "Unknown error")
                        print(f"\n❌ Error: {error_msg}")
                        return False

                    elif msg_type == "progress":
                        # Show progress if available
                        status = data.get("status", "")
                        if status:
                            print(f"\n  [{status}]", end="", flush=True)

                except asyncio.TimeoutError:
                    print("\n\n⏱️  Response timeout")
                    return False

            print(f"\n\n📊 Response Length: {len(full_response)} characters")
            print(f"✅ TEST PASSED: Successfully completed conversation!")
            return True

    except websockets.exceptions.WebSocketException as e:
        print(f"\n❌ WebSocket error: {e}")
        return False
    except Exception as e:
        print(f"\n❌ Error: {e}")
        import traceback
        traceback.print_exc()
        return False


def main():
    """Main entry point"""
    print("\n" + "="*70)
    print("🧪 Archie Conversation Test")
    print("="*70)

    success = asyncio.run(test_conversation())
    sys.exit(0 if success else 1)


if __name__ == "__main__":
    main()