from argparse import ArgumentParser
from secrets import randbelow

from utils.check_case import check_case
from utils.format_line import format_line
from orchestration.param_generator import ParamGenerator


def _random_secret() -> int:
    # 2..9999 включительно (меньше 10000, исключаем 0/1)
    return 2 + randbelow(9998)


def run_auto(count: int) -> None:
    """
    Автоматический режим:
      - случайно генерируем (p, q, g);
      - случайно выбираем x, y (<10000);
      - вычисляем K и проверяем через check_case;
      - печатаем результат через format_line
    """
    for i in range(1, count + 1):
        p, q, g = ParamGenerator.generate()  # q выводим для информации
        x = _random_secret()
        y = _random_secret()

        # Вычислим K заранее, чтобы передать как эталон в check_case
        # (тем самым используем общий хелпер единообразно)
        ok, (A, B, K_a, K_b) = check_case(p, g, x, y, K_expected=None)
        # Когда K_expected=None, наша старая версия check_case
        # сравнивала строго с эталоном. Если у тебя именно так,
        # просто заменим на повторный вызов с «эталоном = K_a»
        if K_a != K_b:
            ok = False
        else:
            ok, (A, B, K_a, K_b) = check_case(p, g, x, y, K_expected=K_a)

        print(f"[auto #{i}] p={p}  q={q}  g={g}")
        print(format_line(i, A, B, K_a, K_b, ok))


def run_manual(p: int, g: int, x: int, y: int, k: int | None) -> None:
    """
    Ручной режим:
      - принимаем p, g, x, y;
      - если указан k — проверяем с эталоном;
      - если k не указан — просто считаем и показываем
    """
    if k is None:
        ok, (A, B, K_a, K_b) = check_case(p, g, x, y, K_expected=None)
        ok = (K_a == K_b)
        print(f"[manual] p={p}  g={g}")
        print(format_line(1, A, B, K_a, K_b, ok))
    else:
        ok, (A, B, K_a, K_b) = check_case(p, g, x, y, K_expected=k)
        print(f"[manual] p={p}  g={g}  K_expected={k}")
        print(format_line(1, A, B, K_a, K_b, ok))


def parse_args():
    ap = ArgumentParser(
        description=(
            "Demo Diffie–Hellman: режимы manual и auto\n"
            "Auto: случайно генерирует (p,q,g) и секреты (<10000)\n"
            "Manual: используется заданный набор параметров"
        )
    )
    ap.add_argument(
        "--mode",
        choices=["auto", "manual"],
        default="auto",
        help="Режим работы (auto|manual). По умолчанию auto",
    )
    ap.add_argument(
        "--count",
        type=int,
        default=5,
        help="Сколько генераций в auto-режиме (по умолчанию 5)",
    )
    ap.add_argument("--p", type=int, help="Модуль (prime_modulus) в manual")
    ap.add_argument("--g", type=int, help="Генератор (generator) в manual")
    ap.add_argument("--x", type=int, help="Секрет Алисы (<10000) в manual")
    ap.add_argument("--y", type=int, help="Секрет Боба (<10000) в manual")
    ap.add_argument(
        "--k",
        type=int,
        help="Эталонный общий ключ для проверки в manual (необязательно)",
    )
    return ap.parse_args()


def main():
    args = parse_args()
    if args.mode == "auto":
        run_auto(count=args.count)
        return

    # manual
    missing = [n for n, v in [("p", args.p), ("g", args.g),
                              ("x", args.x), ("y", args.y)] if v is None]
    if missing:
        need = ", ".join(missing)
        raise SystemExit(f"manual-режим: необходимо указать параметры: {need}")
    run_manual(p=args.p, g=args.g, x=args.x, y=args.y, k=args.k)


if __name__ == "__main__":
    main()