#!/usr/bin/env python3
"""
Test suite for FAISSx server authentication and tenant isolation.

This test suite verifies that the FAISSx server correctly:
1. Rejects requests without API keys when authentication is enabled
2. Rejects requests with invalid API keys
3. Accepts requests with valid API keys
4. Properly isolates tenants from each other's data
5. Allows health checks (ping) without authentication
"""

import subprocess
import time
import unittest
from pathlib import Path

import numpy as np

from faissx import client as faiss


class TestServerAuthentication(unittest.TestCase):
    """Test authentication and tenant isolation functionality."""

    @classmethod
    def setUpClass(cls):
        """Start a test server with authentication enabled."""
        cls.auth_keys = {"123": "tenant-1", "abc": "tenant-2"}
        cls.server_process = None
        cls._start_test_server()
        time.sleep(2)  # Give server time to start

    @classmethod
    def tearDownClass(cls):
        """Stop the test server."""
        if cls.server_process:
            cls.server_process.terminate()
            cls.server_process.wait()

    @classmethod
    def _start_test_server(cls):
        """Start the FAISSx server for testing."""
        # Use a different port to avoid conflicts
        port = 45679
        auth_keys_str = "123:tenant-1,abc:tenant-2"

        cmd = [
            "faissx.server", "run",
            "--port", str(port),
            "--enable-auth",
            "--auth-keys", auth_keys_str,
            "--log-level", "WARNING"
        ]

        cls.server_process = subprocess.Popen(
            cmd,
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL
        )
        cls.server_port = port

    def test_no_auth_key_rejected(self):
        """Test that requests without API keys are rejected."""
        with self.assertRaises(Exception) as context:
            faiss.configure(
                server=f"tcp://localhost:{self.server_port}",
                api_key=None,
                tenant_id=None
            )

            index = faiss.IndexFlatL2(128)
            vectors = np.random.rand(10, 128).astype(np.float32)
            index.add(vectors)

        # Client may send empty string when api_key=None, resulting in "Invalid API key"
        error_msg = str(context.exception)
        self.assertTrue(
            "API key required" in error_msg or "Invalid API key" in error_msg,
            f"Expected auth error, got: {error_msg}"
        )

    def test_invalid_auth_key_rejected(self):
        """Test that requests with invalid API keys are rejected."""
        with self.assertRaises(Exception) as context:
            faiss.configure(
                server=f"tcp://localhost:{self.server_port}",
                api_key="invalid-key",
                tenant_id="some-tenant"
            )

            index = faiss.IndexFlatL2(128)
            vectors = np.random.rand(10, 128).astype(np.float32)
            index.add(vectors)

        self.assertIn("Invalid API key", str(context.exception))

    def test_valid_auth_key_accepted_tenant1(self):
        """Test that valid API key for tenant-1 is accepted."""
        faiss.configure(
            server=f"tcp://localhost:{self.server_port}",
            api_key="123",
            tenant_id="tenant-1"
        )

        index = faiss.IndexFlatL2(128)
        vectors = np.random.rand(10, 128).astype(np.float32)
        index.add(vectors)

        # Verify we can search
        query = np.random.rand(1, 128).astype(np.float32)
        distances, indices = index.search(query, 5)

        self.assertEqual(index.ntotal, 10)
        self.assertEqual(len(distances[0]), 5)

    def test_valid_auth_key_accepted_tenant2(self):
        """Test that valid API key for tenant-2 is accepted."""
        faiss.configure(
            server=f"tcp://localhost:{self.server_port}",
            api_key="abc",
            tenant_id="tenant-2"
        )

        index = faiss.IndexFlatL2(128)
        vectors = np.random.rand(10, 128).astype(np.float32)
        index.add(vectors)

        # Verify we can search
        query = np.random.rand(1, 128).astype(np.float32)
        distances, indices = index.search(query, 5)

        self.assertEqual(index.ntotal, 10)
        self.assertEqual(len(distances[0]), 5)

    def test_tenant_isolation(self):
        """Test that tenants cannot access each other's data."""
        # Create data with tenant-1
        faiss.configure(
            server=f"tcp://localhost:{self.server_port}",
            api_key="123",
            tenant_id="tenant-1"
        )

        index1 = faiss.IndexFlatL2(128)
        vectors1 = np.random.rand(10, 128).astype(np.float32)
        index1.add(vectors1)

        self.assertEqual(index1.ntotal, 10)

        # Switch to tenant-2 - should not see tenant-1's data
        faiss.configure(
            server=f"tcp://localhost:{self.server_port}",
            api_key="abc",
            tenant_id="tenant-2"
        )

        index2 = faiss.IndexFlatL2(128)
        self.assertEqual(index2.ntotal, 0)  # Should be empty

        # Add different data to tenant-2
        vectors2 = np.random.rand(5, 128).astype(np.float32)
        index2.add(vectors2)
        self.assertEqual(index2.ntotal, 5)  # Should only see tenant-2's data

    def test_ping_without_auth(self):
        """Test that ping works without authentication (health check)."""
        # This would require a direct ZMQ client since our python client
        # always sends auth info. For now, we'll test that ping works with auth.
        faiss.configure(
            server=f"tcp://localhost:{self.server_port}",
            api_key="123",
            tenant_id="tenant-1"
        )

        # Just test that we can create an index (indirect ping test)
        index = faiss.IndexFlatL2(128)
        self.assertIsNotNone(index)


class TestServerAuthenticationCLI(unittest.TestCase):
    """Test CLI argument parsing for authentication."""

    def test_auth_keys_format(self):
        """Test that the correct auth keys format is documented."""
        # This is more of a documentation test
        expected_format = "key1:tenant1,key2:tenant2"
        auth_keys = {}

        # Simulate CLI parsing
        for key_pair in expected_format.split(","):
            api_key, tenant_id = key_pair.strip().split(":")
            auth_keys[api_key] = tenant_id

        expected = {"key1": "tenant1", "key2": "tenant2"}
        self.assertEqual(auth_keys, expected)


if __name__ == "__main__":
    # Check if we're in the project root
    if not Path("faissx").exists():
        print("Error: This test must be run from the project root directory")
        exit(1)

    unittest.main()
