#!/usr/bin/env python3
"""
Authorization gates for sensitive endpoints.

Controls access to metrics, receipts, and admin endpoints based on:
- FLOWTORCH_PUBLIC mode (0=private, 1=public)
- Token-based authentication for private deployments
- Per-IP rate limiting for admin endpoints (P6-P8 hardening)
"""
import os
import time
import json
import threading
from collections import defaultdict
from fastapi import Header, HTTPException, Request
from fastapi.responses import JSONResponse

# Use monotonic clock for rate limiting (immune to NTP jumps)
_now = time.monotonic

# Proxy trust configuration (P2)
import ipaddress

TRUST_PROXY = os.getenv("FLOWTORCH_TRUST_PROXY", "0") == "1"
TRUSTED_PROXY_HOPS = max(0, int(os.getenv("FLOWTORCH_TRUSTED_PROXY_HOPS", "1")))

# Helper to normalize IPs early (needed for TRUSTED_PROXIES parsing)
def _norm_ip_early(token: str) -> str | None:
    """Normalize IP address token, return None if invalid."""
    try:
        return str(ipaddress.ip_address(token))
    except Exception:
        return None

# Parse TRUSTED_PROXIES as CIDR blocks or single IPs
TRUSTED_PROXIES_RAW = [p.strip() for p in os.getenv("FLOWTORCH_TRUSTED_PROXIES", "").split(",") if p.strip()]
_TRUSTED_NETS = []
for p in TRUSTED_PROXIES_RAW:
    try:
        _TRUSTED_NETS.append(ipaddress.ip_network(p, strict=False))
    except Exception:
        # Try as single IP
        ni = _norm_ip_early(p)
        if ni:
            _TRUSTED_NETS.append(ipaddress.ip_network(ni + "/32" if ":" not in ni else ni + "/128", strict=False))


def _peer_trusted(ip: str) -> bool:
    """Check if IP is in trusted proxy list or localhost."""
    try:
        ipo = ipaddress.ip_address(ip)
        # Check localhost
        if ip in {"127.0.0.1", "::1"}:
            return True
        # Check trusted networks
        return any(ipo in n for n in _TRUSTED_NETS)
    except Exception:
        return False


# Public/private mode
PUBLIC = os.getenv("FLOWTORCH_PUBLIC", "0") == "1"

# Tokens for gated endpoints
METRICS_TOKEN = os.getenv("FLOWTORCH_METRICS_TOKEN", "")

# Per-IP rate limiter for admin endpoints (token bucket)
ADMIN_RATE_LIMIT_QPS = float(os.getenv("FLOWTORCH_ADMIN_RATE_LIMIT_QPS", "10"))  # QPS per IP
ADMIN_RATE_BURST = max(1, int(os.getenv("FLOWTORCH_ADMIN_RATE_BURST", str(int(ADMIN_RATE_LIMIT_QPS)))))  # Burst capacity
ADMIN_IP_BUCKET_CAP = 10000  # Soft cap on per-worker IP bucket size (LRU eviction)

# Token bucket storage: {ip: (tokens, last_check_time)}
_admin_ip_bucket: dict[str, tuple[float, float]] = {}
_admin_bucket_lock = threading.Lock()
_LAST_EVICT_TS = 0.0  # Throttle eviction to prevent O(n log n) on every hot path


