#!/usr/bin/env python3
"""
Stress test for server_dashscope.py and server_openai.py

This test measures:
- Maximum concurrent request handling capacity
- Latency metrics (p50, p95, p99, min, max, avg)
- Throughput (requests per second)
- Success/failure rates

Instructions:
1. Start the DashScope server:
   uvicorn fastapi_server.server_dashscope:app --host 0.0.0.0 --port 8588
2. Start the OpenAI server:
   uvicorn fastapi_server.server_openai:app --host 0.0.0.0 --port 8589
3. Set environment variables:
   FAST_AUTH_KEYS=test-key-1,test-key-2
4. Run this test:
   pytest tests/test_stress_servers.py -v
   or
   python tests/test_stress_servers.py [--stream | --no-stream]
   
   Options:
   --stream      Test only streaming endpoints
   --no-stream   Test only non-streaming endpoints
   (no option)   Test both streaming and non-streaming (default)
"""
import os
import sys
import time
import asyncio
import statistics
import argparse
import uuid
from dataclasses import dataclass, field
from typing import List, Optional, Dict
from collections import defaultdict
import httpx
from loguru import logger

# Load environment variables (matching test_dashscope_client.py and test_openai_client.py)
from dotenv import load_dotenv
load_dotenv()

# Server URLs (matching test files)
DS_BASE_URL = os.getenv("DS_BASE_URL", "http://127.0.0.1:8588/api/")
OPENAI_BASE_URL = os.getenv("OPENAI_BASE_URL", "http://127.0.0.1:8589/v1")

# Normalize base URLs (remove trailing slashes)
DASHSCOPE_BASE_URL = DS_BASE_URL.rstrip("/")
OPENAI_BASE_URL = OPENAI_BASE_URL.rstrip("/")

# API Key (matching test files - use first key if comma-separated)
FAST_AUTH_KEYS = os.getenv("FAST_AUTH_KEYS", "test-key")
API_KEY = FAST_AUTH_KEYS.split(",")[0] if FAST_AUTH_KEYS else "test-key"


@dataclass
class RequestResult:
    """Result of a single request."""
    success: bool
    latency_ms: float
    status_code: Optional[int] = None
    error: Optional[str] = None
    response_size: int = 0


@dataclass
class StressTestResult:
    """Results from a stress test run."""
    server_name: str
    endpoint: str
    concurrency: int
    total_requests: int
    successful_requests: int
    failed_requests: int
    latencies_ms: List[float] = field(default_factory=list)
    throughput_rps: float = 0.0
    duration_seconds: float = 0.0
    
    @property
    def success_rate(self) -> float:
        """Calculate success rate as percentage."""
        if self.total_requests == 0:
            return 0.0
        return (self.successful_requests / self.total_requests) * 100
    
    @property
    def avg_latency_ms(self) -> float:
        """Calculate average latency."""
        if not self.latencies_ms:
            return 0.0
        return statistics.mean(self.latencies_ms)
    
    @property
    def min_latency_ms(self) -> float:
        """Calculate minimum latency."""
        if not self.latencies_ms:
            return 0.0
        return min(self.latencies_ms)
    
    @property
    def max_latency_ms(self) -> float:
        """Calculate maximum latency."""
        if not self.latencies_ms:
            return 0.0
        return max(self.latencies_ms)
    
    @property
    def p50_latency_ms(self) -> float:
        """Calculate 50th percentile latency."""
        if not self.latencies_ms:
            return 0.0
        return statistics.median(self.latencies_ms)
    
    @property
    def p95_latency_ms(self) -> float:
        """Calculate 95th percentile latency."""
        if not self.latencies_ms:
            return 0.0
        return self._percentile(self.latencies_ms, 95)
    
    @property
    def p99_latency_ms(self) -> float:
        """Calculate 99th percentile latency."""
        if not self.latencies_ms:
            return 0.0
        return self._percentile(self.latencies_ms, 99)
    
    @staticmethod
    def _percentile(data: List[float], percentile: int) -> float:
        """Calculate percentile value."""
        sorted_data = sorted(data)
        index = (percentile / 100) * (len(sorted_data) - 1)
        if index.is_integer():
            return sorted_data[int(index)]
        lower = sorted_data[int(index)]
        upper = sorted_data[int(index) + 1]
        return lower + (upper - lower) * (index - int(index))


