#!/usr/bin/env python3
"""tests/reboot_probe.py — тест рестарта пода с in-flight сообщениями (план Соннета P0).

Фаза 1 (phase1): 10 очередей x 20 сообщений (проверка восстановления из Redis)
+ основная очередь со 100 сообщениями, из которых 50 выдаются (in-flight).
Фаза 2 (phase2, ПОСЛЕ kubectl delete pod --grace-period=0):
- /health поднимается;
- во всех 10 очередях должны быть ровно по 20 сообщений;
- из основной должны выдаться все 100 уникальных тел (50 in-flight + 50 видимых).
"""
import json
import os
import sys
import time
import uuid

import boto3
import urllib.request
from botocore.config import Config

ENDPOINT = os.environ.get("ENDPOINT_URL", "https://sqs.containerk8s.dev.nubes.ru")
REGION = os.environ.get("REGION", "us-east-1")
STATE = "/tmp/reboot_state.json"

sqs = boto3.client("sqs", endpoint_url=ENDPOINT, region_name=REGION,
                   config=Config(connect_timeout=10, read_timeout=30, retries={"max_attempts": 3}))


def wait_health(timeout=180):
    t0 = time.time()
    while time.time() - t0 < timeout:
        try:
            with urllib.request.urlopen(ENDPOINT + "/health", timeout=10) as r:
                if r.status == 200:
                    return True
        except Exception:
            pass
        time.sleep(3)
    return False


def drain(url):
    got = []
    while True:
        msgs = sqs.receive_message(QueueUrl=url, MaxNumberOfMessages=10).get("Messages", [])
        if not msgs:
            break
        got.extend(msgs)
    return got


def phase1():
    uid = str(uuid.uuid4())[:8]
    state = {"uid": uid, "restore_queues": [], "main_queue": None, "bodies": []}

    # 10 очередей по 20 сообщений — проверка восстановления из Redis
    for i in range(10):
        u = sqs.create_queue(QueueName="reboot-restore-%d-%s" % (i, uid))["QueueUrl"]
        for j in range(20):
            sqs.send_message(QueueUrl=u, MessageBody="r-%d-%d-%s" % (i, j, uid))
        state["restore_queues"].append(u)
        print("restore queue %d ready" % i, flush=True)

    # основная: 100 сообщений, 50 выдаём (in-flight)
    u = sqs.create_queue(QueueName="reboot-main-%s" % uid)["QueueUrl"]
    bodies = ["m-%03d-%s" % (i, uid) for i in range(100)]
    for b in bodies:
        sqs.send_message(QueueUrl=u, MessageBody=b)
    first = sqs.receive_message(QueueUrl=u, MaxNumberOfMessages=10, WaitTimeSeconds=1).get("Messages", [])
    # добираем до 50 in-flight (Max=10 за вызов)
    while len(first) < 50:
        batch = sqs.receive_message(QueueUrl=u, MaxNumberOfMessages=10, WaitTimeSeconds=1).get("Messages", [])
        if not batch:
            break
        first.extend(batch)
    print("in-flight на момент рестарта: %d" % len(first), flush=True)
    state["main_queue"] = u
    state["bodies"] = bodies
    json.dump(state, open(STATE, "w"))
    print("PHASE1_DONE", flush=True)


def phase2():
    st = json.load(open(STATE))
    print("ждём /health после рестарта...", flush=True)
    if not wait_health():
        print("HEALTH_FAIL: /health не поднялся", flush=True)
        sys.exit(1)
    print("health ok", flush=True)

    # восстановление 10 очередей
    restore_ok = True
    for u in st["restore_queues"]:
        msgs = drain(u)
        if len(msgs) != 20:
            restore_ok = False
            print("RESTORE FAIL: %s -> %d сообщений" % (u.split("/")[-1], len(msgs)), flush=True)
    print("restore: %s" % ("OK (10x20)" if restore_ok else "FAIL"), flush=True)

    # основная: все 100 уникальных
    got = drain(st["main_queue"])
    got_bodies = {m["Body"] for m in got}
    want = set(st["bodies"])
    missing = want - got_bodies
    dups = len(got) - len(got_bodies)
    print("main: получено %d, уникальных %d, пропущено %d, дублей %d"
          % (len(got), len(got_bodies), len(missing), dups), flush=True)
    print("MAIN_RESULT: %s" % ("PASS" if not missing else "FAIL (потеря!)"), flush=True)

    # cleanup
    for u in st["restore_queues"] + [st["main_queue"]]:
        try:
            sqs.delete_queue(QueueUrl=u)
        except Exception:
            pass
    print("PHASE2_DONE", flush=True)


if __name__ == "__main__":
    if len(sys.argv) != 2 or sys.argv[1] not in ("phase1", "phase2"):
        print("usage: reboot_probe.py phase1|phase2")
        sys.exit(2)
    if sys.argv[1] == "phase1":
        phase1()
    else:
        phase2()