def _tb_check(ip: str, bucket: dict, qps: float, burst: int, cap: int = 10000, lock: threading.Lock = None) -> bool:
    """
    Generic token bucket rate limiter with LRU eviction.
    Returns True if request allowed, False if rate-limited.
    Uses monotonic clock (immune to NTP jumps) for rate calculations.

    Memory-bounded: Evicts oldest IPs when bucket size exceeds cap.
    Thread-safe: Uses lock if provided to prevent race conditions.
    """
    now = _now()

    if lock:
        with lock:
            tokens, last = bucket.get(ip, (float(burst), now))
            elapsed = max(0.0, now - last)
            tokens = min(float(burst), tokens + elapsed * qps)

            if tokens < 1.0:
                bucket[ip] = (tokens, now)
                return False

            bucket[ip] = (tokens - 1.0, now)

            # LRU eviction: throttled to max once per second to prevent O(n log n) on hot path
            global _LAST_EVICT_TS
            if len(bucket) > cap and (now - _LAST_EVICT_TS) > 1.0:
                _LAST_EVICT_TS = now
                sorted_ips = sorted(bucket.items(), key=lambda x: x[1][1])  # sort by last_check_time
                evict_count = max(1, len(bucket) // 10)  # evict at least 1, up to 10%
                for evict_ip, _ in sorted_ips[:evict_count]:
                    bucket.pop(evict_ip, None)

            return True
    else:
        # No lock provided - fallback (non-thread-safe)
        tokens, last = bucket.get(ip, (float(burst), now))
        elapsed = max(0.0, now - last)
        tokens = min(float(burst), tokens + elapsed * qps)

        if tokens < 1.0:
            bucket[ip] = (tokens, now)
            return False

        bucket[ip] = (tokens - 1.0, now)

        # LRU eviction: throttled to max once per second to prevent O(n log n) on hot path
        global _LAST_EVICT_TS
        if len(bucket) > cap and (now - _LAST_EVICT_TS) > 1.0:
            _LAST_EVICT_TS = now
            sorted_ips = sorted(bucket.items(), key=lambda x: x[1][1])
            evict_count = max(1, len(bucket) // 10)
            for evict_ip, _ in sorted_ips[:evict_count]:
                bucket.pop(evict_ip, None)

        return True


def _check_ip_rate_limit(client_ip: str) -> bool:
    """
    Per-IP rate limiter for admin endpoints.
    Uses ADMIN_RATE_LIMIT_QPS, ADMIN_RATE_BURST, and ADMIN_IP_BUCKET_CAP.
    Thread-safe via _admin_bucket_lock.
    """
    return _tb_check(client_ip, _admin_ip_bucket, ADMIN_RATE_LIMIT_QPS, ADMIN_RATE_BURST, ADMIN_IP_BUCKET_CAP, _admin_bucket_lock)


async def require_metrics_auth(request: Request, x_metrics_token: str = Header(default="")):
    """
    Dependency that gates access to metrics/receipts/observability endpoints.

    In PUBLIC mode (FLOWTORCH_PUBLIC=1): allows all requests (with rate limiting)
    In private mode (default): requires x-metrics-token header matching
    FLOWTORCH_METRICS_TOKEN env var (with rate limiting).

    Per-IP rate limiting: FLOWTORCH_ADMIN_RATE_LIMIT_QPS (default 10 QPS)
    Uses proxy-aware IP resolution (P2 trust rules) for consistent rate limiting.

    Raises:
        HTTPException: 403 if token missing/invalid in private mode
        HTTPException: 429 if rate-limited

    Usage:
        @app.get("/metrics.txt", dependencies=[Depends(require_metrics_auth)])
        async def metrics_text(): ...
    """
    # Get client IP (proxy-aware)
    client_ip = proxy_get_client_ip(request) or "unknown"

    # Check rate limit first (applies to both public and private mode)
    if not _check_ip_rate_limit(client_ip):
        # Use proper 429 contract (matches _shed() in serve_flowtorch_prod.py)
        retry_after_sec = 1
        now = time.time()
        payload = {
            "error": "rate_limited",
            "detail": "too_many_requests",
            "next_allowed_at": round(now + retry_after_sec, 3),
            "granularity": "ip",
            "source": "admin_rate_limit"
        }
        headers = {
            "Retry-After": str(retry_after_sec),
            "Cache-Control": "no-store",
            "X-Content-Type-Options": "nosniff",
            "Content-Type": "application/json",  # mirror 504 contract
        }
        return JSONResponse(content=payload, status_code=429, headers=headers)  # guard-ok:429-contract

    if PUBLIC:
        # Public mode - no auth required (but rate-limited)
        return

    # Private mode - require token
    if not METRICS_TOKEN or x_metrics_token != METRICS_TOKEN:
        raise HTTPException(status_code=403, detail="forbidden")


def _norm_ip(token: str) -> str | None:
    """Normalize IP address token, return None if invalid."""
    try:
        return str(ipaddress.ip_address(token))
    except Exception:
        return None


def _is_private(ip: str) -> bool:
    """Check if IP is private/loopback/reserved."""
    try:
        ipo = ipaddress.ip_address(ip)
        return (ipo.is_private or ipo.is_loopback or ipo.is_link_local or
                ipo.is_reserved or ipo.is_multicast or ipo.is_unspecified)
    except Exception:
        return True


def proxy_get_client_ip(request: Request) -> str | None:
    """
    Proxy-aware client IP resolution (P2 hardened).

    Security features:
    - Only trusts X-Forwarded-For when socket peer is in TRUSTED_PROXIES
    - Normalizes and validates all IP tokens
    - Prefers public IPs over private
    - Guards against header bloat (2KB cap)

    NOTE: We intentionally ignore RFC 7239 `Forwarded` header to avoid mixed parsing.
    Only X-Forwarded-For is supported and normalized here.
    """
    if not request or not request.client:
        return None

    direct_ip = _norm_ip(request.client.host) or "0.0.0.0"

    if not TRUST_PROXY:
        return direct_ip

    # SECURITY: We only honor XFF if the direct socket peer is trusted.
    # This prevents spoofing attacks where untrusted clients inject XFF headers.
    if not _peer_trusted(direct_ip):
        return direct_ip

    # Cap header length to prevent bloat attacks
    raw_xff = (request.headers.get("x-forwarded-for") or "")[:2048]
    parts = [p.strip() for p in raw_xff.split(",") if p.strip()]

    # Normalize and drop invalids
    parts = [_norm_ip(p) for p in parts]
    parts = [p for p in parts if p]

    if not parts:
        return direct_ip

    # SECURITY: Trim trusted proxy hops from the right.
    # HOPS trimming assumes a trusted chain behind the peer (verified above).
    if TRUSTED_PROXY_HOPS and len(parts) > TRUSTED_PROXY_HOPS:
        parts = parts[:-TRUSTED_PROXY_HOPS]

    # Prefer the left-most globally routable IP if available
    # If none are globally routable (all are RFC-reserved or private), return leftmost candidate
    for cand in parts:
        if not _is_private(cand):
            return cand

    # Nothing globally routable → use first available (may be RFC-reserved or private, last resort)
    return parts[0]
