# Copyright © 2026 Contrast Security, Inc.
# See https://www.contrastsecurity.com/enduser-terms-0317a for more details.
from __future__ import annotations

import os
import resource
import subprocess
import sys
from socket import gethostname
from threading import Thread

from contrast_fireball import MemoryMetrics, ServerInventory

from contrast.reporting import Reporter
from contrast.utils.configuration_utils import get_platform
from contrast.utils.decorators import fail_loudly, fail_quietly
from contrast_vendor import structlog as logging

logger = logging.getLogger("contrast")


REPORT_RESOURCE_ID_THREAD_NAME = "ContrastReportServerInventory"


def report_server_runtime(
    reporting_client: Reporter, discover_cloud_resource_enabled: bool = True
) -> None:
    """
    Report runtime environment info to TS in a background thread.
    """
    logger.debug("Starting detect runtime environment background thread")
    Thread(
        name=REPORT_RESOURCE_ID_THREAD_NAME,
        daemon=True,
        target=_report_runtime_environment,
        args=(reporting_client, discover_cloud_resource_enabled),
    ).start()


@fail_loudly("Failed to report operating system runtime environment details.")
def _report_runtime_environment(
    reporting_client: Reporter, discover_cloud_resource_enabled
) -> None:
    provider = None
    resource_id = None

    if discover_cloud_resource_enabled:
        provider, resource_id = _detect_cloud_resource_id()

    memory_metrics = get_memory_metrics()

    reporting_client.new_server_inventory(
        ServerInventory(
            operating_system=get_platform(),
            runtime_path=sys.executable,
            runtime_version=sys.version.split(" ")[0],
            hostname=gethostname(),
            is_kubernetes=_is_kubernetes(),
            is_docker=_is_docker(),
            cloud_provider=provider,
            cloud_resource_id=resource_id,
            memory_metrics=MemoryMetrics(memory_metrics.process_memory_limit_bytes),
        )
    )


def _is_kubernetes():
    return "KUBERNETES_SERVICE_HOST" in os.environ


@fail_quietly()
def _is_docker():
    docker_env_file = "/.dockerenv"
    cgroup_file = "/proc/self/cgroup"

    if os.path.isfile(docker_env_file):
        return True

    try:
        with open(cgroup_file) as cgroup:
            return "/docker" in cgroup.read().strip()
    except FileNotFoundError:
        logger.debug("The file %s was not found", cgroup_file)

    return False


def _detect_cloud_resource_id():
    extract_from = {
        "aws": extract_aws_resource_id,
        "azure": extract_azure_resource_id,
        # "gcp": extract_gcp_resource_id, # TODO: PYT-3346
    }

    for provider, extract_id_func in extract_from.items():
        try:
            resource_id = extract_id_func()
        except ResourceIdError as e:
            logger.debug(
                "did not extract resource ID",
                provider=provider,
                exc_info=e,
            )
            continue

        logger.debug(
            "discovered resource ID",
            cloud_provider=provider,
            resource_id=resource_id,
        )
        return provider, resource_id
    return None, None


METADATA_ENDPOINT_ADDRESS = "169.254.169.254"
TOTAL_RETRIES = 5
RETRY_BACKOFF_FACTOR = 1  # wait 0s, 1s, 2s, 4s, 8s, ..., on request retry
NO_PROXIES = {"http": "", "https": ""}
TIMEOUT_SECONDS = 10


class ResourceIdError(Exception):
    pass


def extract_azure_resource_id() -> str:
    """
    Extract the resource ID from the Azure Instance Metadata Service.
    """
    import requests
    from requests.adapters import HTTPAdapter
    from urllib3.exceptions import MaxRetryError
    from urllib3.util import Retry

    # This endpoint comes from https://learn.microsoft.com/en-us/azure/virtual-machines/instance-metadata-service?tabs=linux#route-parameters
    RESOURCE_ID_ENDPOINT = f"http://{METADATA_ENDPOINT_ADDRESS}/metadata/instance/compute/resourceId?api-version=2023-07-01&format=text"

    # https://learn.microsoft.com/en-us/azure/virtual-machines/instance-metadata-service?tabs=linux#errors-and-debugging
    AZURE_RETRY_CODES = {410, 429, 500}

    with requests.Session() as session:
        retries = Retry(
            total=TOTAL_RETRIES,
            backoff_factor=RETRY_BACKOFF_FACTOR,
            status_forcelist=Retry.RETRY_AFTER_STATUS_CODES.union(AZURE_RETRY_CODES),
        )
        session.mount("http://", HTTPAdapter(max_retries=retries))

        try:
            response = session.get(
                RESOURCE_ID_ENDPOINT,
                proxies=NO_PROXIES,
                headers={"Metadata": "true"},
                timeout=TIMEOUT_SECONDS,
            )
            response.raise_for_status()
        except (MaxRetryError, requests.RequestException) as e:
            raise ResourceIdError("Failed to get Azure resource ID") from e
        return response.text


