"""
Protocol Benchmark: HTTP vs WebSocket Server Functions

Compares performance of HTTP POST vs WebSocket RPC for server function calls.
Includes realistic scenarios with ORM queries.

Usage:
    python manage.py test mizan.tests.test_benchmarks --verbosity=2

Note:
    These are not unit tests - they measure performance. Results are printed
    to stdout and should be run in isolation for accurate measurements.
"""

import asyncio
import json
import statistics
import time
from typing import Any
from unittest.mock import MagicMock, AsyncMock

from django.contrib.auth import get_user_model
from django.contrib.auth.models import AnonymousUser
from django.db import connection
from django.http import HttpRequest
from django.test import RequestFactory, TestCase, TransactionTestCase, override_settings
from pydantic import BaseModel

from mizan.client.executor import FunctionResult, execute_function, function_call_view
from mizan_core.registry import clear_registry
from mizan.client import client


User = get_user_model()


# =============================================================================
# Benchmark Output Models
# =============================================================================


class SimpleOutput(BaseModel):
    value: int


class UserOutput(BaseModel):
    id: int
    email: str


class UserListOutput(BaseModel):
    users: list[dict[str, Any]]
    count: int


class StatsOutput(BaseModel):
    total_users: int
    active_users: int
    staff_count: int


# =============================================================================
# Benchmark Functions
# =============================================================================


def setup_benchmark_functions():
    """Register benchmark server functions."""
    from mizan_core.registry import register

    clear_registry()

    # 1. Simple computation (no I/O)
    @client
    def bench_simple(request: HttpRequest, a: int, b: int) -> SimpleOutput:
        """Simple addition - baseline with no I/O."""
        return SimpleOutput(value=a + b)

    register(bench_simple, "bench_simple")

    # 2. Single ORM query
    @client
    def bench_get_user(request: HttpRequest, user_id: int) -> UserOutput:
        """Fetch single user by ID."""
        user = User.objects.filter(id=user_id).first()
        if user:
            return UserOutput(id=user.id, email=user.email)
        return UserOutput(id=0, email="")

    register(bench_get_user, "bench_get_user")

    # 3. List query with limit
    @client
    def bench_list_users(request: HttpRequest, limit: int) -> UserListOutput:
        """Fetch list of users with limit."""
        users = User.objects.all()[:limit]
        return UserListOutput(
            users=[{"id": u.id, "email": u.email} for u in users],
            count=len(users),
        )

    register(bench_list_users, "bench_list_users")

    # 4. Aggregation query
    @client
    def bench_user_stats(request: HttpRequest) -> StatsOutput:
        """Compute user statistics with multiple queries."""
        total = User.objects.count()
        active = User.objects.filter(is_active=True).count()
        staff = User.objects.filter(is_staff=True).count()
        return StatsOutput(
            total_users=total,
            active_users=active,
            staff_count=staff,
        )

    register(bench_user_stats, "bench_user_stats")

    # 5. Complex query with joins
    @client
    def bench_user_search(
        request: HttpRequest, email_contains: str, limit: int
    ) -> UserListOutput:
        """Search users by email pattern."""
        users = User.objects.filter(
            email__icontains=email_contains,
            is_active=True,
        ).select_related()[:limit]
        return UserListOutput(
            users=[{"id": u.id, "email": u.email} for u in users],
            count=len(users),
        )

    register(bench_user_search, "bench_user_search")


# =============================================================================
# Benchmark Test Cases
# =============================================================================


