# Copyright Coralogix Ltd.
#
# 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.

import os
import logging
from typing import Any
from dataclasses import dataclass

from opentelemetry import trace
from opentelemetry.sdk.resources import SERVICE_NAME, Resource
from opentelemetry.sdk.trace import TracerProvider, SpanLimits
from opentelemetry.sdk.trace.export import (
    BatchSpanProcessor,
    SimpleSpanProcessor,
    SpanProcessor,
)
from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter

from llm_tracekit.core._config import enable_capture_content

logger = logging.getLogger(__name__)


@dataclass
class ExportConfig:
    endpoint: str | None = None
    headers: dict[str, Any] | None = None


def generate_exporter_config(
    coralogix_token: str | None,
    coralogix_endpoint: str | None,
    application_name: str | None,
    subsystem_name: str | None,
) -> ExportConfig:
    if coralogix_token is None:
        coralogix_token = os.environ.get("CX_TOKEN")
    if coralogix_endpoint is None:
        coralogix_endpoint = os.environ.get("CX_ENDPOINT")
    if application_name is None:
        application_name = os.environ.get("CX_APPLICATION_NAME", "Unknown")
    if subsystem_name is None:
        subsystem_name = os.environ.get("CX_SUBSYSTEM_NAME", "Unknown")

    headers = {
        "authorization": f"Bearer {coralogix_token}",
        "cx-application-name": application_name,
        "cx-subsystem-name": subsystem_name,
    }

    return ExportConfig(endpoint=coralogix_endpoint, headers=headers)


def setup_export_to_coralogix(
    service_name: str,
    coralogix_token: str | None = None,
    coralogix_endpoint: str | None = None,
    application_name: str | None = None,
    subsystem_name: str | None = None,
    use_batch_processor: bool = True,
    capture_content: bool = True,
    processors: list[SpanProcessor] | None = None,
    span_attribute_count_limit: int = 512,
):
    """
    Setup OpenAI spans to be exported to Coralogix.

    Args:
        service_name: The service name.
        coralogix_token: The Coralogix token. Defaults to os.environ["CX_TOKEN"]
        coralogix_endpoint: The Coralogix endpoint. Defaults to os.environ["CX_ENDPOINT"]
        application_name: The Coralogix application name. Defaults to os.environ["CX_APPLICATION_NAME"]
        subsystem_name: The Coralogix subsystem name. Defaults to os.environ["CX_SUBSYSTEM_NAME"]
        use_batch_processor: Whether to use a batch processor or a simple processor.
        capture_content: Whether to capture the content of the messages.
        processors: Optional list of SpanProcessor instances to add to the tracer provider before the exporter processor.
        span_attribute_count_limit: The maximum number of span attributes.
    """

    if capture_content:
        enable_capture_content()

    exporter_config = generate_exporter_config(
        coralogix_token=coralogix_token,
        coralogix_endpoint=coralogix_endpoint,
        application_name=application_name,
        subsystem_name=subsystem_name,
    )

    env_limit_raw = os.environ.get("OTEL_SPAN_ATTRIBUTE_COUNT_LIMIT")
    effective_limit = span_attribute_count_limit
    if env_limit_raw is not None:
        try:
            env_limit = int(env_limit_raw)
        except ValueError:
            logger.warning(
                "Invalid OTEL_SPAN_ATTRIBUTE_COUNT_LIMIT=%r; using configured limit %s",
                env_limit_raw,
                span_attribute_count_limit,
            )
        else:
            effective_limit = env_limit
            if env_limit < span_attribute_count_limit:
                logger.warning(
                    "OTEL_SPAN_ATTRIBUTE_COUNT_LIMIT=%s is lower than requested %s; spans may drop attributes.",
                    env_limit,
                    span_attribute_count_limit,
                )

    span_attribute_limit = SpanLimits(
        max_span_attributes=effective_limit,
    )

    # set up a tracer provider to send spans to coralogix.
    tracer_provider = TracerProvider(
        resource=Resource.create({SERVICE_NAME: service_name}),
        span_limits=span_attribute_limit,
    )

    # add any custom span processors before configuring the exporter processor
    if processors:
        for span_processor in processors:
            tracer_provider.add_span_processor(span_processor)

    # set up an OTLP exporter to send spans to coralogix directly.
    exporter = OTLPSpanExporter(
        endpoint=exporter_config.endpoint, headers=exporter_config.headers
    )

    # set up a span processor to send spans to the exporter
    span_processor = (
        BatchSpanProcessor(exporter)
        if use_batch_processor
        else SimpleSpanProcessor(exporter)
    )

    # add the span processor to the tracer provider
    tracer_provider.add_span_processor(span_processor)
    trace.set_tracer_provider(tracer_provider)
