#!/usr/bin/env python3
"""ДЛИННОЕ сравнение shared-SQS vs Yandex YMQ — с ЛОКАЛИ, по таймеру.

Схема (честная — обе очереди в одних сетевых условиях, раунды попеременно):
  1. Старт: seed K сообщений в каждую очередь.
  2. Раунд для каждой очереди: receive(1) -> delete -> send(1) (восполнение пула).
  3. Раунды чередуются (наш, яндекс, наш, яндекс...) до истечения DURATION.

Метрики: p50/p95/max для send/receive/delete, ошибки по типам (с таймаутами),
op/s. Прогресс каждые 60с пишется в LOG и stdout.

Env: SHARED_AK, SHARED_SK, YMQ_AK, YMQ_SK, YMQ_QUEUE_URL, DURATION (сек, default 1800)
"""
import os
import sys
import time
import uuid

import boto3
from botocore.config import Config

DURATION = int(os.environ.get("DURATION", "1800"))
K = 10
BODY = "x" * 512
LOG = "tests/long_compare.log"

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

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


def mk(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)


class Stat:
    def __init__(self):
        self.vals = {"send": [], "receive": [], "delete": []}
        self.errs = {}   # type -> count
        self.ops = {"send": 0, "receive": 0, "delete": 0}

    def op(self, name, fn):
        self.ops[name] += 1
        t0 = time.time()
        try:
            fn()
            self.vals[name].append(time.time() - t0)
        except Exception as e:
            k = type(e).__name__
            self.errs[k] = self.errs.get(k, 0) + 1


def stat_line(st, label):
    out = [label]
    for name in ("send", "receive", "delete"):
        v = sorted(st.vals[name])
        n = len(v)
        if n:
            out.append("%s: n=%d p50=%.0fms p95=%.0fms max=%.0fms" %
                       (name, n, v[n // 2] * 1000, v[int(n * .95)] * 1000, v[-1] * 1000))
        else:
            out.append("%s: n=0" % name)
    out.append("errors=%s" % (st.errs if st.errs else "0"))
    return " | ".join(out)


def logp(msg):
    line = "%s %s" % (time.strftime("%H:%M:%S"), msg)
    print(line, flush=True)
    try:
        with open(LOG, "a") as f:
            f.write(line + "\n")
    except Exception:
        pass


def main():
    shared = mk(SHARED_EP, "us-east-1", os.environ["SHARED_AK"], os.environ["SHARED_SK"])
    ymq = mk(YMQ_EP, "ru-central1", os.environ["YMQ_AK"], os.environ["YMQ_SK"])
    ymq_url = os.environ["YMQ_QUEUE_URL"]

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

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

    drain(shared, shared_url)
    drain(ymq, ymq_url)

    ss = Stat()
    ys = Stat()
    sts = {"shared": (shared, shared_url, ss), "ymq": (ymq, ymq_url, ys)}

    # seed K в обе очереди
    for key in ("shared", "ymq"):
        c, u, st = sts[key]
        for _ in range(K):
            st.op("send", lambda c=c, u=u: c.send_message(QueueUrl=u, MessageBody=BODY))

    logp("LONG COMPARE START: duration=%ds K=%d shared=%s" % (DURATION, K, shared_url))
    logp("ymq=%s" % ymq_url)

    deadline = time.time() + DURATION
    last_prog = time.time()
    rounds = 0

    def round_for(key):
        nonlocal rounds
        c, u, st = sts[key]
        # receive (замер) + получить сообщения для delete
        got = []
        t0 = time.time()
        st.ops["receive"] += 1
        try:
            r = c.receive_message(QueueUrl=u, MaxNumberOfMessages=1, VisibilityTimeout=30)
            st.vals["receive"].append(time.time() - t0)
            got = r.get("Messages", [])
        except Exception as e:
            st.errs[type(e).__name__] = st.errs.get(type(e).__name__, 0) + 1
        for m in got:
            st.op("delete", lambda c=c, u=u, m=m: c.delete_message(QueueUrl=u, ReceiptHandle=m["ReceiptHandle"]))
        st.op("send", lambda c=c, u=u: c.send_message(QueueUrl=u, MessageBody=BODY))

    while time.time() < deadline:
        round_for("shared")
        round_for("ymq")
        rounds += 1
        if time.time() - last_prog >= 60:
            last_prog = time.time()
            logp("progress %d/%ds" % (int(time.time() - (deadline - DURATION)), DURATION))
            logp("  " + stat_line(ss, "SHARED"))
            logp("  " + stat_line(ys, "YMQ"))

    logp("LONG COMPARE DONE: rounds=%d" % rounds)
    logp(stat_line(ss, "FINAL SHARED"))
    logp(stat_line(ys, "FINAL YMQ"))

    # cleanup
    drain(shared, shared_url)
    drain(ymq, ymq_url)
    try:
        shared.delete_queue(QueueUrl=shared_url)
        logp("cleanup: shared temp queue deleted, both drained")
    except Exception as e:
        logp("cleanup err: %s" % str(e)[:100])


if __name__ == "__main__":
    main()