#!/usr/bin/env python3
"""Redis-backed rate limiting validation helper.

This script is a small operational aid for teams evaluating rate limiting policy
before production rollout. It does not implement your production limiter. Instead,
it helps you:

- simulate request bursts against a Redis key strategy,
- inspect how counters or token-bucket style state evolves,
- confirm whether a proposed threshold would allow or reject traffic,
- optionally write a lightweight, time-boxed test key into Redis.

Supported modes:

1. fixed-window: increments a counter for a given identity key and checks it
   against a request limit within a time window.
2. token-bucket: tracks a refillable token pool in Redis.

Safety notes:
- No credentials are hard-coded.
- Redis connection settings must be provided via CLI flags.
- The script is read-friendly and avoids destructive defaults.
- Use it in a non-production environment first.

Examples:

  # Check fixed-window behavior for a user key
  python rate_limit_validate.py fixed-window \
      --redis-host localhost --redis-port 6379 \
      --key rate:user:123 --limit 100 --window-seconds 60 \
      --requests 120

  # Check token-bucket behavior for a route+token key
  python rate_limit_validate.py token-bucket \
      --redis-host localhost --redis-port 6379 \
      --key rate:token:abc123:/v1/export \
      --capacity 20 --refill-rate 2.0 --requests 35

  # Dry-run without writing to Redis
  python rate_limit_validate.py fixed-window \
      --key rate:ip:203.0.113.10 --limit 20 --window-seconds 10 \
      --requests 25 --dry-run
"""

from __future__ import annotations

import argparse
import json
import sys
import time
from dataclasses import dataclass
from typing import Any, Dict, Optional


try:
    import redis  # type: ignore
except ImportError:
    redis = None


@dataclass
class RedisConfig:
    host: str
    port: int
    db: int
    password: Optional[str]
    socket_timeout: float


def positive_int(value: str) -> int:
    try:
        parsed = int(value)
    except ValueError as exc:
        raise argparse.ArgumentTypeError(f"invalid integer value: {value}") from exc
    if parsed < 0:
        raise argparse.ArgumentTypeError("value must be >= 0")
    return parsed


def positive_float(value: str) -> float:
    try:
        parsed = float(value)
    except ValueError as exc:
        raise argparse.ArgumentTypeError(f"invalid numeric value: {value}") from exc
    if parsed < 0:
        raise argparse.ArgumentTypeError("value must be >= 0")
    return parsed


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Validate Redis-backed rate limiting behavior for Node.js APIs."
    )
    subparsers = parser.add_subparsers(dest="mode", required=True)

    common = argparse.ArgumentParser(add_help=False)
    common.add_argument("--redis-host", default="127.0.0.1", help="Redis host")
    common.add_argument("--redis-port", type=positive_int, default=6379, help="Redis port")
    common.add_argument("--redis-db", type=positive_int, default=0, help="Redis database number")
    common.add_argument("--redis-password", default=None, help="Redis password (optional)")
    common.add_argument(
        "--socket-timeout",
        type=positive_float,
        default=2.0,
        help="Redis socket timeout in seconds",
    )
    common.add_argument(
        "--key",
        required=True,
        help="Rate-limit identity key, for example rate:user:123 or rate:ip:203.0.113.10",
    )
    common.add_argument(
        "--dry-run",
        action="store_true",
        help="Simulate decisions without writing changes to Redis",
    )
    common.add_argument(
        "--json-output",
        action="store_true",
        help="Print results as JSON",
    )

    fixed = subparsers.add_parser(
        "fixed-window",
        parents=[common],
        help="Validate fixed-window request counting",
    )
    fixed.add_argument("--limit", type=positive_int, required=True, help="Allowed requests per window")
    fixed.add_argument(
        "--window-seconds",
        type=positive_int,
        required=True,
        help="Window size in seconds",
    )
    fixed.add_argument(
        "--requests",
        type=positive_int,
        default=1,
        help="Number of requests to simulate",
    )
    fixed.add_argument(
        "--sleep-ms",
        type=positive_int,
        default=0,
        help="Optional delay between simulated requests",
    )

    token = subparsers.add_parser(
        "token-bucket",
        parents=[common],
        help="Validate token-bucket request allowance",
    )
    token.add_argument("--capacity", type=positive_int, required=True, help="Bucket capacity")
    token.add_argument(
        "--refill-rate",
        type=positive_float,
        required=True,
        help="Token refill rate per second",
    )
    token.add_argument(
        "--requests",
        type=positive_int,
        default=1,
        help="Number of requests to simulate",
    )
    token.add_argument(
        "--sleep-ms",
        type=positive_int,
        default=0,
        help="Optional delay between simulated requests",
    )

    return parser