async def make_dashscope_request(
    client: httpx.AsyncClient,
    app_id: str = "test-app",
    session_id: str = "test-session",
    stream: bool = False,
    message: str = "Hello, how are you?",
) -> RequestResult:
    """Make a request to the DashScope server."""
    # Use /api/v1/... if base URL contains /api/, otherwise /v1/...
    # The server supports both endpoints
    if "/api" in DASHSCOPE_BASE_URL:
        url = f"{DASHSCOPE_BASE_URL}/v1/apps/{app_id}/sessions/{session_id}/responses"
    else:
        url = f"{DASHSCOPE_BASE_URL}/api/v1/apps/{app_id}/sessions/{session_id}/responses"
    headers = {"Authorization": f"Bearer {API_KEY}"}
    payload = {
        "input": {
            "session_id": session_id,
            "messages": [
                {"role": "user", "content": message}
            ]
        },
        "stream": stream,
    }
    
    start_time = time.perf_counter()
    try:
        if stream:
            response_size = 0
            async with client.stream("POST", url, headers=headers, json=payload, timeout=120.0) as response:
                if response.status_code != 200:
                    error_text = await response.aread()
                    return RequestResult(
                        success=False,
                        latency_ms=(time.perf_counter() - start_time) * 1000,
                        status_code=response.status_code,
                        error=f"HTTP {response.status_code}",
                    )
                async for line in response.aiter_lines():
                    if line.startswith("data: "):
                        response_size += len(line)
                        # For stress testing, we can stop after receiving first chunk to measure latency
                        # Uncomment the break below if you want to measure time-to-first-byte only
                        # break
        else:
            response = await client.post(url, headers=headers, json=payload, timeout=60.0)
            response_size = len(response.content)
            if response.status_code != 200:
                return RequestResult(
                    success=False,
                    latency_ms=(time.perf_counter() - start_time) * 1000,
                    status_code=response.status_code,
                    error=response.text[:200],
                )
        
        latency_ms = (time.perf_counter() - start_time) * 1000
        return RequestResult(
            success=True,
            latency_ms=latency_ms,
            status_code=200,
            response_size=response_size,
        )
    except Exception as e:
        latency_ms = (time.perf_counter() - start_time) * 1000
        return RequestResult(
            success=False,
            latency_ms=latency_ms,
            error=str(e)[:200],
        )


async def make_openai_request(
    client: httpx.AsyncClient,
    stream: bool = False,
    message: str = "Hello, how are you?",
    thread_id: str = "test-thread",
) -> RequestResult:
    """Make a request to the OpenAI-compatible server."""
    url = f"{OPENAI_BASE_URL}/v1/chat/completions"
    headers = {"Authorization": f"Bearer {API_KEY}"}
    payload = {
        "model": "gpt-3.5-turbo",
        "messages": [
            {"role": "user", "content": message}
        ],
        "stream": stream,
        "thread_id": thread_id,
    }
    
    start_time = time.perf_counter()
    try:
        if stream:
            response_size = 0
            async with client.stream("POST", url, headers=headers, json=payload, timeout=120.0) as response:
                if response.status_code != 200:
                    error_text = await response.aread()
                    return RequestResult(
                        success=False,
                        latency_ms=(time.perf_counter() - start_time) * 1000,
                        status_code=response.status_code,
                        error=f"HTTP {response.status_code}",
                    )
                async for line in response.aiter_lines():
                    if line.startswith("data: "):
                        response_size += len(line)
                        # For stress testing, we can stop after receiving first chunk to measure latency
                        # Uncomment the break below if you want to measure time-to-first-byte only
                        # break
        else:
            response = await client.post(url, headers=headers, json=payload, timeout=60.0)
            response_size = len(response.content)
            if response.status_code != 200:
                return RequestResult(
                    success=False,
                    latency_ms=(time.perf_counter() - start_time) * 1000,
                    status_code=response.status_code,
                    error=response.text[:200],
                )
        
        latency_ms = (time.perf_counter() - start_time) * 1000
        return RequestResult(
            success=True,
            latency_ms=latency_ms,
            status_code=200,
            response_size=response_size,
        )
    except Exception as e:
        latency_ms = (time.perf_counter() - start_time) * 1000
        return RequestResult(
            success=False,
            latency_ms=latency_ms,
            error=str(e)[:200],
        )


