"""Time/random/datetime/uuid/secrets replay patches.

Controls ALL sources of non-determinism during replay:
- time.time(), time.monotonic(), time.perf_counter()
- datetime.datetime.now(), datetime.datetime.utcnow()
- random.random(), randint(), randrange(), choice(), sample(), shuffle(), etc.
- uuid.uuid4(), uuid.uuid1()
- os.urandom(), secrets.token_bytes/hex/urlsafe
"""

from __future__ import annotations

import datetime as _dt_module
import asyncio as _asyncio
import os
import random
import time
import uuid as _uuid_module
from typing import Any, Callable, Dict, Optional

from autodebug_replay.transcripts import TranscriptIndex

_FALLBACK_TIME = 1700000000.0  # fallback epoch when transcript exhausted
_fallback_counter = 0

# Store all originals for uninstall
_originals: Dict[str, Any] = {}


def _reset():
    """Reset module state between replay sessions."""
    global _fallback_counter
    _fallback_counter = 0


# ---------------------------------------------------------------------------
# Time wrappers
# ---------------------------------------------------------------------------

def _next_time(index: TranscriptIndex) -> float:
    """Get next time value from transcript or fallback."""
    global _fallback_counter
    if not index.has_more("TIME"):
        _fallback_counter += 1
        return _FALLBACK_TIME + _fallback_counter
    record = index.expect("TIME")
    return float(record.get("value", _FALLBACK_TIME))


def _make_time_wrapper(index: TranscriptIndex) -> Callable[[], float]:
    def wrapper() -> float:
        return _next_time(index)
    return wrapper


def _make_replay_datetime_class(index: TranscriptIndex):
    """Create a datetime.datetime subclass that replays now()/utcnow().

    datetime.datetime is a C type — can't set attributes directly.
    Instead, replace with a Python subclass (same approach as freezegun).
    """
    _orig_datetime = _dt_module.datetime

    class _ReplayDatetime(_orig_datetime):
        @classmethod
        def now(cls, tz=None):
            ts = _next_time(index)
            base = _orig_datetime.fromtimestamp(ts, tz=tz)
            return cls(base.year, base.month, base.day,
                       base.hour, base.minute, base.second,
                       base.microsecond, base.tzinfo)

        @classmethod
        def utcnow(cls):
            ts = _next_time(index)
            base = _orig_datetime.utcfromtimestamp(ts)
            return cls(base.year, base.month, base.day,
                       base.hour, base.minute, base.second,
                       base.microsecond)

    return _ReplayDatetime


# ---------------------------------------------------------------------------
# RNG wrappers — type-specific to preserve return type contracts
# ---------------------------------------------------------------------------

def _next_rng_value(index: TranscriptIndex, fallback: Any = 0.5) -> Any:
    """Get next RNG value from transcript or fallback."""
    if not index.has_more("RNG"):
        return fallback
    record = index.expect("RNG")
    return record.get("value", fallback)


def _make_rng_float(index: TranscriptIndex) -> Callable:
    """For random.random(), random.uniform() — returns float."""
    def wrapper(*args, **kwargs) -> float:
        return float(_next_rng_value(index, 0.5))
    return wrapper


def _make_rng_int(index: TranscriptIndex) -> Callable:
    """For random.randint(), random.randrange() — returns int."""
    def wrapper(*args, **kwargs) -> int:
        val = _next_rng_value(index, args[0] if args else 0)
        return int(val)
    return wrapper


def _make_rng_choice(index: TranscriptIndex) -> Callable:
    """For random.choice() — returns element from sequence."""
    def wrapper(seq, *args, **kwargs):
        val = _next_rng_value(index, 0)
        if isinstance(val, int) and seq:
            return seq[val % len(seq)]
        # If transcript stored the actual chosen value, return it
        return val if val != 0.5 else (seq[0] if seq else None)
    return wrapper


def _make_rng_choices(index: TranscriptIndex) -> Callable:
    """For random.choices() — returns list of elements."""
    def wrapper(population, weights=None, *, cum_weights=None, k=1):
        val = _next_rng_value(index, None)
        if isinstance(val, list):
            return val
        return list(population)[:k]
    return wrapper