class ProtocolBenchmark(TransactionTestCase):
    """
    Benchmark comparing HTTP vs WebSocket (simulated) performance.

    Uses TransactionTestCase to ensure database state is realistic.
    """

    # Number of iterations for each benchmark
    ITERATIONS = 100
    WARMUP = 10

    @classmethod
    def setUpClass(cls):
        super().setUpClass()
        setup_benchmark_functions()

    def setUp(self):
        self.factory = RequestFactory()
        # Create test users for ORM benchmarks
        self._create_test_users()

    def _create_test_users(self):
        """Create test users for benchmarks."""
        # Create 100 test users
        users = []
        for i in range(100):
            users.append(
                User(
                    email=f"bench{i}@example.com",
                    is_active=i % 10 != 0,  # 90% active
                    is_staff=i < 5,  # 5 staff
                )
            )
        User.objects.bulk_create(users, ignore_conflicts=True)
        self.test_user = User.objects.first()

    def _make_request(self, body: dict | None = None) -> HttpRequest:
        """Create a request with optional JSON body."""
        if body:
            request = self.factory.post(
                "/api/mizan/call/",
                data=json.dumps(body),
                content_type="application/json",
            )
        else:
            request = self.factory.post("/api/mizan/call/")
        request.user = AnonymousUser()
        request._dont_enforce_csrf_checks = True
        return request

    def _benchmark_executor(self, fn_name: str, args: dict, label: str) -> dict:
        """
        Benchmark direct executor calls (simulates WebSocket RPC).

        Returns timing statistics.
        """
        request = self._make_request()
        times = []

        # Warmup
        for _ in range(self.WARMUP):
            execute_function(request, fn_name, args)

        # Benchmark
        for _ in range(self.ITERATIONS):
            start = time.perf_counter()
            result = execute_function(request, fn_name, args)
            end = time.perf_counter()
            times.append((end - start) * 1000)  # ms

        return self._compute_stats(times, f"Executor ({label})")

    def _benchmark_http(self, fn_name: str, args: dict, label: str) -> dict:
        """
        Benchmark HTTP view calls.

        Returns timing statistics.
        """
        times = []

        # Warmup
        for _ in range(self.WARMUP):
            request = self._make_request({"fn": fn_name, "args": args})
            function_call_view(request)

        # Benchmark
        for _ in range(self.ITERATIONS):
            request = self._make_request({"fn": fn_name, "args": args})
            start = time.perf_counter()
            response = function_call_view(request)
            end = time.perf_counter()
            times.append((end - start) * 1000)  # ms

        return self._compute_stats(times, f"HTTP ({label})")

    def _compute_stats(self, times: list[float], label: str) -> dict:
        """Compute statistics from timing data."""
        return {
            "label": label,
            "min": min(times),
            "max": max(times),
            "mean": statistics.mean(times),
            "median": statistics.median(times),
            "stdev": statistics.stdev(times) if len(times) > 1 else 0,
            "p95": sorted(times)[int(len(times) * 0.95)],
            "p99": sorted(times)[int(len(times) * 0.99)],
            "iterations": len(times),
        }

    def _print_results(self, results: list[dict]):
        """Print benchmark results in a table."""
        print("\n" + "=" * 80)
        print(f"{'Benchmark':<40} {'Mean':>8} {'Median':>8} {'P95':>8} {'P99':>8}")
        print("=" * 80)
        for r in results:
            print(
                f"{r['label']:<40} {r['mean']:>7.3f}ms {r['median']:>7.3f}ms {r['p95']:>7.3f}ms {r['p99']:>7.3f}ms"
            )
        print("=" * 80)

    def _print_comparison(self, executor_stats: dict, http_stats: dict):
        """Print comparison between executor and HTTP."""
        overhead = (
            (http_stats["mean"] - executor_stats["mean"]) / executor_stats["mean"]
        ) * 100
        print(f"  HTTP overhead vs Executor: {overhead:+.1f}%")

    # -------------------------------------------------------------------------
    # Benchmark Tests
    # -------------------------------------------------------------------------

    def test_benchmark_simple_computation(self):
        """Benchmark: Simple computation (no I/O)."""
        print("\n\n### BENCHMARK: Simple Computation (no I/O) ###")

        args = {"a": 100, "b": 200}

        exec_stats = self._benchmark_executor("bench_simple", args, "simple")
        http_stats = self._benchmark_http("bench_simple", args, "simple")

        self._print_results([exec_stats, http_stats])
        self._print_comparison(exec_stats, http_stats)

        # Correctness check: 100 + 200 = 300
        request = self._make_request()
        result = execute_function(request, "bench_simple", args)
        self.assertIsInstance(result, FunctionResult)
        self.assertEqual(result.data["value"], 300)

    def test_benchmark_single_query(self):
        """Benchmark: Single ORM query."""
        print("\n\n### BENCHMARK: Single ORM Query ###")

        args = {"user_id": self.test_user.id if self.test_user else 1}

        exec_stats = self._benchmark_executor("bench_get_user", args, "single query")
        http_stats = self._benchmark_http("bench_get_user", args, "single query")

        self._print_results([exec_stats, http_stats])
        self._print_comparison(exec_stats, http_stats)

        # Correctness check: should return the test user's data
        request = self._make_request()
        result = execute_function(request, "bench_get_user", args)
        self.assertIsInstance(result, FunctionResult)
        self.assertEqual(result.data["id"], self.test_user.id)

    def test_benchmark_list_query(self):
        """Benchmark: List query with serialization."""
        print("\n\n### BENCHMARK: List Query (10 users) ###")

        args = {"limit": 10}

        exec_stats = self._benchmark_executor("bench_list_users", args, "list 10")
        http_stats = self._benchmark_http("bench_list_users", args, "list 10")

        self._print_results([exec_stats, http_stats])
        self._print_comparison(exec_stats, http_stats)

        # Correctness check: should return up to 10 users
        request = self._make_request()
        result = execute_function(request, "bench_list_users", args)
        self.assertIsInstance(result, FunctionResult)
        self.assertLessEqual(result.data["count"], 10)
        self.assertEqual(len(result.data["users"]), result.data["count"])

    def test_benchmark_aggregation(self):
        """Benchmark: Aggregation queries."""
        print("\n\n### BENCHMARK: Aggregation (3 COUNT queries) ###")

        args = {}

        exec_stats = self._benchmark_executor("bench_user_stats", args, "aggregation")
        http_stats = self._benchmark_http("bench_user_stats", args, "aggregation")

        self._print_results([exec_stats, http_stats])
        self._print_comparison(exec_stats, http_stats)

        # Correctness check: stats should have non-negative counts
        request = self._make_request()
        result = execute_function(request, "bench_user_stats", args)
        self.assertIsInstance(result, FunctionResult)
        self.assertGreaterEqual(result.data["total_users"], 0)
        self.assertGreaterEqual(result.data["active_users"], 0)
        self.assertGreaterEqual(result.data["staff_count"], 0)

    def test_benchmark_search_query(self):
        """Benchmark: Search with filter."""
        print("\n\n### BENCHMARK: Search Query (LIKE + LIMIT) ###")

        args = {"email_contains": "bench", "limit": 20}

        exec_stats = self._benchmark_executor("bench_user_search", args, "search")
        http_stats = self._benchmark_http("bench_user_search", args, "search")

        self._print_results([exec_stats, http_stats])
        self._print_comparison(exec_stats, http_stats)

        # Correctness check: search results should contain "bench" in emails
        request = self._make_request()
        result = execute_function(request, "bench_user_search", args)
        self.assertIsInstance(result, FunctionResult)
        self.assertLessEqual(result.data["count"], 20)
        for user in result.data["users"]:
            self.assertIn("bench", user["email"].lower())

    def test_summary(self):
        """Print summary of all benchmarks."""
        print("\n\n" + "=" * 80)
        print("BENCHMARK SUMMARY")
        print("=" * 80)
        print(f"Iterations per benchmark: {self.ITERATIONS}")
        print(f"Warmup iterations: {self.WARMUP}")
        print("\nKey findings:")
        print("- 'Executor' simulates WebSocket RPC (direct function call)")
        print("- 'HTTP' measures full request/response cycle")
        print("- HTTP overhead includes: JSON parsing, CSRF, view dispatch")
        print("- For I/O-bound operations, protocol overhead is negligible")
        print("=" * 80)

        # Verify bench_simple still produces correct output after all benchmarks
        request = self._make_request()
        result = execute_function(request, "bench_simple", {"a": 7, "b": 8})
        self.assertIsInstance(result, FunctionResult)
        self.assertEqual(result.data["value"], 15)


