# Copyright 2023-2026 Airbus, CS Group
#
# 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.

"""OpenTelemetry utility"""

import inspect
import json
import os
import pkgutil
import sys
from collections.abc import Iterator
from typing import Any

import fastapi
import opentelemetry.instrumentation
import requests
from opentelemetry import trace
from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
from opentelemetry.instrumentation.asyncio import AsyncioInstrumentor
from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor
from opentelemetry.instrumentation.httpx import HTTPXClientInstrumentor
from opentelemetry.instrumentation.instrumentor import BaseInstrumentor  # type: ignore
from opentelemetry.instrumentation.requests import RequestsInstrumentor
from opentelemetry.sdk.resources import Resource  # type: ignore
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor
from opentelemetry.trace import Status, StatusCode
from opentelemetry.trace.span import NonRecordingSpan, Span, SpanContext, TraceFlags
from opentelemetry.util._decorator import _agnosticcontextmanager

from rs_dpr_service.utils.logging import Logging
from rs_dpr_service.utils.settings import env_bool

default_logger = Logging.default(__name__)

FROM_PYTEST = False


# Show details of http headers and body/content in tempo/grafana ?
# Don't store results in global variables because the env var values can change
# after this module was loaded.
def trace_headers():
    """Trace request headers ?"""
    return env_bool("OTEL_PYTHON_REQUESTS_TRACE_HEADERS", default=False)


def trace_body():
    """Trace request bodies and response contents ?"""
    return env_bool("OTEL_PYTHON_REQUESTS_TRACE_BODY", default=False)


def parse_data(data) -> str:
    """Convert data to a string representation"""

    if not data:
        return ""

    # Try to decode bytes
    if isinstance(data, bytes):
        data = data.decode("utf-8")

    # Try to convert to a dict
    try:
        data = dict(data)
    except Exception:  # pylint: disable=broad-exception-caught # nosec
        pass

    # Or to parse to a dict
    try:
        data = json.loads(data)
    except Exception:  # pylint: disable=broad-exception-caught # nosec
        pass

    # If we have a dict
    if isinstance(data, dict):

        # Decode bytes
        data = {
            key.decode("utf-8") if isinstance(key, bytes) else key: (
                value.decode("utf-8") if isinstance(value, bytes) else value
            )
            for key, value in data.items()
        }

        # Convert to strings
        data = {str(key): str(value) for key, value in data.items()}

        # Apply json formatting
        data = json.dumps(data, indent=2)

    return data or ""


def requests_hook(span: Span, request: requests.PreparedRequest, response: requests.Response | None = None):
    """
    Callback function invoked by RequestsInstrumentor. It implements the hooks:

      - request_hook: invoked right after a span is created.
      - response_hook: invoked right before the span has finished processing a response.

    See: https://opentelemetry-python-contrib.readthedocs.io/en/latest/instrumentation/requests/requests.html
    """
    if not (span and span.is_recording()):
        return

    # Copy this attribute by adding a '_' prefix to the name,
    # so it appears at the top in the grafana UI, it's more readable
    span.set_attribute("_url", span.attributes.get("http.url"))  # type: ignore

    if trace_headers():
        span.set_attribute("http.request.headers", parse_data(request.headers))
        if response:
            span.set_attribute("http.response.headers", parse_data(response.headers))

    if trace_body():
        span.set_attribute("http.request.body", parse_data(request.body))
        if response:
            span.set_attribute("http.response.content", parse_data(response.content))


def fastapi_hook(span: Span, scope: dict[str, Any], message=None):
    """
    Callback function invoked by FastAPIInstrumentor. It implements the hooks:

      - server_request_hook: called with the server span and ASGI scope object for every incoming request.
      - client_request_hook: called with the internal span, and ASGI scope and event which are sent as dictionaries
                             for when the method receive is called.
      - client_response_hook: called with the internal span, and ASGI scope and event which are sent as dictionaries
                              for when the method send is called.

    See: https://opentelemetry-python-contrib.readthedocs.io/en/latest/instrumentation/fastapi/fastapi.html
    """
    if not (span and span.is_recording()):
        return

    # Copy this attribute by adding a '_' prefix to the name,
    # so it appears at the top in the grafana UI, it's more readable
    span.set_attribute("_path", str(scope.get("path")))

    if trace_headers():
        span.set_attribute("http.scope.headers", parse_data(scope.get("headers")))
        if message:
            span.set_attribute("http.message.headers", parse_data(message.get("headers")))

    if trace_body() and message:
        span.set_attribute("http.message.body", parse_data(message.get("body")))