def _make_rng_sample(index: TranscriptIndex) -> Callable:
    """For random.sample() — returns list."""
    def wrapper(population, k, *args, **kwargs):
        val = _next_rng_value(index, None)
        if isinstance(val, list):
            return val
        return list(population)[:k]
    return wrapper


def _make_rng_shuffle(index: TranscriptIndex) -> Callable:
    """For random.shuffle() — mutates list in place, returns None."""
    def wrapper(x, *args, **kwargs):
        val = _next_rng_value(index, None)
        if isinstance(val, list) and len(val) == len(x):
            x[:] = val
        return None
    return wrapper


def _make_rng_gauss(index: TranscriptIndex) -> Callable:
    """For random.gauss(), random.normalvariate() — returns float."""
    def wrapper(*args, **kwargs) -> float:
        return float(_next_rng_value(index, 0.0))
    return wrapper


# ---------------------------------------------------------------------------
# UUID wrappers
# ---------------------------------------------------------------------------

def _make_uuid4(index: TranscriptIndex) -> Callable:
    _counter = [0]

    def wrapper() -> _uuid_module.UUID:
        if index.has_more("UUID"):
            record = index.expect("UUID")
            return _uuid_module.UUID(str(record.get("value")))
        # Deterministic fallback
        _counter[0] += 1
        return _uuid_module.UUID(int=_counter[0])
    return wrapper


def _make_uuid1(index: TranscriptIndex) -> Callable:
    _counter = [0]

    def wrapper(node=None, clock_seq=None) -> _uuid_module.UUID:
        if index.has_more("UUID"):
            record = index.expect("UUID")
            return _uuid_module.UUID(str(record.get("value")))
        _counter[0] += 1
        return _uuid_module.UUID(int=_counter[0])
    return wrapper


# ---------------------------------------------------------------------------
# os.urandom / secrets wrappers
# ---------------------------------------------------------------------------

def _make_urandom(index: TranscriptIndex) -> Callable:
    _counter = [0]

    def wrapper(n: int) -> bytes:
        if index.has_more("RNG"):
            record = index.expect("RNG")
            val = record.get("value")
            if isinstance(val, str):
                try:
                    return bytes.fromhex(val)[:n]
                except ValueError:
                    pass
        # Deterministic fallback
        _counter[0] += 1
        return (_counter[0]).to_bytes(max(n, 8), "big")[-n:]
    return wrapper


def _make_secrets_token_bytes(index: TranscriptIndex) -> Callable:
    _counter = [0]

    def wrapper(nbytes: Optional[int] = None) -> bytes:
        if nbytes is None:
            nbytes = 32
        if index.has_more("RNG"):
            record = index.expect("RNG")
            val = record.get("value")
            if isinstance(val, str):
                try:
                    return bytes.fromhex(val)[:nbytes]
                except ValueError:
                    pass
        _counter[0] += 1
        return (_counter[0]).to_bytes(max(nbytes, 8), "big")[-nbytes:]
    return wrapper


def _make_secrets_token_hex(index: TranscriptIndex) -> Callable:
    token_bytes_fn = _make_secrets_token_bytes(index)

    def wrapper(nbytes: Optional[int] = None) -> str:
        return token_bytes_fn(nbytes).hex()
    return wrapper


def _make_secrets_token_urlsafe(index: TranscriptIndex) -> Callable:
    import base64
    token_bytes_fn = _make_secrets_token_bytes(index)

    def wrapper(nbytes: Optional[int] = None) -> str:
        return base64.urlsafe_b64encode(token_bytes_fn(nbytes)).decode("ascii").rstrip("=")
    return wrapper


# ---------------------------------------------------------------------------
# Install / Uninstall
# ---------------------------------------------------------------------------

