"""
Test script for VIJ AI Assistant - LIVE SPACE
Run this to verify your Hugging Face Space is working correctly
"""
import requests
import json
import time

# Your live Hugging Face Space URL
API_URL = "https://vishwas896-vish-ai-space.hf.space"

def print_header(text):
    print("\n" + "=" * 60)
    print(f"  {text}")
    print("=" * 60)

def test_root():
    """Test root endpoint"""
    print_header("Testing Root Endpoint")
    try:
        response = requests.get(f"{API_URL}/")
        print(f"✅ Status: {response.status_code}")
        print(f"Response: {json.dumps(response.json(), indent=2)}")
        return True
    except Exception as e:
        print(f"❌ Error: {e}")
        return False

def test_health():
    """Test health endpoint"""
    print_header("Testing Health Endpoint")
    try:
        response = requests.get(f"{API_URL}/health")
        print(f"✅ Status: {response.status_code}")
        print(f"Response: {json.dumps(response.json(), indent=2)}")
        return True
    except Exception as e:
        print(f"❌ Error: {e}")
        return False

def test_system():
    """Test system status endpoint"""
    print_header("Testing System Status")
    try:
        response = requests.get(f"{API_URL}/system")
        print(f"✅ Status: {response.status_code}")
        print(f"Response: {json.dumps(response.json(), indent=2)}")
        return True
    except Exception as e:
        print(f"❌ Error: {e}")
        return False

def test_vij_info():
    """Test VIJ info endpoint"""
    print_header("Testing VIJ Info Endpoint")
    try:
        response = requests.get(f"{API_URL}/api/vij_info")
        print(f"✅ Status: {response.status_code}")
        data = response.json()
        print(f"Categories: {data.get('categories', [])}")
        print(f"Knowledge entries: {len(data.get('knowledge_base', {}))}")
        return True
    except Exception as e:
        print(f"❌ Error: {e}")
        return False

def test_assist():
    """Test quick assist endpoint"""
    print_header("Testing Quick Assist Endpoint")
    try:
        payload = {
            "text": "What services does VIJ provide?"
        }
        print(f"📤 Sending: {payload['text']}")
        response = requests.post(
            f"{API_URL}/api/assist",
            json=payload,
            timeout=30
        )
        print(f"✅ Status: {response.status_code}")
        data = response.json()
        print(f"📥 Reply: {data.get('reply', 'No reply')}")
        print(f"VIJ context used: {data.get('vij_context_used', False)}")
        return True
    except Exception as e:
        print(f"❌ Error: {e}")
        return False

def test_streaming():
    """Test streaming chat endpoint"""
    print_header("Testing Streaming Chat Endpoint")
    try:
        payload = {
            "text": "Tell me about VIJ's founder",
            "task": "assistant"
        }
        print(f"📤 Sending: {payload['text']}")
        print("🔄 Streaming response...")
        
        response = requests.post(
            f"{API_URL}/api/chat_stream",
            json=payload,
            stream=True,
            timeout=60
        )
        
        full_text = ""
        token_count = 0
        
        for line in response.iter_lines():
            if line:
                line = line.decode('utf-8')
                if line.startswith('data: '):
                    try:
                        data = json.loads(line[6:])
                        if data.get('done'):
                            full_text = data.get('complete_text', full_text)
                            break
                        else:
                            token = data.get('token', '')
                            full_text += token
                            token_count += 1
                            print(token, end='', flush=True)
                    except json.JSONDecodeError:
                        pass
        
        print(f"\n\n✅ Streamed {token_count} tokens")
        print(f"📥 Complete response: {full_text[:200]}...")
        return True
    except Exception as e:
        print(f"❌ Error: {e}")
        return False

def run_all_tests():
    """Run all tests"""
    print("\n" + "🧪" * 30)
    print("  VIJ AI Assistant - Test Suite")
    print("🧪" * 30)
    
    tests = [
        ("Root Endpoint", test_root),
        ("Health Check", test_health),
        ("System Status", test_system),
        ("VIJ Info", test_vij_info),
        ("Quick Assist", test_assist),
        ("Streaming Chat", test_streaming),
    ]
    
    results = []
    for test_name, test_func in tests:
        try:
            result = test_func()
            results.append((test_name, result))
            time.sleep(1)  # Small delay between tests
        except Exception as e:
            print(f"❌ {test_name} failed with exception: {e}")
            results.append((test_name, False))
    
    # Summary
    print_header("Test Summary")
    passed = sum(1 for _, result in results if result)
    total = len(results)
    
    for test_name, result in results:
        status = "✅ PASS" if result else "❌ FAIL"
        print(f"{status} - {test_name}")
    
    print(f"\n📊 Results: {passed}/{total} tests passed")
    
    if passed == total:
        print("🎉 All tests passed! Your VIJ AI Space is working perfectly!")
    else:
        print("⚠️  Some tests failed. Check the logs above for details.")

if __name__ == "__main__":
    print("\n🌐 Testing your LIVE Hugging Face Space:", API_URL)
    print("   Space: https://huggingface.co/spaces/Vishwas896/vish-ai-space\n")
    
    input("Press Enter to start tests...")
    run_all_tests()