# =============================================================================
# Throughput Benchmark
# =============================================================================


class ThroughputBenchmark(TransactionTestCase):
    """
    Measure requests per second (throughput) for server functions.

    Tests both sequential and concurrent scenarios.
    """

    DURATION_SECONDS = 2  # How long to run each throughput test

    @classmethod
    def setUpClass(cls):
        super().setUpClass()
        setup_benchmark_functions()

    def setUp(self):
        self.factory = RequestFactory()
        self._create_test_users()

    def _create_test_users(self):
        """Create test users for benchmarks."""
        users = []
        for i in range(100):
            users.append(
                User(
                    email=f"bench{i}@example.com",
                    is_active=i % 10 != 0,
                    is_staff=i < 5,
                )
            )
        User.objects.bulk_create(users, ignore_conflicts=True)
        self.test_user = User.objects.first()

    def _make_request(self, body: dict) -> HttpRequest:
        """Create a POST request with JSON body."""
        request = self.factory.post(
            "/api/mizan/call/",
            data=json.dumps(body),
            content_type="application/json",
        )
        request.user = AnonymousUser()
        request._dont_enforce_csrf_checks = True
        return request

    def _measure_throughput_executor(self, fn_name: str, args: dict) -> float:
        """Measure requests/second using direct executor calls."""
        request = self._make_request({"fn": fn_name, "args": args})

        # Warmup
        for _ in range(10):
            execute_function(request, fn_name, args)

        # Measure
        count = 0
        start = time.perf_counter()
        deadline = start + self.DURATION_SECONDS

        while time.perf_counter() < deadline:
            execute_function(request, fn_name, args)
            count += 1

        elapsed = time.perf_counter() - start
        return count / elapsed

    def _measure_throughput_http(self, fn_name: str, args: dict) -> float:
        """Measure requests/second using HTTP view calls."""
        # Warmup
        for _ in range(10):
            request = self._make_request({"fn": fn_name, "args": args})
            function_call_view(request)

        # Measure
        count = 0
        start = time.perf_counter()
        deadline = start + self.DURATION_SECONDS

        while time.perf_counter() < deadline:
            request = self._make_request({"fn": fn_name, "args": args})
            function_call_view(request)
            count += 1

        elapsed = time.perf_counter() - start
        return count / elapsed

    def _print_throughput(self, label: str, executor_rps: float, http_rps: float):
        """Print throughput results."""
        print(f"\n{label}:")
        print(f"  Executor (WebSocket): {executor_rps:,.0f} req/s")
        print(f"  HTTP:                 {http_rps:,.0f} req/s")
        print(f"  Ratio:                {executor_rps/http_rps:.1f}x")

    def test_throughput_simple(self):
        """Throughput: Simple computation (no I/O)."""
        print("\n\n### THROUGHPUT: Simple Computation ###")

        executor_rps = self._measure_throughput_executor(
            "bench_simple", {"a": 1, "b": 2}
        )
        http_rps = self._measure_throughput_http("bench_simple", {"a": 1, "b": 2})

        self._print_throughput("Simple (no I/O)", executor_rps, http_rps)

        # Correctness check
        request = self._make_request({"fn": "bench_simple", "args": {"a": 1, "b": 2}})
        result = execute_function(request, "bench_simple", {"a": 1, "b": 2})
        self.assertIsInstance(result, FunctionResult)
        self.assertEqual(result.data["value"], 3)

    def test_throughput_single_query(self):
        """Throughput: Single ORM query."""
        print("\n\n### THROUGHPUT: Single ORM Query ###")

        args = {"user_id": self.test_user.id if self.test_user else 1}

        executor_rps = self._measure_throughput_executor("bench_get_user", args)
        http_rps = self._measure_throughput_http("bench_get_user", args)

        self._print_throughput("Single Query", executor_rps, http_rps)

        # Correctness check
        request = self._make_request({"fn": "bench_get_user", "args": args})
        result = execute_function(request, "bench_get_user", args)
        self.assertIsInstance(result, FunctionResult)
        self.assertEqual(result.data["id"], self.test_user.id)

    def test_throughput_list_query(self):
        """Throughput: List query."""
        print("\n\n### THROUGHPUT: List Query (10 users) ###")

        executor_rps = self._measure_throughput_executor(
            "bench_list_users", {"limit": 10}
        )
        http_rps = self._measure_throughput_http("bench_list_users", {"limit": 10})

        self._print_throughput("List Query", executor_rps, http_rps)

        # Correctness check
        request = self._make_request({"fn": "bench_list_users", "args": {"limit": 10}})
        result = execute_function(request, "bench_list_users", {"limit": 10})
        self.assertIsInstance(result, FunctionResult)
        self.assertLessEqual(result.data["count"], 10)

    def test_throughput_aggregation(self):
        """Throughput: Aggregation queries."""
        print("\n\n### THROUGHPUT: Aggregation ###")

        executor_rps = self._measure_throughput_executor("bench_user_stats", {})
        http_rps = self._measure_throughput_http("bench_user_stats", {})

        self._print_throughput("Aggregation", executor_rps, http_rps)

        # Correctness check
        request = self._make_request({"fn": "bench_user_stats", "args": {}})
        result = execute_function(request, "bench_user_stats", {})
        self.assertIsInstance(result, FunctionResult)
        self.assertGreaterEqual(result.data["total_users"], 0)

    def test_throughput_summary(self):
        """Print throughput summary."""
        print("\n\n" + "=" * 80)
        print("THROUGHPUT SUMMARY")
        print("=" * 80)
        print(f"Test duration: {self.DURATION_SECONDS}s per scenario")
        print("\nNotes:")
        print("- These are single-threaded sequential measurements")
        print("- Real throughput scales with worker processes (gunicorn -w N)")
        print("- Database queries are the bottleneck, not protocol overhead")
        print("- Async workers (uvicorn) can handle more concurrent connections")
        print("=" * 80)

        # Verify bench_simple still produces correct output after all throughput tests
        request = self._make_request({"fn": "bench_simple", "args": {"a": 10, "b": 20}})
        result = execute_function(request, "bench_simple", {"a": 10, "b": 20})
        self.assertIsInstance(result, FunctionResult)
        self.assertEqual(result.data["value"], 30)
