"""Time and RNG patching for deterministic capture."""

from __future__ import annotations

import random
import time
from typing import Callable

from autodebug_probe.recorder import get_recorder

ORIGINAL_TIME = time.time
ORIGINAL_MONOTONIC = time.monotonic
ORIGINAL_PERF_COUNTER = time.perf_counter
ORIGINAL_SLEEP = time.sleep


def _wrap_time(func: Callable[[], float], name: str) -> Callable[[], float]:
    def wrapper() -> float:
        value = func()
        get_recorder().record_time_rng({"type": "TIME", "call": name, "value": value})
        return value

    return wrapper


def _wrap_sleep(func: Callable[[float], None], name: str) -> Callable[[float], None]:
    def wrapper(secs):
        func(secs)
        get_recorder().record_time_rng({"type": "SLEEP", "call": name, "value": secs})

    return wrapper


def _wrap_rng(func: Callable[..., float], name: str) -> Callable[..., float]:
    def wrapper(*args, **kwargs):
        value = func(*args, **kwargs)
        get_recorder().record_time_rng({"type": "RNG", "call": name, "value": value})
        return value

    return wrapper


def _wrap_shuffle():
    original = random.shuffle

    def wrapper(x, random_func=None):
        if random_func:
            original(x, random_func)
        else:
            original(x)
        get_recorder().record_time_rng({"type": "RNG", "call": "shuffle", "value": len(x)})

    return wrapper


def _wrap_rng_len(func: Callable, name: str) -> Callable:
    def wrapper(*args, **kwargs):
        value = func(*args, **kwargs)
        get_recorder().record_time_rng({"type": "RNG", "call": name, "value": len(value)})
        return value

    return wrapper


def _wrap_gauss(func: Callable) -> Callable:
    def wrapper(mu=0.0, sigma=1.0):
        value = func(mu, sigma)
        get_recorder().record_time_rng(
            {"type": "RNG", "call": "gauss", "value": value, "mu": mu, "sigma": sigma}
        )
        return value

    return wrapper


_installed = False


def install() -> None:
    # Thread-safety note: random module itself is not thread-safe (per CPython docs),
    # so this patching inherits the same guarantees. The _installed guard prevents
    # double-wrapping if install() is called from multiple threads.
    global _installed
    if _installed:
        return
    _installed = True
    time.time = _wrap_time(ORIGINAL_TIME, "time")  # type: ignore[assignment]
    time.monotonic = _wrap_time(ORIGINAL_MONOTONIC, "monotonic")  # type: ignore[assignment]
    time.perf_counter = _wrap_time(ORIGINAL_PERF_COUNTER, "perf_counter")  # type: ignore[assignment]

    random.random = _wrap_rng(random.random, "random")  # type: ignore[assignment]
    random.randint = _wrap_rng(random.randint, "randint")  # type: ignore[assignment]
    random.randrange = _wrap_rng(random.randrange, "randrange")  # type: ignore[assignment]
    random.choice = _wrap_rng(random.choice, "choice")  # type: ignore[assignment]
    random.uniform = _wrap_rng(random.uniform, "uniform")  # type: ignore[assignment]

    time.sleep = _wrap_sleep(ORIGINAL_SLEEP, "sleep")  # type: ignore[assignment]

    random.shuffle = _wrap_shuffle()  # type: ignore[assignment]
    random.sample = _wrap_rng_len(random.sample, "sample")  # type: ignore[assignment]
    random.choices = _wrap_rng_len(random.choices, "choices")  # type: ignore[assignment]
    random.gauss = _wrap_gauss(random.gauss)  # type: ignore[assignment]
