"""
security.py — Qualora Security & Compliance
============================================
Consolidated: auth.py + limits.py + alerts.py

Strictly hardened for Vercel (Serverless) & Enterprise MVC standards.
Features:
- XSS-immune HTTP-Only JWTs + CSRF protection.
- Sliding-window rate limiting with OOM-safe garbage collection.
- Context-aware Webhook dispatching (prevents lambda thread-freezing).
"""

import os
import hmac
import secrets
import time
import random
import logging
import threading
from functools import wraps
from datetime import datetime, timedelta, timezone
from typing import Optional, Dict, Any
from collections import deque
import requests

import jwt
from flask import request, jsonify
from werkzeug.security import generate_password_hash, check_password_hash
from bson import ObjectId
from bson.errors import InvalidId

# Centralized core imports
from core import (
    JWT_SECRET, JWT_EXPIRATION_SECONDS, DEBUG, IS_PRODUCTION,
    AUTH_RATE_LIMIT_PER_MIN, AUDIT_RATE_LIMIT_PER_MIN,
    ABUSE_BLOCK_THRESHOLD, ABUSE_BLOCK_DURATION_SEC,
    get_db, infer_f1_score
)

logger = logging.getLogger(__name__)

# Keyword-based mapping for compliance flag severity classification.
# Edit these lists to tune which LLM `compliance_flags` should escalate.
_COMPLIANCE_FLAG_SEVERITY_KEYWORDS = {
    "critical": [
        "legal", "privacy", "confidential", "pci", "ssn", "gdpr",
        "breach", "data leak", "fraud", "criminal", "safety"
    ],
    "warning": [
        "escalate", "escalation", "supervisor", "manager", "sensitive",
        "policy_violation", "policy", "sla"
    ],
    "info": [
        "note", "format", "minor", "typo", "cosmetic"
    ]
}


# Allow runtime override of the keyword lists via environment variables
# Use comma-separated lists, e.g. COMPLIANCE_FLAGS_CRITICAL="privacy,breach"
def _load_compliance_flag_severity_keywords_from_env():
    try:
        for sev, env_name in (
            ("critical", "COMPLIANCE_FLAGS_CRITICAL"),
            ("warning",  "COMPLIANCE_FLAGS_WARNING"),
            ("info",     "COMPLIANCE_FLAGS_INFO"),
        ):
            raw = os.environ.get(env_name, "").strip()
            if not raw:
                continue
            parts = [p.strip() for p in raw.split(",") if p.strip()]
            if parts:
                _COMPLIANCE_FLAG_SEVERITY_KEYWORDS[sev] = parts
        logger.debug("Loaded compliance flag severity overrides from env (if any).")
    except Exception as e:
        logger.warning("Failed to load compliance flag overrides from env: %s", e)


# Load overrides at import time so the runtime behavior is adjustable via env vars
_load_compliance_flag_severity_keywords_from_env()

# ══════════════════════════════════════════════════════════════════════════════
# AUTHENTICATION & AUTHORIZATION
# ══════════════════════════════════════════════════════════════════════════════

def hash_password(password: str) -> str:
    """Hash a plaintext password using Werkzeug's secure hashing (Vercel-compatible)."""
    return generate_password_hash(password)

def verify_password(hashed: str, password: str) -> bool:
    """Verify a plaintext password against a Werkzeug hash."""
    return check_password_hash(hashed, password)

def create_token(user_id: str, org_id: str, role: str, email: str) -> str:
    """Create a signed JWT token with user identity, org isolation, and role."""
    payload = {
        "user_id": str(user_id),
        "org_id": str(org_id),
        "role": role,
        "email": email,
        "iat": datetime.now(timezone.utc),
        "exp": datetime.now(timezone.utc) + timedelta(seconds=JWT_EXPIRATION_SECONDS),
        "iss": "qualora_enterprise"
    }
    return jwt.encode(payload, JWT_SECRET, algorithm="HS256")

def decode_token(token: str) -> Optional[Dict[str, Any]]:
    """Decode and verify a JWT token, enforcing issuer."""
    if not token or not isinstance(token, str):
        return None
    try:
        return jwt.decode(token, JWT_SECRET, algorithms=["HS256"], issuer="qualora_enterprise")
    except (jwt.ExpiredSignatureError, jwt.InvalidTokenError, Exception):
        return None

