import grpc
import time

from grpc_util.grpc_sam3 import grpc_sam3_img_pb2_grpc, grpc_sam3_img_pb2


def check_server_status(channel):
    try:
        health_stub = grpc_sam3_img_pb2_grpc.HealthCheckStub(channel)
        response = health_stub.Check(grpc_sam3_img_pb2.HealthCheckRequest(service="TaskService"))
        return response.status == grpc_sam3_img_pb2.HealthCheckResponse.ServingStatus.SERVING
    except grpc.RpcError as e:
        print(f"Health check failed: {e}")
        return False


def check_grpc_request(max_retries=3, delay=5):
    channel = None
    retries = 0

    while retries < max_retries:
        try:
            # 创建通道
            channel = grpc.insecure_channel('192.168.110.187:9999')

            # 检查服务器状态
            if not check_server_status(channel):
                raise Exception("Server is not healthy")

            stub = grpc_sam3_img_pb2_grpc.TaskServiceStub(channel)

            # 创建请求消息
            request = grpc_sam3_img_pb2.TaskRequest(
                task_id="d6118954-a170-4e1c-84bd-ddbd3114b354111",
                sn="8UUXN6S00A0CK7",
                content_body=grpc_sam3_img_pb2.ContentBody(
                    img_url="demo/03.png",
                    prompt="cat",
                    confidence=0.5,
                    mqtt_ip="47.108.62.6",
                    mqtt_port=12503,
                    mqtt_topic="thing/product/ai/events"
                )
            )

            # 调用远程方法
            response = stub.ProcessTask(request)
            print(f"Response: task_id={response.task_id}, success={response.success}, message={response.message}")
            return True

        except grpc.RpcError as e:
            retries += 1
            print(f"RPC error occurred (attempt {retries}/{max_retries}): {e}")
            if retries < max_retries:
                print(f"Retrying in {delay} seconds...")
                time.sleep(delay)
        except Exception as e:
            print(f"Error occurred: {e}")
            retries += 1
            if retries < max_retries:
                print(f"Retrying in {delay} seconds...")
                time.sleep(delay)
        finally:
            if channel:
                channel.close()

    print("All retry attempts failed")
    return False


if __name__ == '__main__':
    check_grpc_request()