"""Input validation, rate limiting, and security utilities."""
from __future__ import annotations

import ipaddress
import os
import re
import secrets
import time
import logging
import traceback
from collections import defaultdict
from typing import List, Optional
from urllib.parse import urlparse

from fastapi import HTTPException, Request
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.responses import JSONResponse

logger = logging.getLogger(__name__)

# Size limits
MAX_JSON_PAYLOAD = 1 * 1024 * 1024   # 1 MB
MAX_FILE_UPLOAD = 10 * 1024 * 1024    # 10 MB
MAX_ABSOLUTE_PAYLOAD = 100 * 1024 * 1024  # 100 MB absolute max

_HOSTNAME_RE = re.compile(r"^[a-zA-Z0-9]([a-zA-Z0-9\-]{0,61}[a-zA-Z0-9])?$")
_NODE_ID_RE = re.compile(r"^[0-9a-f]{12}$")
_SAFE_PATH_PARAM_RE = re.compile(r"^[a-zA-Z0-9_\-]+$")
_ATLAS_DB_NAME_RE = re.compile(r"^[a-zA-Z0-9_\.]+$")


def sanitize_hostname(hostname: str) -> str:
    """Validate and return hostname — alphanumeric + hyphens only.

    Strips common suffixes like .local, .lan, .home before validation.
    Raises ValueError on invalid input.
    """
    hostname = hostname.strip()
    # Strip mDNS / local network suffixes
    for suffix in (".local", ".lan", ".home", ".localdomain", ".internal"):
        if hostname.endswith(suffix):
            hostname = hostname[: -len(suffix)]
            break
    if not hostname or len(hostname) > 63:
        raise ValueError(f"Invalid hostname length: {len(hostname)}")
    if not _HOSTNAME_RE.match(hostname):
        raise ValueError(f"Invalid hostname: must be alphanumeric + hyphens, got '{hostname}'")
    return hostname


def validate_node_id(node_id: str) -> str:
    """Validate node_id — must be exactly 12 hex characters.

    Raises ValueError on invalid input.
    """
    node_id = node_id.strip().lower()
    if not _NODE_ID_RE.match(node_id):
        raise ValueError(f"Invalid node_id: must be 12 hex chars, got '{node_id}'")
    return node_id


def sanitize_path_param(value: str) -> str:
    """Sanitize a path parameter to prevent path traversal / injection.

    Strips slashes, dots, backslashes, and non-alphanumeric chars (except _ and -).
    Raises ValueError if result is empty.
    """
    # Strip dangerous characters
    cleaned = value.strip()
    cleaned = cleaned.replace("/", "").replace("\\", "").replace("..", "")
    # Allow only alphanumeric, underscore, hyphen
    if not _SAFE_PATH_PARAM_RE.match(cleaned):
        raise ValueError(f"Invalid path parameter: '{value}'")
    if not cleaned:
        raise ValueError("Empty path parameter")
    return cleaned


# ── Agent URL Validation (SSRF Prevention) ──────────────────────────────────

# Private/reserved IP ranges that are blocked by default
_BLOCKED_NETWORKS = [
    ipaddress.ip_network("127.0.0.0/8"),       # Loopback
    ipaddress.ip_network("::1/128"),            # IPv6 loopback
    ipaddress.ip_network("169.254.0.0/16"),     # Link-local
    ipaddress.ip_network("fe80::/10"),          # IPv6 link-local
    ipaddress.ip_network("0.0.0.0/8"),          # "This" network
]


