# requires: grpc, grpc_tools
"""gRPC 演示：从 user.proto 生成代码 → 启动服务端 → 客户端调用（一元、服务端流、超时、错误码）。

运行：pip install grpcio grpcio-tools && python demo.py
生成的代码放在临时目录，不污染仓库。
"""
import importlib
import sys
import tempfile
import time
from concurrent import futures
from pathlib import Path

import grpc
from grpc_tools import protoc

HERE = Path(__file__).resolve().parent


def generate(out: Path):
    rc = protoc.main(["protoc", f"-I{HERE}", f"--python_out={out}", f"--grpc_python_out={out}", str(HERE / "user.proto")])
    assert rc == 0, "protoc 生成失败"
    sys.path.insert(0, str(out))
    return importlib.import_module("user_pb2"), importlib.import_module("user_pb2_grpc")


def main() -> None:
    with tempfile.TemporaryDirectory() as tmp:
        pb, pb_grpc = generate(Path(tmp))
        users = {i: pb.User(id=i, name=f"用户{i}", email=f"u{i}@example.com") for i in range(1, 6)}

        class UserService(pb_grpc.UserServiceServicer):
            def GetUser(self, request, context):
                if request.id == 99:
                    time.sleep(0.5)  # 模拟慢请求
                user = users.get(request.id)
                if user is None:
                    context.abort(grpc.StatusCode.NOT_FOUND, f"用户 {request.id} 不存在")
                return user

            def ListUsers(self, request, context):
                for i, u in enumerate(users.values()):
                    if i >= request.page_size:
                        break
                    yield u

        server = grpc.server(futures.ThreadPoolExecutor(max_workers=4))
        pb_grpc.add_UserServiceServicer_to_server(UserService(), server)
        port = server.add_insecure_port("127.0.0.1:0")  # 0 = 随机可用端口
        server.start()
        try:
            with grpc.insecure_channel(f"127.0.0.1:{port}") as ch:
                stub = pb_grpc.UserServiceStub(ch)
                u = stub.GetUser(pb.GetUserRequest(id=1), timeout=1.0)  # 每次调用都带超时
                print("一元调用：", u.id, u.name)
                assert u.name == "用户1"

                names = [x.name for x in stub.ListUsers(pb.ListUsersRequest(page_size=3), timeout=1.0)]
                print("服务端流：", names)
                assert len(names) == 3

                try:
                    stub.GetUser(pb.GetUserRequest(id=42), timeout=1.0)
                    raise AssertionError("应该返回 NOT_FOUND")
                except grpc.RpcError as e:
                    print("错误码：", e.code().name, e.details())
                    assert e.code() == grpc.StatusCode.NOT_FOUND

                try:
                    stub.GetUser(pb.GetUserRequest(id=99), timeout=0.1)
                    raise AssertionError("应该超时")
                except grpc.RpcError as e:
                    print("超时：", e.code().name)
                    assert e.code() == grpc.StatusCode.DEADLINE_EXCEEDED
        finally:
            server.stop(grace=None)
    print("gRPC 演示通过")


if __name__ == "__main__":
    main()