async def run_stress_test(
    request_func,
    server_name: str,
    endpoint: str,
    concurrency: int,
    total_requests: int,
    stream: bool = False,
) -> StressTestResult:
    """Run a stress test with specified concurrency and total requests."""
    logger.info(f"Starting stress test: {server_name} - {endpoint} - Concurrency: {concurrency}, Total: {total_requests}, Stream: {stream}")
    
    # Create semaphore to limit concurrency
    semaphore = asyncio.Semaphore(concurrency)
    results: List[RequestResult] = []
    
    async def make_request_with_semaphore():
        async with semaphore:
            return await request_func()
    
    # Create tasks
    tasks = [make_request_with_semaphore() for _ in range(total_requests)]
    
    # Run all requests concurrently
    start_time = time.perf_counter()
    request_results = await asyncio.gather(*tasks, return_exceptions=True)
    end_time = time.perf_counter()
    
    # Process results
    for result in request_results:
        if isinstance(result, Exception):
            results.append(RequestResult(
                success=False,
                latency_ms=0.0,
                error=str(result)[:200],
            ))
        else:
            results.append(result)
    
    # Calculate metrics
    successful = [r for r in results if r.success]
    failed = [r for r in results if not r.success]
    latencies = [r.latency_ms for r in successful]
    
    duration = end_time - start_time
    throughput = len(successful) / duration if duration > 0 else 0
    
    return StressTestResult(
        server_name=server_name,
        endpoint=endpoint,
        concurrency=concurrency,
        total_requests=total_requests,
        successful_requests=len(successful),
        failed_requests=len(failed),
        latencies_ms=latencies,
        throughput_rps=throughput,
        duration_seconds=duration,
    )


def print_results(result: StressTestResult):
    """Print formatted stress test results."""
    print(f"\n{'='*80}")
    print(f"STRESS TEST RESULTS: {result.server_name}")
    print(f"{'='*80}")
    print(f"Endpoint:              {result.endpoint}")
    print(f"Concurrency:            {result.concurrency}")
    print(f"Total Requests:         {result.total_requests}")
    print(f"Successful:             {result.successful_requests} ({result.success_rate:.2f}%)")
    print(f"Failed:                 {result.failed_requests}")
    print(f"Duration:               {result.duration_seconds:.3f}s")
    print(f"Throughput:             {result.throughput_rps:.2f} req/s")
    print(f"\nLatency Metrics (ms):")
    print(f"  Min:                  {result.min_latency_ms:.2f}")
    print(f"  Max:                  {result.max_latency_ms:.2f}")
    print(f"  Average:              {result.avg_latency_ms:.2f}")
    print(f"  Median (p50):         {result.p50_latency_ms:.2f}")
    print(f"  p95:                  {result.p95_latency_ms:.2f}")
    print(f"  p99:                  {result.p99_latency_ms:.2f}")
    
    if result.failed_requests > 0:
        print(f"\nErrors encountered: {result.failed_requests}")
    print(f"{'='*80}\n")