def init_traces(app: fastapi.FastAPI | None, service_name: str, logger=None):  # pylint: disable=too-many-branches
    """
    Init instrumentation of OpenTelemetry traces.

    Args:
        app (fastapi.FastAPI): FastAPI application
        service_name (str): service name
        logger: non-default logger to user
    """

    # See: https://github.com/softwarebloat/python-tracing-demo/tree/main

    logger = logger or default_logger

    # Don't call this line from pytest because it causes errors:
    # Transient error StatusCode.UNAVAILABLE encountered while exporting metrics to localhost:4317, retrying in ..s.
    if not FROM_PYTEST:
        tempo_endpoint = os.getenv("TEMPO_ENDPOINT")
        if not tempo_endpoint:
            logger.warning("'TEMPO_ENDPOINT' variable is missing, cannot initialize OpenTelemetry")
            return

        # TODO: to avoid errors in local mode:
        # Transient error StatusCode.UNAVAILABLE encountered while exporting metrics to localhost:4317, retrying in ..s.
        #
        # The below line does not work either but at least we have less error messages.
        # See: https://pforge-exchange2.astrium.eads.net/jira/browse/RSPY-221?focusedId=162092&
        # page=com.atlassian.jira.plugin.system.issuetabpanels%3Acomment-tabpanel#comment-162092
        #
        # Now we have a single line error, which is less worst:
        # Failed to export metrics to tempo:4317, error code: StatusCode.UNIMPLEMENTED
        os.environ["OTEL_EXPORTER_OTLP_ENDPOINT"] = tempo_endpoint

    otel_resource = Resource(attributes={"service.name": service_name})
    otel_tracer = TracerProvider(resource=otel_resource)
    trace.set_tracer_provider(otel_tracer)

    if not FROM_PYTEST:
        otel_tracer.add_span_processor(BatchSpanProcessor(OTLPSpanExporter(endpoint=tempo_endpoint)))

    if app:
        if trace_headers() or trace_body():
            FastAPIInstrumentor.instrument_app(
                app,
                tracer_provider=otel_tracer,
                server_request_hook=fastapi_hook,
                client_request_hook=fastapi_hook,
                client_response_hook=fastapi_hook,
            )
        else:
            FastAPIInstrumentor.instrument_app(app, tracer_provider=otel_tracer)
        # logger.debug(f"OpenTelemetry instrumentation of 'fastapi.FastAPIInstrumentor'")

    # Instrument all the dependencies under opentelemetry.instrumentation.*
    # NOTE: we need 'poetry run opentelemetry-bootstrap -a install' to install these.

    package = opentelemetry.instrumentation
    prefix = package.__name__ + "."
    classes = set()

    # We need an empty PYTHONPATH if the env var is missing
    os.environ["PYTHONPATH"] = os.getenv("PYTHONPATH", "")

    # Recursively find all package modules
    for _, module_str, _ in pkgutil.walk_packages(path=package.__path__, prefix=prefix, onerror=None):

        # Don't instrument these modules, they have errors, maybe we should see why
        if module_str in [
            "opentelemetry.instrumentation.tortoiseorm",
            "opentelemetry.instrumentation.auto_instrumentation.sitecustomize",
        ]:
            continue

        # Import and find all module classes
        __import__(module_str)
        for _, _class in inspect.getmembers(sys.modules[module_str]):
            if (not inspect.isclass(_class)) or (_class in classes):
                continue

            # Save the class (classes are found several times when imported by other modules)
            classes.add(_class)

            # Don't instrument these classes, they have errors, maybe we should see why
            if _class in [AsyncioInstrumentor, BaseInstrumentor, HTTPXClientInstrumentor]:
                continue

            # If the "instrument" method exists, call it
            _instrument = getattr(_class, "instrument", None)
            if callable(_instrument):
                _class_instance = _class()
                if _class_instance.is_instrumented_by_opentelemetry:
                    continue
                # name = f"{module_str}.{_class.__name__}".removeprefix(prefix)
                # logger.debug(f"OpenTelemetry instrumentation of {name!r}")

                # Handle specific hooks
                if _class == RequestsInstrumentor and (trace_headers() or trace_body()):
                    _class_instance.instrument(
                        tracer_provider=otel_tracer,
                        request_hook=requests_hook,
                        response_hook=requests_hook,
                    )
                elif _class == FastAPIInstrumentor and (trace_headers() or trace_body()):
                    _class_instance.instrument(
                        tracer_provider=otel_tracer,
                        server_request_hook=fastapi_hook,
                        client_request_hook=fastapi_hook,
                        client_response_hook=fastapi_hook,
                    )

                # General case (no hooks)
                else:
                    _class_instance.instrument(tracer_provider=otel_tracer)


@_agnosticcontextmanager
def start_span(
    instrumenting_module_name: str,
    name: str,
    span_context: SpanContext | None = None,
) -> Iterator[Span]:
    """
    Context manager for creating a new main or child OpenTelemetry span and set it
    as the current span in this tracer's context.

    Args:
        instrumenting_module_name: Caller module name, just pass __name__
        name: The name of the span to be created (use a custom name)
        span_context: Parent span context. Only to create a child span.

    Yields:
        The newly-created span.
    """
    tracer = trace.get_tracer(instrumenting_module_name)

    # Create a main span
    if not span_context:
        with tracer.start_as_current_span(name) as span:
            yield span

    # Create a child span
    else:
        main_span_context = SpanContext(
            trace_id=span_context.trace_id,
            span_id=span_context.span_id,
            is_remote=True,
            trace_flags=TraceFlags(TraceFlags.SAMPLED),
        )
        main_span = NonRecordingSpan(main_span_context)
        with trace.use_span(main_span):  # pylint: disable=not-context-manager
            # Optionnaly, we could use the main span instead of creating
            # a new one, to be discussed.
            with tracer.start_as_current_span(name) as span:
                yield span


def record_error(span: Span, e: Exception):
    """
    Record an exception in the span and mark it as ERROR if recording.

    Args:
        span (Span): The tracing span.
        e (Exception): The exception to record.
    """
    if span.is_recording() and e is not None:
        span.record_exception(e)
        span.set_status(Status(StatusCode.ERROR, str(e)))