def generate_csrf_token() -> str:
    return secrets.token_urlsafe(32)

def set_auth_cookies(response, token: str, csrf_token: str):
    """Set secure cookies on the response (access_token is HTTP-Only for XSS immunity)."""
    response.set_cookie(
        "access_token",
        token,
        httponly=True,
        secure=IS_PRODUCTION,
        samesite="Strict",
        max_age=JWT_EXPIRATION_SECONDS,
        path="/",
    )
    # Mirror cookie: NOT Http-Only. Used by auth-init.js to sync state & prevent redirect loops.
    response.set_cookie(
        "qualora_logged_in",
        "true",
        httponly=False,
        secure=IS_PRODUCTION,
        samesite="Strict",
        max_age=JWT_EXPIRATION_SECONDS,
        path="/",
    )
    response.set_cookie(
        "csrf_token",
        csrf_token,
        httponly=False,  # JS must read this to construct headers
        secure=IS_PRODUCTION,
        samesite="Strict",
        max_age=JWT_EXPIRATION_SECONDS,
        path="/",
    )

def clear_auth_cookies(response):
    response.delete_cookie(
        "access_token",
        path="/",
        httponly=True,
        secure=IS_PRODUCTION,
        samesite="Strict",
    )
    response.delete_cookie(
        "qualora_logged_in",
        path="/",
        httponly=False,
        secure=IS_PRODUCTION,
        samesite="Strict",
    )
    response.delete_cookie(
        "csrf_token",
        path="/",
        httponly=False,
        secure=IS_PRODUCTION,
        samesite="Strict",
    )

def csrf_protect(f):
    """Validate CSRF token on state-changing requests using constant-time comparison."""
    @wraps(f)
    def wrapper(*args, **kwargs):
        if request.method in ['GET', 'HEAD', 'OPTIONS']:
            return f(*args, **kwargs)
        
        csrf_header = request.headers.get('X-CSRF-Token', '').strip()
        csrf_cookie = request.cookies.get('csrf_token', '').strip()
        
        if not csrf_header or not csrf_cookie:
            return jsonify({'error': 'CSRF token missing'}), 403
            
        if not hmac.compare_digest(csrf_header, csrf_cookie):
            return jsonify({'error': 'CSRF token mismatch'}), 403
            
        return f(*args, **kwargs)
    return wrapper

def _safe_user_doc(doc: Dict[str, Any]) -> Dict[str, Any]:
    """Remove sensitive fields from user document before sending to client."""
    doc.pop("password", None)
    doc.pop("reset_token", None)
    doc.pop("reset_token_expires", None)
    for key in ("_id", "org_id", "user_id"):
        if key in doc and doc[key] is not None:
            doc[key] = str(doc[key])
    return doc

def safe_oid(oid: Any) -> Optional[ObjectId]:
    if isinstance(oid, ObjectId):
        return oid
    try:
        return ObjectId(str(oid))
    except (InvalidId, TypeError, ValueError):
        return None
def require_auth(f):
    """Require valid JWT token in the access_token cookie (Bearer fallback strictly removed)."""
    @wraps(f)
    def wrapper(*args, **kwargs):
        token = request.cookies.get("access_token")
        if not token:
            return jsonify({"error": "Unauthorized: missing HTTP-Only token"}), 401
            
        payload = decode_token(token)
        if payload is None:
            return jsonify({"error": "Unauthorized: invalid or expired token"}), 401
            
        request.user = payload
        return f(*args, **kwargs)
    return wrapper

def optional_auth(f):
    """Attempt JWT auth via cookie, but allow guest passthrough if missing."""
    @wraps(f)
    def wrapper(*args, **kwargs):
        request.user = None
        token = request.cookies.get("access_token")
        if token:
            payload = decode_token(token)
            if payload is not None:
                request.user = payload
        return f(*args, **kwargs)
    return wrapper

def require_role(allowed_roles: list):
    """Require user to have one of the allowed roles (RBAC)."""
    def decorator(f):
        @wraps(f)
        def wrapper(*args, **kwargs):
            if not hasattr(request, 'user') or request.user is None:
                return jsonify({"error": "Unauthorized"}), 401
            if request.user.get('role') not in allowed_roles:
                return jsonify({"error": f"Forbidden: requires one of {allowed_roles}"}), 403
            return f(*args, **kwargs)
        return wrapper
    return decorator