def install(index: TranscriptIndex) -> None:
    """Install all determinism patches for replay."""
    global _originals
    _reset()

    # Save originals
    _originals = {
        "time.time": time.time,
        "time.monotonic": time.monotonic,
        "time.perf_counter": time.perf_counter,
        "datetime.datetime": _dt_module.datetime,
        "random.random": random.random,
        "random.randint": random.randint,
        "random.randrange": random.randrange,
        "random.choice": random.choice,
        "random.choices": random.choices,
        "random.sample": random.sample,
        "random.shuffle": random.shuffle,
        "random.uniform": random.uniform,
        "random.gauss": random.gauss,
        "random.normalvariate": random.normalvariate,
        "uuid.uuid4": _uuid_module.uuid4,
        "uuid.uuid1": _uuid_module.uuid1,
        "os.urandom": os.urandom,
        "time.sleep": time.sleep,
    }

    # Time patches
    time.time = _make_time_wrapper(index)  # type: ignore[assignment]
    time.monotonic = _make_time_wrapper(index)  # type: ignore[assignment]
    time.perf_counter = _make_time_wrapper(index)  # type: ignore[assignment]
    time.sleep = lambda _seconds=0: None  # type: ignore[assignment]

    # Datetime patches — replace datetime.datetime with subclass
    _dt_module.datetime = _make_replay_datetime_class(index)  # type: ignore[misc]

    # RNG patches — type-specific wrappers
    random.random = _make_rng_float(index)  # type: ignore[assignment]
    random.uniform = _make_rng_float(index)  # type: ignore[assignment]
    random.randint = _make_rng_int(index)  # type: ignore[assignment]
    random.randrange = _make_rng_int(index)  # type: ignore[assignment]
    random.choice = _make_rng_choice(index)  # type: ignore[assignment]
    random.choices = _make_rng_choices(index)  # type: ignore[assignment]
    random.sample = _make_rng_sample(index)  # type: ignore[assignment]
    random.shuffle = _make_rng_shuffle(index)  # type: ignore[assignment]
    random.gauss = _make_rng_gauss(index)  # type: ignore[assignment]
    random.normalvariate = _make_rng_gauss(index)  # type: ignore[assignment]

    # UUID patches
    _uuid_module.uuid4 = _make_uuid4(index)  # type: ignore[assignment]
    _uuid_module.uuid1 = _make_uuid1(index)  # type: ignore[assignment]

    # os.urandom patch
    os.urandom = _make_urandom(index)  # type: ignore[assignment]

    # secrets patches (optional — only if secrets is importable)
    try:
        import secrets
        _originals["secrets.token_bytes"] = secrets.token_bytes
        _originals["secrets.token_hex"] = secrets.token_hex
        _originals["secrets.token_urlsafe"] = secrets.token_urlsafe
        secrets.token_bytes = _make_secrets_token_bytes(index)  # type: ignore[assignment]
        secrets.token_hex = _make_secrets_token_hex(index)  # type: ignore[assignment]
        secrets.token_urlsafe = _make_secrets_token_urlsafe(index)  # type: ignore[assignment]
    except ImportError:
        pass

    try:
        _originals["asyncio.sleep"] = _asyncio.sleep

        async def _replay_async_sleep(_delay=0, result=None):
            await _originals["asyncio.sleep"](0)
            return result

        _asyncio.sleep = _replay_async_sleep  # type: ignore[assignment]
    except Exception:
        pass


def uninstall() -> None:
    """Restore all original functions."""
    global _originals

    for key, original in _originals.items():
        parts = key.rsplit(".", 1)
        if len(parts) != 2:
            continue
        module_path, attr_name = parts

        if module_path == "time":
            setattr(time, attr_name, original)
        elif module_path == "random":
            setattr(random, attr_name, original)
        elif key == "datetime.datetime":
            _dt_module.datetime = original  # type: ignore[misc]
            continue
        elif module_path == "uuid":
            setattr(_uuid_module, attr_name, original)
        elif module_path == "os":
            setattr(os, attr_name, original)
        elif module_path == "asyncio":
            setattr(_asyncio, attr_name, original)
        elif module_path == "secrets":
            try:
                import secrets
                setattr(secrets, attr_name, original)
            except ImportError:
                pass

    _originals.clear()
    _reset()