async def test_dashscope_server(stream_mode: Optional[bool] = None):
    """Test the DashScope server with various concurrency levels.
    
    Args:
        stream_mode: If True, test only streaming. If False, test only non-streaming.
                     If None, test both.
    """
    print("\n" + "="*80)
    print("TESTING DASHSCOPE SERVER")
    if stream_mode is True:
        print("Mode: Streaming only")
    elif stream_mode is False:
        print("Mode: Non-streaming only")
    else:
        print("Mode: Both streaming and non-streaming")
    print("="*80)
    
    # Test configurations: (concurrency, total_requests, stream)
    all_test_configs = [
        (1, 10, False),      # Sequential, non-streaming
        (5, 25, False),      # Low concurrency
        (10, 50, False),     # Medium concurrency
        (20, 100, False),    # High concurrency
        (50, 200, False),    # Very high concurrency
        (1, 10, True),       # Sequential, streaming
        (10, 50, True),      # Medium concurrency, streaming
        (20, 100, True),     # High concurrency, streaming
    ]
    
    # Filter based on stream_mode
    if stream_mode is not None:
        test_configs = [cfg for cfg in all_test_configs if cfg[2] == stream_mode]
    else:
        test_configs = all_test_configs
    
    all_results = []
    
    async with httpx.AsyncClient(timeout=120.0) as client:
        for concurrency, total_requests, stream in test_configs:
            endpoint = f"/v1/apps/{{app_id}}/sessions/{{session_id}}/responses"
            if stream:
                endpoint += " (streaming)"
            
            # Create request function with client bound
            async def request_func():
                # Overall server capacity across many sessions:
                # use a unique session_id (thread_id) per request to avoid per-thread contention
                session_id = str(uuid.uuid4())
                return await make_dashscope_request(
                    client,
                    app_id=f"test-app-{concurrency}",
                    session_id=session_id,
                    stream=stream,
                    message=f"Test message for concurrency {concurrency}",
                )
            
            result = await run_stress_test(
                request_func,
                "DashScope Server",
                endpoint,
                concurrency,
                total_requests,
                stream,
            )
            
            all_results.append(result)
            print_results(result)
            
            # Small delay between test runs
            await asyncio.sleep(1)
    
    return all_results


async def test_openai_server(stream_mode: Optional[bool] = None):
    """Test the OpenAI-compatible server with various concurrency levels.
    
    Args:
        stream_mode: If True, test only streaming. If False, test only non-streaming.
                     If None, test both.
    """
    print("\n" + "="*80)
    print("TESTING OPENAI SERVER")
    if stream_mode is True:
        print("Mode: Streaming only")
    elif stream_mode is False:
        print("Mode: Non-streaming only")
    else:
        print("Mode: Both streaming and non-streaming")
    print("="*80)
    
    # Test configurations: (concurrency, total_requests, stream)
    all_test_configs = [
        (1, 10, False),      # Sequential, non-streaming
        (5, 25, False),      # Low concurrency
        (10, 50, False),     # Medium concurrency
        (20, 100, False),    # High concurrency
        (50, 200, False),    # Very high concurrency
        (1, 10, True),       # Sequential, streaming
        (10, 50, True),      # Medium concurrency, streaming
        (20, 100, True),     # High concurrency, streaming
    ]
    
    # Filter based on stream_mode
    if stream_mode is not None:
        test_configs = [cfg for cfg in all_test_configs if cfg[2] == stream_mode]
    else:
        test_configs = all_test_configs
    
    all_results = []
    
    async with httpx.AsyncClient(timeout=120.0) as client:
        for concurrency, total_requests, stream in test_configs:
            endpoint = "/v1/chat/completions"
            if stream:
                endpoint += " (streaming)"
            
            # Create request function with client bound
            async def request_func():
                # Overall server capacity across many sessions:
                # use a unique thread_id per request to avoid per-thread contention
                thread_id = str(uuid.uuid4())
                return await make_openai_request(
                    client,
                    stream=stream,
                    message=f"Test message for concurrency {concurrency}",
                    thread_id=thread_id,
                )
            
            result = await run_stress_test(
                request_func,
                "OpenAI Server",
                endpoint,
                concurrency,
                total_requests,
                stream,
            )
            
            all_results.append(result)
            print_results(result)
            
            # Small delay between test runs
            await asyncio.sleep(1)
    
    return all_results


def print_summary(results: List[StressTestResult], header: str):
    """Print a summary of test results for a single experiment.
    
    Args:
        results: List of stress test results
        header: Header text to print for this experiment
    """
    print("\n" + "="*80)
    print(header)
    print("="*80)
    
    # Separate streaming and non-streaming results
    streaming_results = [r for r in results if "streaming" in r.endpoint]
    non_streaming_results = [r for r in results if "streaming" not in r.endpoint]
    
    # Print non-streaming results if available
    if non_streaming_results:
        print("\nNon-Streaming:")
        print(f"{'Concurrency':<15} {'Requests':<12} {'Success %':<12} {'Throughput (req/s)':<20} {'Avg Latency (ms)':<18} {'p95 (ms)':<12} {'p99 (ms)':<12}")
        print("-" * 110)
        for result in non_streaming_results:
            print(f"{result.concurrency:<15} {result.total_requests:<12} {result.success_rate:<11.2f}% "
                  f"{result.throughput_rps:<20.2f} {result.avg_latency_ms:<18.2f} "
                  f"{result.p95_latency_ms:<12.2f} {result.p99_latency_ms:<12.2f}")
    
    # Print streaming results if available
    if streaming_results:
        print("\nStreaming:")
        print(f"{'Concurrency':<15} {'Requests':<12} {'Success %':<12} {'Throughput (req/s)':<20} {'Avg Latency (ms)':<18} {'p95 (ms)':<12} {'p99 (ms)':<12}")
        print("-" * 110)
        for result in streaming_results:
            print(f"{result.concurrency:<15} {result.total_requests:<12} {result.success_rate:<11.2f}% "
                  f"{result.throughput_rps:<20.2f} {result.avg_latency_ms:<18.2f} "
                  f"{result.p95_latency_ms:<12.2f} {result.p99_latency_ms:<12.2f}")
    
    print("\n" + "="*80)