def ensure_org_isolation(org_id_from_request: str) -> Optional[ObjectId]:
    """Verify requested org_id matches the authenticated user's org_id (IDOR prevention)."""
    if not hasattr(request, 'user') or not request.user:
        return None
    req_org_id = str(org_id_from_request)
    if str(request.user.get('org_id')) != req_org_id:
        return None
    return safe_oid(req_org_id)

AUTH_RATE_LIMIT = AUTH_RATE_LIMIT_PER_MIN
AUDIT_RATE_LIMIT = AUDIT_RATE_LIMIT_PER_MIN
ABUSE_THRESHOLD = ABUSE_BLOCK_THRESHOLD
ABUSE_BLOCK_SECS = ABUSE_BLOCK_DURATION_SEC
WINDOW_SECONDS = 60

# Thread-safe Standard Dicts (Not defaultdicts) to allow GC and prevent OOM leaks
_lock = threading.RLock()
_ip_windows: dict = {}
_org_windows: dict = {}
_failed_auth: dict = {}
_blocked_ips: dict = {}

_stats_lock = threading.Lock()
_rate_limit_hits = {"auth": 0, "audit": 0}
_abuse_blocks = 0

def _now() -> float:
    return time.monotonic()

def _lazy_garbage_collect():
    """Probabilistic memory sweep to purge empty dictionary keys and prevent OOM."""
    if random.random() > 0.05:
        return

    with _lock:
        now = _now()
        window_cutoff = now - WINDOW_SECONDS
        abuse_cutoff = now - ABUSE_BLOCK_SECS
        
        stale_ips = [ip for ip, dq in _ip_windows.items() if not dq or dq[-1] < window_cutoff]
        for ip in stale_ips: _ip_windows.pop(ip, None)
            
        stale_orgs = [org for org, dq in _org_windows.items() if not dq or dq[-1] < window_cutoff]
        for org in stale_orgs: _org_windows.pop(org, None)
            
        stale_fails = [ip for ip, fails in _failed_auth.items() if not fails or fails[-1] < abuse_cutoff]
        for ip in stale_fails: _failed_auth.pop(ip, None)
            
        stale_blocks = [ip for ip, unblock_at in _blocked_ips.items() if unblock_at < now]
        for ip in stale_blocks: _blocked_ips.pop(ip, None)

def _purge_window(dq: deque, cutoff: float):
    while dq and dq[0] < cutoff:
        dq.popleft()

def get_remote_ip() -> str:
    """Extract real IP securely, preventing X-Forwarded-For spoofing."""
    real_ip = request.headers.get("x-real-ip")
    if real_ip:
        return real_ip.strip()
        
    fwd = request.headers.get("X-Forwarded-For", "")
    if fwd:
        return fwd.split(",")[-1].strip()
        
    return request.remote_addr or "unknown"

def is_ip_blocked(ip: str) -> bool:
    with _lock:
        unblock_at = _blocked_ips.get(ip)
        if unblock_at is None:
            return False
        if _now() >= unblock_at:
            _blocked_ips.pop(ip, None)
            return False
        return True

def record_failed_auth(ip: str):
    global _abuse_blocks
    with _lock:
        cutoff = _now() - ABUSE_BLOCK_SECS
        fails = _failed_auth.get(ip, [])
        fails = [t for t in fails if t > cutoff]
        fails.append(_now())
        _failed_auth[ip] = fails
        
        if len(_failed_auth[ip]) >= ABUSE_THRESHOLD:
            _blocked_ips[ip] = _now() + ABUSE_BLOCK_SECS
            logger.warning("[RateLimit] IP %s blocked for %ds (too many failed auths)", ip, ABUSE_BLOCK_SECS)
            with _stats_lock:
                _abuse_blocks += 1

def record_success_auth(ip: str):
    with _lock:
        _failed_auth.pop(ip, None)

def check_ip_rate(ip: str, limit: int = AUTH_RATE_LIMIT) -> bool:
    with _lock:
        dq = _ip_windows.setdefault(ip, deque())
        cutoff = _now() - WINDOW_SECONDS
        _purge_window(dq, cutoff)
        if len(dq) >= limit:
            return False
        dq.append(_now())
        return True