def validate_agent_url(
    url: str,
    allowed_networks: Optional[List[str]] = None,
    allow_localhost: bool = False,
) -> str:
    """Validate an agent URL to prevent SSRF attacks.

    Checks:
    - Must be http:// or https://
    - Valid hostname or IP address
    - Port in valid range (1-65535)
    - Not loopback/link-local unless allow_localhost=True
    - If allowed_networks is set, IP must be in one of those CIDRs

    Returns the validated URL. Raises ValueError on invalid input.
    """
    if not url or not url.strip():
        raise ValueError("Empty agent URL")

    parsed = urlparse(url.strip())

    # Scheme check
    if parsed.scheme not in ("http", "https"):
        raise ValueError(f"Invalid URL scheme: '{parsed.scheme}'. Must be http or https.")

    # Must have a hostname
    hostname = parsed.hostname
    if not hostname:
        raise ValueError("URL missing hostname")

    # Port validation
    port = parsed.port
    if port is not None and (port < 1 or port > 65535):
        raise ValueError(f"Invalid port: {port}")

    # Resolve hostname to IP for network checks
    try:
        ip = ipaddress.ip_address(hostname)
    except ValueError:
        # It's a hostname, not an IP — resolve it
        import socket
        try:
            resolved = socket.getaddrinfo(hostname, port, proto=socket.IPPROTO_TCP)
            if not resolved:
                raise ValueError(f"Cannot resolve hostname: '{hostname}'")
            ip = ipaddress.ip_address(resolved[0][4][0])
        except socket.gaierror:
            raise ValueError(f"Cannot resolve hostname: '{hostname}'")

    # Block dangerous networks unless explicitly allowed
    if not allow_localhost:
        for blocked in _BLOCKED_NETWORKS:
            if ip in blocked:
                raise ValueError(
                    f"Agent URL resolves to blocked network ({ip}). "
                    f"Set allow_localhost=True to permit local addresses."
                )

    # If allowed_networks is specified, enforce CIDR allowlist
    if allowed_networks is not None:
        in_allowed = False
        for cidr in allowed_networks:
            try:
                net = ipaddress.ip_network(cidr, strict=False)
                if ip in net:
                    in_allowed = True
                    break
            except ValueError:
                continue
        if not in_allowed:
            raise ValueError(
                f"Agent URL IP ({ip}) not in allowed networks: {allowed_networks}"
            )

    return url.strip()


# ── Rate Limiter ────────────────────────────────────────────────────────────

class RateLimiter:
    """Simple in-memory rate limiter: max N requests per minute per IP."""

    def __init__(self, max_requests: int = 100, window_s: int = 60):
        self.max_requests = max_requests
        self.window_s = window_s
        self._requests: dict[str, list[float]] = defaultdict(list)
        self._request_count: int = 0

    def check(self, client_ip: str) -> bool:
        """Return True if the request is allowed, False if rate-limited."""
        now = time.time()
        cutoff = now - self.window_s

        # Periodic cleanup every 100 requests to prevent memory leak
        self._request_count += 1
        if self._request_count % 100 == 0:
            self.cleanup()

        # Prune old entries
        timestamps = self._requests[client_ip]
        self._requests[client_ip] = [t for t in timestamps if t > cutoff]

        if len(self._requests[client_ip]) >= self.max_requests:
            return False

        self._requests[client_ip].append(now)
        return True

    def cleanup(self):
        """Remove stale IPs (call periodically to avoid memory leak)."""
        now = time.time()
        cutoff = now - self.window_s
        stale = [ip for ip, ts in self._requests.items() if not ts or ts[-1] < cutoff]
        for ip in stale:
            del self._requests[ip]


# ── Error Sanitizer ─────────────────────────────────────────────────────────

class ErrorSanitizer:
    """Sanitize error messages to prevent information leakage in production."""

    # Patterns that indicate internal details
    _PATH_RE = re.compile(r"(/[a-zA-Z0-9_./-]+){3,}")  # File paths with 3+ segments
    _TRACEBACK_RE = re.compile(r"Traceback \(most recent call last\).*", re.DOTALL)
    _LINE_RE = re.compile(r"File \"[^\"]+\", line \d+")

    @classmethod
    def sanitize(cls, error: str, production: bool = True) -> str:
        """Sanitize an error message for external consumption.

        In production mode, strips file paths, tracebacks, and internal details.
        In development mode, returns the original error.
        """
        if not production:
            return error

        # Remove tracebacks
        sanitized = cls._TRACEBACK_RE.sub("[internal error]", error)
        # Remove file/line references
        sanitized = cls._LINE_RE.sub("[internal]", sanitized)
        # Remove absolute paths
        sanitized = cls._PATH_RE.sub("[path]", sanitized)

        return sanitized

    @classmethod
    def safe_error_response(cls, exception: Exception, production: bool = True) -> str:
        """Create a safe error string from an exception."""
        if production:
            # Return only the exception type and a generic message
            exc_type = type(exception).__name__
            # Allow known safe exception types to pass their message through
            safe_types = {"ValueError", "TypeError", "KeyError", "HTTPException"}
            if exc_type in safe_types:
                return cls.sanitize(str(exception), production=True)
            return "An internal error occurred. Check server logs for details."
        return str(exception)


# ── Validation Middleware ───────────────────────────────────────────────────

