#!/usr/bin/env python3
"""
Reproduction script for _ConcurrentRpcLimiter bug in gRPC Python aio server.

Usage:
    pip install grpcio grpcio-tools
    python repro.py
"""

import asyncio
import logging
import sys
import os
import tempfile

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(message)s')
logger = logging.getLogger(__name__)
logging.getLogger('grpc').setLevel(logging.WARNING)

PROTO_CONTENT = '''
syntax = "proto3";
package test;
service TestService {
  rpc Test (TestRequest) returns (TestResponse) {}
}
message TestRequest { string data = 1; }
message TestResponse { string result = 1; }
'''

def setup_proto():
    import subprocess
    proto_dir = tempfile.mkdtemp()
    proto_file = os.path.join(proto_dir, "test.proto")
    with open(proto_file, "w") as f:
        f.write(PROTO_CONTENT)
    subprocess.run([
        sys.executable, "-m", "grpc_tools.protoc",
        f"-I{proto_dir}", f"--python_out={proto_dir}",
        f"--grpc_python_out={proto_dir}", proto_file
    ], check=True)
    sys.path.insert(0, proto_dir)
    return proto_dir

async def main():
    import grpc
    from grpc import aio

    proto_dir = setup_proto()
    import test_pb2, test_pb2_grpc

    concurrent_count = 0
    max_concurrent_seen = 0
    lock = asyncio.Lock()

    class TestServicer(test_pb2_grpc.TestServiceServicer):
        async def Test(self, request, context):
            nonlocal concurrent_count, max_concurrent_seen
            async with lock:
                concurrent_count += 1
                if concurrent_count > max_concurrent_seen:
                    max_concurrent_seen = concurrent_count
            logger.info(f"[HANDLER] START {request.data}, concurrent={concurrent_count}")
            await asyncio.sleep(1.0)
            async with lock:
                concurrent_count -= 1
            logger.info(f"[HANDLER] END {request.data}")
            return test_pb2.TestResponse(result=f"ok-{request.data}")

    server = aio.server(maximum_concurrent_rpcs=1)
    test_pb2_grpc.add_TestServiceServicer_to_server(TestServicer(), server)
    port = server.add_insecure_port("127.0.0.1:0")
    await server.start()

    channel = aio.insecure_channel(f"127.0.0.1:{port}")
    stub = test_pb2_grpc.TestServiceStub(channel)
    await asyncio.sleep(0.5)

    async def send(id, wave):
        try:
            await stub.Test(test_pb2.TestRequest(data=f"w{wave}r{id}"), timeout=30)
            return "SUCCESS"
        except grpc.aio.AioRpcError as e:
            return e.code().name

    # Wave 1
    logger.info("WAVE 1: 3 concurrent requests")
    w1 = await asyncio.gather(*[send(i, 1) for i in range(3)])
    await asyncio.sleep(2)

    # Wave 2
    logger.info("WAVE 2: 3 more concurrent requests")
    w2 = await asyncio.gather(*[send(i, 2) for i in range(3)])
    await asyncio.sleep(2)

    logger.info(f"Wave 1: {w1.count('SUCCESS')} success, {w1.count('RESOURCE_EXHAUSTED')} exhausted")
    logger.info(f"Wave 2: {w2.count('SUCCESS')} success, {w2.count('RESOURCE_EXHAUSTED')} exhausted")
    logger.info(f"Max concurrent seen: {max_concurrent_seen}")

    if max_concurrent_seen > 1:
        logger.error("BUG CONFIRMED: max concurrent exceeded limit!")

    await server.stop(0)
    import shutil
    shutil.rmtree(proto_dir, ignore_errors=True)

if __name__ == "__main__":
    asyncio.run(main())