# Copyright (c) Streamlit Inc. (2018-2022) Snowflake Inc. (2022-2026)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Utility functions for the Starlette server implementation."""

from __future__ import annotations

import binascii
import os
import time
from typing import Final

from streamlit.web.server.starlette.starlette_server_config import (
    AUTH_COOKIE_MAX_AGE_SECONDS,
)

# Allow a small server-side grace period beyond the browser cookie max-age so
# recently expired auth cookies can still be decoded during edge-case timing.
_AUTH_COOKIE_GRACE_PERIOD_SECONDS: Final = 24 * 60 * 60
_DEFAULT_SIGNED_VALUE_MAX_AGE_DAYS: Final = (
    AUTH_COOKIE_MAX_AGE_SECONDS + _AUTH_COOKIE_GRACE_PERIOD_SECONDS
) / 86400


def parse_range_header(range_header: str, total_size: int) -> tuple[int, int]:
    """Parse the Range header and return the start and end byte positions.

    This is used for serving media files with range requests.

    Parameters
    ----------
    range_header : str
        The value of the Range header (e.g. "bytes=0-1023").

    total_size : int
        The total size of the resource in bytes.

    Returns
    -------
    tuple[int, int]
        A tuple containing (start, end) byte positions.
    """
    if total_size <= 0:
        raise ValueError("empty content")

    units, sep, range_spec = range_header.partition("=")
    if units.strip().lower() != "bytes" or sep == "" or "," in range_spec:
        raise ValueError("invalid range")

    range_spec = range_spec.strip()
    if range_spec.startswith("-"):
        try:
            suffix = int(range_spec[1:])
        except ValueError:
            raise ValueError("invalid suffix range") from None
        if suffix <= 0:
            raise ValueError("invalid suffix range")
        if suffix >= total_size:
            return 0, total_size - 1
        return total_size - suffix, total_size - 1

    start_str, sep, end_str = range_spec.partition("-")
    if not start_str:
        raise ValueError("missing range start")

    start = int(start_str)
    if start < 0 or start >= total_size:
        raise ValueError("start out of range")

    if sep == "" or not end_str:
        end = total_size - 1
    else:
        end = int(end_str)
        if end < start:
            raise ValueError("end before start")
        end = min(end, total_size - 1)

    return start, end


def websocket_mask(mask: bytes, data: bytes) -> bytes:
    """Mask or unmask data for WebSocket transmission per RFC 6455.

    Each byte of data is XORed with mask[i % 4]. This operation is
    bidirectional - applying it twice with the same mask returns the
    original data.

    Parameters
    ----------
    mask : bytes
        A 4-byte masking key.
    data : bytes
        The data to mask or unmask.

    Returns
    -------
    bytes
        The masked/unmasked data.
    """
    if len(mask) != 4:
        raise ValueError("mask must be 4 bytes")

    result = bytearray(len(data))
    for i, byte in enumerate(data):
        result[i] = byte ^ mask[i % 4]
    return bytes(result)


def create_signed_value(
    secret: str,
    name: str,
    value: str | bytes,
) -> bytes:
    """Create a signed cookie value using itsdangerous.

    Parameters
    ----------
    secret
        The secret key used for signing.
    name
        The cookie name (used as salt for additional security).
    value
        The value to sign.

    Returns
    -------
    bytes
        The signed value as bytes.
    """
    from itsdangerous import URLSafeTimedSerializer

    serializer = URLSafeTimedSerializer(secret, salt=name)
    if isinstance(value, bytes):
        value = value.decode("utf-8")
    return serializer.dumps(value).encode("utf-8")