def check_org_rate(org_id: str, limit: int = AUDIT_RATE_LIMIT) -> bool:
    if not org_id:
        return True
    with _lock:
        dq = _org_windows.setdefault(org_id, deque())
        cutoff = _now() - WINDOW_SECONDS
        _purge_window(dq, cutoff)
        if len(dq) >= limit:
            return False
        dq.append(_now())
        return True

def auth_rate_limit(f):
    @wraps(f)
    def wrapper(*args, **kwargs):
        _lazy_garbage_collect()
        ip = get_remote_ip()
        
        if is_ip_blocked(ip):
            logger.warning("[RateLimit] Blocked IP %s attempted access", ip)
            with _stats_lock:
                _rate_limit_hits["auth"] += 1
            return jsonify({"error": "Too many failed attempts — try again later"}), 429
            
        if not check_ip_rate(ip, AUTH_RATE_LIMIT):
            logger.warning("[RateLimit] IP %s exceeded auth rate limit", ip)
            with _stats_lock:
                _rate_limit_hits["auth"] += 1
            return jsonify({"error": "Rate limit exceeded — slow down"}), 429
            
        return f(*args, **kwargs)
    return wrapper

def audit_rate_limit(f):
    @wraps(f)
    def wrapper(*args, **kwargs):
        _lazy_garbage_collect()
        
        org_id = ""
        user = getattr(request, "user", None)
        if user:
            org_id = user.get("org_id", "")
        if not org_id:
            org_id = request.args.get("org_id", "") or (request.json or {}).get("org_id", "")

        if not check_org_rate(org_id, AUDIT_RATE_LIMIT):
            logger.warning("[RateLimit] Org %s exceeded audit rate limit", org_id)
            with _stats_lock:
                _rate_limit_hits["audit"] += 1
            return jsonify({"error": "Audit rate limit exceeded for your organisation"}), 429
            
        return f(*args, **kwargs)
    return wrapper

def get_rate_limit_status() -> dict:
    with _lock:
        with _stats_lock:
            return {
                "auth_limits_tracked": len(_ip_windows),
                "audit_limits_tracked": len(_org_windows),
                "blocked_ips": list(_blocked_ips.keys()),
                "total_auth_limit_hits": _rate_limit_hits["auth"],
                "total_audit_limit_hits": _rate_limit_hits["audit"],
                "total_abuse_blocks": _abuse_blocks,
            }

# ══════════════════════════════════════════════════════════════════════════════
# COMPLIANCE ALERTS
# ══════════════════════════════════════════════════════════════════════════════

def _safe_to_oid(id_val: any) -> ObjectId | str:
    if isinstance(id_val, ObjectId):
        return id_val
    try:
        return ObjectId(str(id_val))
    except (InvalidId, TypeError, ValueError):
        return str(id_val)

