#!/usr/bin/env python3
"""Сравнение shared-SQS vs Yandex YMQ (очередь newsqs).

Короткий замер: последовательные операции, N send → N receive+delete.
Метрики: p50/p95/max латентности send/receive/delete, число ошибок.

Требуемые переменные окружения:
  SHARED_AK, SHARED_SK          — креды shared-SQS
  YMQ_AK, YMQ_SK                — креды Yandex YMQ (SA newsqs)
  YMQ_QUEUE_URL                 — URL очереди Яндекса (newsqs)

Запуск: python3 tests/compare_ymq_vs_shared.py [N]
"""
import os
import sys
import time
import uuid

import boto3
from botocore.config import Config

N = int(sys.argv[1]) if len(sys.argv) > 1 else 50
BODY = "x" * 512

SHARED_EP = "https://sqs.containerk8s.dev.nubes.ru"
YMQ_EP = "https://message-queue.api.cloud.yandex.net"

SHARED_AK = os.environ["SHARED_AK"]
SHARED_SK = os.environ["SHARED_SK"]
YMQ_AK = os.environ["YMQ_AK"]
YMQ_SK = os.environ["YMQ_SK"]
YMQ_QUEUE_URL = os.environ["YMQ_QUEUE_URL"]

CFG = Config(connect_timeout=15, read_timeout=30, retries={"max_attempts": 0})


def mk_client(endpoint, region, ak, sk):
    return boto3.client(
        "sqs",
        endpoint_url=endpoint,
        region_name=region,
        aws_access_key_id=ak,
        aws_secret_access_key=sk,
        config=CFG,
    )


def stats(name, values):
    v = sorted(values)
    if not v:
        print("%s: n=0" % name)
        return
    n = len(v)
    print("%s: n=%d p50=%.0fms p95=%.0fms max=%.0fms" %
          (name, n, v[n // 2] * 1000, v[int(n * .95)] * 1000, v[-1] * 1000))


def bench(client, queue_url, label):
    errs = 0
    lat_send, lat_recv, lat_del = [], [], []
    t_start = time.time()

    for i in range(N):
        t0 = time.time()
        try:
            client.send_message(QueueUrl=queue_url, MessageBody=BODY)
            lat_send.append(time.time() - t0)
        except Exception as e:
            errs += 1
            if errs <= 3:
                print("  send err: %s %s" % (type(e).__name__, str(e)[:80]))

    got = 0
    tries = 0
    while got < N and tries < N * 6:
        tries += 1
        t0 = time.time()
        try:
            r = client.receive_message(
                QueueUrl=queue_url, MaxNumberOfMessages=1, VisibilityTimeout=30)
            lat_recv.append(time.time() - t0)
            msgs = r.get("Messages", [])
        except Exception as e:
            errs += 1
            if errs <= 3:
                print("  recv err: %s %s" % (type(e).__name__, str(e)[:80]))
            continue
        for m in msgs:
            got += 1
            t0 = time.time()
            try:
                client.delete_message(QueueUrl=queue_url, ReceiptHandle=m["ReceiptHandle"])
                lat_del.append(time.time() - t0)
            except Exception as e:
                errs += 1
                if errs <= 3:
                    print("  del err: %s %s" % (type(e).__name__, str(e)[:80]))

    total = time.time() - t_start
    print("== %s (N=%d) ==" % (label, N))
    stats("  send   ", lat_send)
    stats("  receive", lat_recv)
    stats("  delete ", lat_del)
    print("  received=%d/%d errors=%d total=%.1fs" % (got, N, errs, total))
    print("  throughput: send=%.1f op/s, recv+del=%.1f op/s" %
          (N / total if total else 0, got / total if total else 0))
    return errs


def drain(client, queue_url):
    while True:
        r = client.receive_message(QueueUrl=queue_url, MaxNumberOfMessages=10)
        msgs = r.get("Messages", [])
        if not msgs:
            break
        for m in msgs:
            try:
                client.delete_message(QueueUrl=queue_url, ReceiptHandle=m["ReceiptHandle"])
            except Exception:
                pass


def main():
    shared = mk_client(SHARED_EP, "us-east-1", SHARED_AK, SHARED_SK)
    ymq = mk_client(YMQ_EP, "ru-central1", YMQ_AK, YMQ_SK)

    # Временная очередь shared-SQS
    qname = "cmp-%s" % uuid.uuid4().hex[:8]
    shared_url = shared.create_queue(QueueName=qname)["QueueUrl"]
    print("shared queue: %s" % shared_url)

    # YMQ newsqs: перед тестом чистим
    drain(ymq, YMQ_QUEUE_URL)
    print("ymq queue: %s (drained)" % YMQ_QUEUE_URL)

    print()
    e1 = bench(shared, shared_url, "shared-SQS (наш сервис)")
    print()
    e2 = bench(ymq, YMQ_QUEUE_URL, "Yandex YMQ (newsqs)")

    # Очистка
    drain(ymq, YMQ_QUEUE_URL)
    try:
        shared.delete_queue(QueueUrl=shared_url)
        print("\ncleanup: shared temp queue deleted, ymq drained")
    except Exception as ex:
        print("\ncleanup err: %s" % str(ex)[:80])

    print("\nRESULT: errors shared=%d ymq=%d" % (e1, e2))


if __name__ == "__main__":
    main()