class ValidationMiddleware(BaseHTTPMiddleware):
    """Middleware that enforces payload size limits and rate limiting."""

    def __init__(
        self,
        app,
        max_json_size: int = MAX_JSON_PAYLOAD,
        max_upload_size: int = MAX_FILE_UPLOAD,
        rate_limit: int = 100,
    ):
        super().__init__(app)
        self.max_json_size = max_json_size
        self.max_upload_size = max_upload_size
        self.rate_limiter = RateLimiter(max_requests=rate_limit)

    async def dispatch(self, request: Request, call_next):
        # Rate limiting
        client_ip = request.client.host if request.client else "unknown"
        if not self.rate_limiter.check(client_ip):
            return JSONResponse(
                status_code=429,
                content={"detail": "Rate limit exceeded"},
            )

        # Content-length check for non-GET/HEAD/OPTIONS
        if request.method in ("POST", "PUT", "PATCH"):
            content_length = request.headers.get("content-length")
            if content_length:
                try:
                    size = int(content_length)
                except ValueError:
                    return JSONResponse(
                        status_code=400,
                        content={"detail": "Invalid Content-Length header"},
                    )

                # Reject negative or zero content-length
                if size < 0:
                    return JSONResponse(
                        status_code=400,
                        content={"detail": "Invalid Content-Length: must be non-negative"},
                    )

                # Absolute maximum (100MB)
                if size > MAX_ABSOLUTE_PAYLOAD:
                    return JSONResponse(
                        status_code=413,
                        content={"detail": f"Payload exceeds absolute maximum ({MAX_ABSOLUTE_PAYLOAD} bytes)"},
                    )

                content_type = request.headers.get("content-type", "")

                if "multipart" in content_type:
                    if size > self.max_upload_size:
                        return JSONResponse(
                            status_code=413,
                            content={"detail": f"Upload too large ({size} bytes, max {self.max_upload_size})"},
                        )
                else:
                    if size > self.max_json_size:
                        return JSONResponse(
                            status_code=413,
                            content={"detail": f"Payload too large ({size} bytes, max {self.max_json_size})"},
                        )

        return await call_next(request)


# ── Security Headers Middleware ─────────────────────────────────────────────

class SecurityHeadersMiddleware(BaseHTTPMiddleware):
    """Add security headers to all responses, including CSP with nonce."""

    async def dispatch(self, request: Request, call_next):
        # Generate a per-request nonce for CSP
        nonce = secrets.token_urlsafe(16)
        request.state.csp_nonce = nonce

        response = await call_next(request)
        response.headers["X-Content-Type-Options"] = "nosniff"
        response.headers["X-Frame-Options"] = "DENY"
        response.headers["X-XSS-Protection"] = "1; mode=block"
        response.headers["Referrer-Policy"] = "strict-origin-when-cross-origin"
        response.headers["Content-Security-Policy"] = (
            "default-src 'self'; "
            f"script-src 'self' 'nonce-{nonce}' https://unpkg.com; "
            "style-src 'self' 'unsafe-inline' https://cdn.jsdelivr.net; "
            "img-src 'self' data:; "
            "connect-src 'self'; "
            "frame-ancestors 'none'"
        )
        return response


# ── CSRF Middleware ─────────────────────────────────────────────────────────

class CSRFMiddleware(BaseHTTPMiddleware):
    """Simple CSRF protection — checks Origin/Referer on state-changing requests."""

    def __init__(self, app, allowed_origins: Optional[list[str]] = None):
        super().__init__(app)
        self.allowed_origins = set(allowed_origins or [])

    async def dispatch(self, request: Request, call_next):
        if request.method in ("POST", "PUT", "PATCH", "DELETE"):
            # API requests with JSON content-type are exempt (SOP protects these)
            content_type = request.headers.get("content-type", "")
            if "application/json" in content_type:
                return await call_next(request)

            # For form submissions, check Origin or Referer
            origin = request.headers.get("origin", "")
            referer = request.headers.get("referer", "")

            if origin:
                if self.allowed_origins and origin not in self.allowed_origins:
                    return JSONResponse(
                        status_code=403,
                        content={"detail": "CSRF check failed: invalid origin"},
                    )
            elif referer:
                # Extract origin from referer
                parsed = urlparse(referer)
                ref_origin = f"{parsed.scheme}://{parsed.netloc}"
                if self.allowed_origins and ref_origin not in self.allowed_origins:
                    return JSONResponse(
                        status_code=403,
                        content={"detail": "CSRF check failed: invalid referer"},
                    )

        return await call_next(request)