def evaluate_and_save_alerts(audit_result: dict, audit_id: str, org_id: str) -> list:
    """Evaluate audit results and dynamically dispatch webhooks without blocking Vercel."""
    triggered = []
    
    # ── Rule 1: Compliance Risk 
    if audit_result.get("compliance_risk") == "Red":
        triggered.append({
            "alert_type": "compliance_breach", "severity": "critical",
            "message": "Critical compliance breach detected.",
            "triggered_by": ["compliance_risk=Red"]
        })
        
    # ── Rule 2: Low Agent Performance (using mathematical inferencer)
    f1 = infer_f1_score(audit_result)
    if f1 < 0.6:
        triggered.append({
            "alert_type": "low_performance", "severity": "warning",
            "message": f"Agent F1 score is dangerously low: {f1:.2f}",
            "triggered_by": [f"f1_score={f1:.2f}"]
        })
        
    # ── Rule 3: Systematic Failures 
    evals = audit_result.get("success_evaluation", {})
    failures = [k for k, v in evals.items() if isinstance(v, dict) and v.get("pass") is False]
    if len(failures) >= 3:
        triggered.append({
            "alert_type": "systematic_failure", "severity": "critical",
            "message": f"Systematic failure across {len(failures)} criteria.",
            "triggered_by": failures
        })
        
    # ── Rule 4: Explicit compliance flags from the LLM output
    # The LLM may return fine-grained `compliance_flags` describing specific
    # policy violations (e.g. "privacy breach", "escalation required"). Treat
    # those as alertable items — map certain keywords to critical severity.
    try:
        flags = audit_result.get("compliance_flags", []) or []
        # Normalize to a list of strings
        if isinstance(flags, str):
            flags = [flags]
        elif isinstance(flags, dict):
            # dictionary form -> take truthy keys
            flags = [k for k, v in flags.items() if v]
        if flags:
            logger.info("Compliance flags present: %s", flags)
            seen = set()
            for flag in flags:
                if not flag:
                    continue
                fstr = str(flag).strip()
                if not fstr or fstr in seen:
                    continue
                seen.add(fstr)
                raw = fstr.lower()
                severity = "warning"
                # Escalate to critical for high-risk keywords
                if any(k in raw for k in ("legal", "privacy", "confidential", "pci", "ssn", "gdpr", "breach", "criminal", "fraud", "data leak", "safety")):
                    severity = "critical"
                triggered.append({
                    "alert_type": "compliance_flag",
                    "severity": severity,
                    "message": f"Compliance flag detected: {fstr}",
                    "triggered_by": [f"compliance_flag:{fstr}"]
                })
    except Exception as _err:
        logger.warning("Failed to evaluate compliance_flags for alerts: %s", _err)

    if not triggered:
        return []
        
    # ── Persist to DB 
    db = get_db()
    if db is not None:
        try:
            now = datetime.now(timezone.utc)
            alert_ids = []
            for alert in triggered:
                doc = {"org_id": _safe_to_oid(org_id), "audit_id": _safe_to_oid(audit_id), "acknowledged": False, "created_at": now, **alert}
                res = db.alerts.insert_one(doc)
                alert_ids.append(res.inserted_id)
            db.audits.update_one({"_id": _safe_to_oid(audit_id), "org_id": _safe_to_oid(org_id)}, {"$set": {"alerts_triggered": alert_ids}})
            logger.info(f"Saved {len(alert_ids)} alerts for audit {audit_id}")
        except Exception as e:
            logger.error("Failed to save alerts to database: %s", str(e))
            
    # ── Context-aware Dispatch 
    webhook_url = os.environ.get("WEBHOOK_URL", "").strip()
    slack_url = os.environ.get("SLACK_URL", "").strip()
    discord_url = os.environ.get("DISCORD_URL", "").strip()
    
    active_urls = [url for url in [webhook_url, slack_url, discord_url] if url]
    
    if active_urls:
        if os.environ.get("VERCEL_ENV"):
            # SERVERLESS: Must wait for IO before lambda freezes
            from concurrent.futures import ThreadPoolExecutor, wait
            with ThreadPoolExecutor(max_workers=len(active_urls)) as executor:
                futures = [executor.submit(_fire_webhook, url, triggered, audit_id, org_id) for url in active_urls]
                wait(futures, timeout=6.0)
        else:
            # SERVER: Fire-and-forget
            for url in active_urls:
                threading.Thread(target=_fire_webhook, args=(url, triggered, audit_id, org_id), daemon=True).start()
                
    return triggered

def _fire_webhook(url: str, alerts: list, audit_id: str, org_id: str) -> None:
    """Format and transmit payloads to external webhooks securely."""
    try:
        is_slack = "hooks.slack.com" in url
        is_discord = "discord.com/api/webhooks" in url
        
        payload = {
            "source": "qualora", 
            "org_id": str(org_id), 
            "audit_id": str(audit_id), 
            "alerts": alerts, 
            "timestamp": datetime.now(timezone.utc).isoformat()
        }
        
        if is_slack or is_discord:
            txt = f"🚨 *Qualora Alert Triggered*\nOrg: `{org_id}` | Audit: `{audit_id}`\n"
            for alert in alerts:
                emoji = "🔴" if alert['severity'] == "critical" else "🟠"
                txt += f"{emoji} **{alert['alert_type']}** ({alert['severity']}): {alert['message']}\n"
            payload = {"text": txt}
            if is_discord:
                payload["content"] = txt
                payload.pop("text", None)
                
        response = requests.post(url, json=payload, timeout=5)
        if response.status_code not in (200, 204):
            logger.warning(f"Webhook {url} returned {response.status_code}: {response.text[:100]}")
    except requests.exceptions.Timeout:
        logger.warning(f"Webhook {url} timed out.")
    except Exception as e:
        logger.error(f"Failed to fire webhook {url}: {e}")