# Copyright The OpenTelemetry Authors
# SPDX-License-Identifier: Apache-2.0

import os
import re
import typing
from collections.abc import MutableMapping

from opentelemetry.propagators.textmap import Getter, Setter


def _normalize_key(key: str) -> str:
    result = re.sub(r"[^A-Za-z0-9_]", "_", key).upper()
    if result and result[0].isdigit():
        result = "_" + result
    return result


class EnvironmentGetter(Getter[typing.Mapping[str, str]]):
    """Getter implementation for extracting context and baggage from environment variables.

    EnvironmentGetter creates a normalized lookup from the current environment
    variables at initialization time and provides simple data access without validation.

    Per the OpenTelemetry specification, environment variables are treated as immutable
    within a process. For environments where context-carrying environment variables
    change between logical requests (e.g., AWS Lambda's _X_AMZN_TRACE_ID), create a
    new EnvironmentGetter instance at the start of each request.

    Example usage:
        getter = EnvironmentGetter()
        traceparent = getter.get({}, "traceparent")
    """

    def __init__(self):
        # Create a normalized lookup from current environment
        # Per spec: "creates an in-memory copy of the current environment variables"
        self.carrier: dict[str, str] = {
            _normalize_key(k): v for k, v in os.environ.items()
        }

    def get(
        self, carrier: typing.Mapping[str, str], key: str
    ) -> list[str] | None:
        """Get a value from the environment carrier for the given key.

        Args:
            carrier: Not used; maintained for interface compatibility with Getter[CarrierT]
            key: The key to look up (case-insensitive)

        Returns:
            A list with a single string value if the key exists, None otherwise.
        """
        val = self.carrier.get(_normalize_key(key))
        if val is None:
            return None
        return [val]

    def keys(self, carrier: typing.Mapping[str, str]) -> list[str]:
        """Get all keys from the environment carrier.

        Args:
            carrier: Not used; maintained for interface compatibility with Getter[CarrierT]

        Returns:
            List of all environment variable keys (normalized).
        """
        return list(self.carrier.keys())


class EnvironmentSetter(Setter[MutableMapping[str, str]]):
    """Setter implementation for building environment variable dictionaries.

    EnvironmentSetter builds a dictionary of environment variables that
    can be passed to utilities like subprocess.run()

    Example usage:
        setter = EnvironmentSetter()
        env_vars = {}
        setter.set(env_vars, "traceparent", "00-trace-id-span-id-01")
        subprocess.run(myCommand, env=env_vars)
    """

    def set(
        self, carrier: MutableMapping[str, str], key: str, value: str
    ) -> None:
        """Set a value in the carrier dictionary for the given key.

        Args:
            carrier: Dictionary to store environment variables
            key: The key to set (will be converted to uppercase)
            value: The value to set
        """
        carrier[_normalize_key(key)] = value