def extract_aws_resource_id() -> str:
    """
    Generate an AWS ARN for the EC2 instance using its Identity Document.

    See:
    https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/instance-identity-documents.html
    """
    logger.debug("Attempting to extract AWS resource ID")

    import requests
    import requests.adapters
    import urllib3.exceptions
    import urllib3.util

    retry_config = urllib3.util.Retry(
        total=TOTAL_RETRIES,
        backoff_factor=RETRY_BACKOFF_FACTOR,
        # force a retry for status codes supporting Retry-After, even if the header is
        # missing
        status_forcelist=urllib3.util.Retry.RETRY_AFTER_STATUS_CODES,
    )

    with requests.Session() as session:
        session.mount(
            "http://", requests.adapters.HTTPAdapter(max_retries=retry_config)
        )
        token = _get_aws_token(session)
        # a token is required for IMDSv2, but if we can't get one, still try IMDSv1 (no
        # token)
        headers = {"X-aws-ec2-metadata-token": token} if token else {}
        try:
            response = session.get(
                f"http://{METADATA_ENDPOINT_ADDRESS}/latest/dynamic/instance-identity/document",
                headers=headers,
                proxies=NO_PROXIES,
                timeout=TIMEOUT_SECONDS,
            )
            response.raise_for_status()
            identity_doc = response.json()
        except (
            urllib3.exceptions.MaxRetryError,
            requests.RequestException,
        ) as e:
            raise ResourceIdError("Failed to get AWS resource ID") from e

    logger.debug("Retrieved AWS identity document", aws_identity_doc=identity_doc)
    try:
        region = identity_doc["region"]
        account_id = identity_doc["accountId"]
        instance_id = identity_doc["instanceId"]
    except KeyError as e:
        raise ResourceIdError(
            "Failed to get AWS resource ID. "
            "Missing a required field in identity document."
        ) from e

    arn = ":".join(
        [
            "arn",
            "aws",
            "ec2",
            region,
            account_id,
            f"instance/{instance_id}",
        ]
    )
    return arn


def _get_aws_token(session) -> str | None:
    logger.debug("Retrieving AWS token")

    import requests
    import urllib3.exceptions

    try:
        token_response = session.put(
            f"http://{METADATA_ENDPOINT_ADDRESS}/latest/api/token",
            headers={
                "X-aws-ec2-metadata-token-ttl-seconds": "300",
            },
            proxies=NO_PROXIES,
            timeout=TIMEOUT_SECONDS,
        )
        token_response.raise_for_status()
        return token_response.text
    except (
        urllib3.exceptions.MaxRetryError,
        requests.RequestException,
    ) as e:
        logger.debug(
            "Unable to retrieve token for AWS IMDSv2 - proceeding without it",
            error=str(e),
        )
    return None


CGROUP_V1_PATH = "/sys/fs/cgroup/memory/memory.limit_in_bytes"
CGROUP_V2_PATH = "/sys/fs/cgroup/memory.max"


def get_memory_metrics() -> MemoryMetrics:
    """
    Retrieve memory information for the current machine / process.
    """
    virtual_mem = _total_virtual_memory_bytes()
    cgroup_mem = _cgroup_memory_limit_bytes()
    _, process_hard_limit_mem = _process_rlimits_bytes()

    total_mem = virtual_mem
    if virtual_mem is not None and cgroup_mem is not None:
        total_mem = min(virtual_mem, cgroup_mem)

    process_mem = total_mem
    if total_mem is not None and process_hard_limit_mem is not None:
        process_mem = min(total_mem, process_hard_limit_mem)

    return MemoryMetrics(process_memory_limit_bytes=process_mem)


LINUX = sys.platform.startswith("linux")
MACOS = sys.platform.startswith("darwin")


@fail_quietly(return_value=None)
def _total_virtual_memory_bytes() -> int | None:
    """
    Get the total memory available on the machine in bytes. Based on the implementation
    of `virtual_memory` from https://github.com/giampaolo/psutil.
    """
    if LINUX:
        return _linux_virtual_memory()
    if MACOS:
        return _macos_virtual_memory()
    return None


def _linux_virtual_memory() -> int:
    """
    See https://github.com/giampaolo/psutil/blob/e7754af7400c0836a721befaba0f54c402c4ecb4/psutil/_pslinux.py#L351-L369
    """
    mems = {}
    # Buffering prevents several issues. See the larger comment in psutil.
    with open("/proc/meminfo", "rb", buffering=32 * 1024) as f:
        for line in f:
            fields = line.split()
            mems[fields[0]] = int(fields[1]) * 1024

    return mems[b"MemTotal:"]


def _macos_virtual_memory() -> int:
    """
    Get the total memory available on macOS in bytes.
    """
    output = subprocess.check_output(["sysctl", "-n", "hw.memsize"])
    return int(output.strip())


@fail_quietly(return_value=None)
def _cgroup_memory_limit_bytes() -> int | None:
    """
    Get the memory limit in bytes imposed by cgroups (used by Docker), or None if not
    found or unlimited.
    """
    for path in (CGROUP_V2_PATH, CGROUP_V1_PATH):
        if os.path.exists(path):
            with open(path) as f:
                val = f.read().strip()
            if not val.isdigit():
                return None
            limit = int(val)
            if limit < 0:
                return None
            return limit
    return None


@fail_quietly(return_value=(None, None))
def _process_rlimits_bytes() -> tuple[int | None, int | None]:
    """
    Get the current memory limits for this process in bytes, according to getrlimit.
    Values that cannot be found or are unlimited are reported as None.
    """
    soft_limit, hard_limit = resource.getrlimit(resource.RLIMIT_AS)
    return (
        soft_limit if soft_limit != resource.RLIM_INFINITY else None,
        hard_limit if hard_limit != resource.RLIM_INFINITY else None,
    )