async def main(stream_mode: Optional[bool] = None):
    """Main function to run all stress tests.
    
    Args:
        stream_mode: If True, test only streaming. If False, test only non-streaming.
                     If None, test both.
    """
    print("\n" + "="*80)
    print("STRESS TEST FOR FASTAPI SERVERS")
    print("="*80)
    print(f"DashScope Server URL: {DS_BASE_URL}")
    print(f"OpenAI Server URL: {OPENAI_BASE_URL}")
    print(f"API Key: {API_KEY[:8]}..." if len(API_KEY) > 8 else f"API Key: {API_KEY}")
    if stream_mode is True:
        print("Testing Mode: Streaming only")
    elif stream_mode is False:
        print("Testing Mode: Non-streaming only")
    else:
        print("Testing Mode: Both streaming and non-streaming")
    print("="*80)
    
    # Check if servers are reachable
    async with httpx.AsyncClient(timeout=5.0) as client:
        try:
            # Health endpoint is at root, not under /api/
            # Extract base URL without /api/ path
            if "/api" in DASHSCOPE_BASE_URL:
                base_without_api = DASHSCOPE_BASE_URL.split("/api")[0]
            else:
                base_without_api = DASHSCOPE_BASE_URL.rstrip("/")
            response = await client.get(f"{base_without_api}/health")
            if response.status_code != 200:
                logger.warning(f"DashScope server health check failed: {response.status_code}")
        except Exception as e:
            logger.error(f"Cannot reach DashScope server at {DASHSCOPE_BASE_URL}: {e}")
            logger.info("Please start the server: uvicorn fastapi_server.server_dashscope:app --host 0.0.0.0 --port 8588")
        
        # try:
        #     response = await client.get(f"{OPENAI_BASE_URL}/health")
        #     if response.status_code != 200:
        #         logger.warning(f"OpenAI server health check failed: {response.status_code}")
        # except Exception as e:
        #     logger.error(f"Cannot reach OpenAI server at {OPENAI_BASE_URL}: {e}")
        #     logger.info("Please start the server: uvicorn fastapi_server.server_openai:app --host 0.0.0.0 --port 8589")
    
    # Run stress tests
    dashscope_results = await test_dashscope_server(stream_mode)
    # openai_results = await test_openai_server(stream_mode)
    
    # Print summaries
    print_summary(dashscope_results, "DASHSCOPE SERVER SUMMARY")
    # print_summary(openai_results, "OPENAI SERVER SUMMARY")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Stress test for FastAPI servers (DashScope and OpenAI compatible)",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  python tests/test_stress_servers.py              # Test both streaming and non-streaming
  python tests/test_stress_servers.py --stream     # Test only streaming endpoints
  python tests/test_stress_servers.py --no-stream  # Test only non-streaming endpoints
        """
    )
    group = parser.add_mutually_exclusive_group()
    group.add_argument(
        "--stream",
        action="store_true",
        help="Test only streaming endpoints"
    )
    group.add_argument(
        "--no-stream",
        action="store_true",
        dest="no_stream",
        help="Test only non-streaming endpoints"
    )
    
    args = parser.parse_args()
    
    # Determine stream_mode from arguments
    if args.stream:
        stream_mode = True
    elif args.no_stream:
        stream_mode = False
    else:
        stream_mode = None  # Test both
    
    asyncio.run(main(stream_mode))