def make_client(cfg: RedisConfig):
    if redis is None:
        raise RuntimeError(
            "redis package is not installed. Install it with: pip install redis"
        )
    return redis.Redis(
        host=cfg.host,
        port=cfg.port,
        db=cfg.db,
        password=cfg.password,
        socket_timeout=cfg.socket_timeout,
        decode_responses=True,
    )


def fixed_window_simulation(client, key: str, limit: int, window_seconds: int, requests: int, dry_run: bool, sleep_ms: int) -> Dict[str, Any]:
    results = []
    allowed = 0
    rejected = 0
    ttl_set = False

    for index in range(requests):
        if sleep_ms:
            time.sleep(sleep_ms / 1000.0)

        if dry_run:
            current = index + 1
            ttl = window_seconds
        else:
            current = client.incr(key)
            if current == 1:
                client.expire(key, window_seconds)
                ttl_set = True
            ttl = client.ttl(key)

        decision = current <= limit
        if decision:
            allowed += 1
        else:
            rejected += 1

        results.append(
            {
                "request_number": index + 1,
                "counter": current,
                "allowed": decision,
                "ttl_seconds": ttl,
            }
        )

    return {
        "mode": "fixed-window",
        "key": key,
        "limit": limit,
        "window_seconds": window_seconds,
        "requests_simulated": requests,
        "allowed": allowed,
        "rejected": rejected,
        "ttl_was_set": ttl_set,
        "events": results,
    }


def token_bucket_simulation(client, key: str, capacity: int, refill_rate: float, requests: int, dry_run: bool, sleep_ms: int) -> Dict[str, Any]:
    state_key = f"{key}:state"
    results = []
    allowed = 0
    rejected = 0

    now = time.time()
    tokens = float(capacity)
    last_refill = now

    for index in range(requests):
        if sleep_ms:
            time.sleep(sleep_ms / 1000.0)
        now = time.time()
        elapsed = max(0.0, now - last_refill)
        tokens = min(capacity, tokens + elapsed * refill_rate)
        last_refill = now

        if dry_run:
            current_tokens = tokens
        else:
            current_tokens = tokens
            payload = json.dumps({"tokens": current_tokens, "ts": last_refill})
            client.set(state_key, payload, ex=max(1, int((capacity / refill_rate) if refill_rate > 0 else 60)))

        decision = current_tokens >= 1.0
        if decision:
            tokens = current_tokens - 1.0
            allowed += 1
        else:
            rejected += 1

        results.append(
            {
                "request_number": index + 1,
                "tokens_before_request": round(current_tokens, 3),
                "allowed": decision,
                "tokens_after_request": round(tokens, 3),
            }
        )

    if not dry_run:
        client.set(state_key, json.dumps({"tokens": round(tokens, 3), "ts": last_refill}))

    return {
        "mode": "token-bucket",
        "key": key,
        "capacity": capacity,
        "refill_rate_per_second": refill_rate,
        "requests_simulated": requests,
        "allowed": allowed,
        "rejected": rejected,
        "state_key": state_key,
        "events": results,
    }


def main() -> int:
    parser = build_parser()
    args = parser.parse_args()

    cfg = RedisConfig(
        host=args.redis_host,
        port=args.redis_port,
        db=args.redis_db,
        password=args.redis_password,
        socket_timeout=args.socket_timeout,
    )

    try:
        client = None if args.dry_run else make_client(cfg)
        if client is not None:
            client.ping()

        if args.mode == "fixed-window":
            output = fixed_window_simulation(
                client=client,
                key=args.key,
                limit=args.limit,
                window_seconds=args.window_seconds,
                requests=args.requests,
                dry_run=args.dry_run,
                sleep_ms=args.sleep_ms,
            )
        elif args.mode == "token-bucket":
            output = token_bucket_simulation(
                client=client,
                key=args.key,
                capacity=args.capacity,
                refill_rate=args.refill_rate,
                requests=args.requests,
                dry_run=args.dry_run,
                sleep_ms=args.sleep_ms,
            )
        else:
            parser.error(f"unsupported mode: {args.mode}")
            return 2

        if args.json_output:
            print(json.dumps(output, indent=2, sort_keys=True))
        else:
            print(f"Mode: {output['mode']}")
            print(f"Key: {output['key']}")
            print(f"Allowed: {output['allowed']}")
            print(f"Rejected: {output['rejected']}")
            print("Events:")
            for event in output["events"]:
                print(json.dumps(event, sort_keys=True))

        return 0

    except KeyboardInterrupt:
        print("Interrupted.", file=sys.stderr)
        return 130
    except Exception as exc:
        print(f"error: {exc}", file=sys.stderr)
        return 1


if __name__ == "__main__":
    raise SystemExit(main())