def decode_signed_value(
    secret: str,
    name: str,
    value: str | bytes,
    max_age_days: float = _DEFAULT_SIGNED_VALUE_MAX_AGE_DAYS,
) -> bytes | None:
    """Decode a signed cookie value using itsdangerous.

    Parameters
    ----------
    secret
        The secret key used for signing.
    name
        The cookie name (used as salt for additional security).
    value
        The signed value to decode.
    max_age_days
        Maximum age of the cookie in days. Defaults to the auth cookie lifetime
        plus a 1-day grace period.

    Returns
    -------
    bytes | None
        The decoded value as bytes, or None if invalid/expired.
    """
    from itsdangerous import BadSignature, SignatureExpired, URLSafeTimedSerializer

    if not value:
        return None

    try:
        if isinstance(value, bytes):
            value = value.decode("utf-8")

        serializer = URLSafeTimedSerializer(secret, salt=name)
        decoded = serializer.loads(value, max_age=int(max_age_days * 86400))
        if isinstance(decoded, str):
            return decoded.encode("utf-8")
        if isinstance(decoded, bytes):  # pragma: no cover - defensive
            return decoded
        # Unexpected type from deserializer — treat as invalid
        return None  # pragma: no cover - defensive
    except (BadSignature, SignatureExpired, UnicodeDecodeError):
        return None


def generate_xsrf_token_string(
    token_bytes: bytes | None = None, timestamp: int | None = None
) -> str:
    """Generate a version 2 XSRF token string.

    Format: 2|mask|masked_token|timestamp

    Parameters
    ----------
    token_bytes
        The raw token bytes to encode. If None, generates 16 random bytes.
    timestamp
        The Unix timestamp to include in the token. If None, uses current time.

    Returns
    -------
    str
        The encoded XSRF token string in version 2 format.
    """
    if token_bytes is None:
        token_bytes = os.urandom(16)
    if timestamp is None:
        timestamp = int(time.time())

    mask = os.urandom(4)
    masked_token = websocket_mask(mask, token_bytes)
    return "2|{}|{}|{}".format(
        binascii.b2a_hex(mask).decode("ascii"),
        binascii.b2a_hex(masked_token).decode("ascii"),
        timestamp,
    )


def decode_xsrf_token_string(
    cookie_value: str,
) -> tuple[bytes | None, int | None]:
    """Decode an XSRF token string.

    Supports version 2 (masked) and version 1 (unmasked) tokens.

    Parameters
    ----------
    cookie_value
        The XSRF token cookie value to decode.

    Returns
    -------
    tuple[bytes | None, int | None]
        A tuple of (token_bytes, timestamp). Both values are None if decoding fails.
    """
    if not cookie_value:
        return None, None

    value = cookie_value.strip("\"'")
    if not value:
        return None, None

    try:
        # V2 tokens:
        if value.startswith("2|"):
            _, mask_hex, masked_hex, timestamp_str = value.split("|")
            mask = binascii.a2b_hex(mask_hex.encode("ascii"))
            masked = binascii.a2b_hex(masked_hex.encode("ascii"))
            token = websocket_mask(mask, masked)
            return token, int(timestamp_str)

        # V1 tokens:
        # TODO(lukasmasuch): This is likely unused in Streamlit since only V2 tokens
        # are used. We might be able to just remove this part.
        token = binascii.a2b_hex(value.encode("ascii"))
        if (
            not token
        ):  # pragma: no cover - defensive, non-empty hex never decodes to empty
            return None, None
        # V1 tokens don't have an embedded timestamp, so we use current time
        # as a placeholder. This timestamp is informational only and not used
        # for token validation.
        return token, int(time.time())
    except (binascii.Error, ValueError):
        return None, None


def generate_random_hex_string(num_bytes: int = 32) -> str:
    """Generate a cryptographically secure random hex string.

    Parameters
    ----------
    num_bytes
        Number of random bytes to generate (default: 32).
        The resulting hex string will be twice this length.

    Returns
    -------
    str
        A hex-encoded random string.
    """
    return binascii.b2a_hex(os.urandom(num_bytes)).decode("ascii")


def validate_xsrf_token(supplied_token: str | None, xsrf_cookie: str | None) -> bool:
    """Validate the XSRF token from the WebSocket subprotocol against the cookie.

    This ensures the frontend can share XSRF logic between WebSocket handshake
    and HTTP uploads.
    """

    if not supplied_token or not xsrf_cookie:
        return False

    # Decode the supplied token from the subprotocol
    supplied_token_bytes, _ = decode_xsrf_token_string(supplied_token)
    # Decode the expected token from the cookie
    expected_token_bytes, _ = decode_xsrf_token_string(xsrf_cookie)

    if not supplied_token_bytes or not expected_token_bytes:
        return False

    import hmac

    return hmac.compare_digest(supplied_token_bytes, expected_token_bytes)
