# Copyright 2024 Klira AI
#
# 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.

"""KliraOTLPSpanExporter — OTLP exporter with Bearer token auth and PHI scanning.

This exporter wraps the standard OTLPSpanExporter to add:
1. Bearer token authentication via the Klira API key
2. PHI/PII de-identification via Presidio when anonymization=True
   (protobuf mutation — see Phase 0.4 contract in phi_pipeline.py)
"""

from __future__ import annotations

import fnmatch
import logging
from typing import TYPE_CHECKING, Any, Sequence

from opentelemetry.sdk.trace import ReadableSpan
from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult
from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter

if TYPE_CHECKING:
    from klira.sdk.contracts.phi_pipeline import PhiMethod

logger = logging.getLogger(__name__)


class KliraOTLPSpanExporter(SpanExporter):
    """OTLP span exporter with Klira-specific authentication and PHI scanning.

    Wraps the standard OTLPSpanExporter, adding:
    - Bearer token auth via Authorization header
    - PHI de-identification via protobuf mutation in export()

    When anonymization=True, the export path:
    1. Serializes spans to protobuf (mutable ExportTraceServiceRequest).
    2. Walks all span attributes for PHI_SCANNABLE_ATTRIBUTES keys.
    3. Calls Presidio to scan and de-identify text values.
    4. Overwrites StringValue in-place in the protobuf.
    5. Appends klira.phi.* attributes per span.
    6. Sends the modified protobuf.

    Failure mode: fail-closed on init (raises if Presidio unavailable when
    anonymization=True). Fail-open per-span at export time — if Presidio fails
    on a span, it is exported unchanged with a WARNING log.
    klira.phi.detected absent = platform flags for manual review.
    """

    def __init__(
        self,
        endpoint: str,
        api_key: str,
        anonymization: bool = False,
        phi_method: PhiMethod | None = None,
        phi_export_entity_details: bool = True,
    ) -> None:
        self._endpoint = endpoint
        self._api_key = api_key
        self._anonymization = anonymization
        self._phi_export_entity_details = phi_export_entity_details
        self._scanner: Any = None

        if not anonymization and not phi_export_entity_details:
            logger.debug(
                "phi_export_entity_details=False has no effect when "
                "anonymization is disabled"
            )

        if anonymization:
            from klira.sdk.contracts.phi_pipeline import PhiMethod as _PhiMethod

            method = phi_method or _PhiMethod.REDACT
            from klira.sdk.healthcare.phi import PhiScanner

            self._scanner = PhiScanner(method=method)
            # Eagerly initialize — ImportError/RuntimeError propagates to caller.
            # When a user explicitly requests anonymization=True, silent downgrade
            # would be a HIPAA violation.
            self._scanner._ensure_initialized()

        self._headers = {
            "Authorization": f"Bearer {api_key}",
        }
        self._inner = OTLPSpanExporter(
            endpoint=endpoint,
            headers=self._headers,
        )

    def __repr__(self) -> str:
        return (
            f"KliraOTLPSpanExporter("
            f"endpoint={self._endpoint!r}, "
            f"anonymization={self._anonymization!r}, "
            f"phi_export_entity_details={self._phi_export_entity_details!r})"
        )

    def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
        """Export spans to the Klira platform.

        When anonymization is enabled, serializes to protobuf, runs Presidio
        on mutable protobuf StringValue fields, appends klira.phi.* attributes,
        then sends modified protobuf. Otherwise delegates directly.
        """
        if not spans:
            return SpanExportResult.SUCCESS

        if not self._anonymization or self._scanner is None:
            return self._inner.export(spans)

        return self._export_with_phi_scan(spans)

    def _export_with_phi_scan(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
        """Export spans with PHI scanning via protobuf mutation.

        Serializes spans to protobuf, walks attributes, runs Presidio,
        overwrites StringValue in-place, appends klira.phi.* attributes.
        """
        from opentelemetry.exporter.otlp.proto.common._internal.trace_encoder import (
            encode_spans,
        )
        from opentelemetry.proto.common.v1.common_pb2 import (  # type: ignore[import-untyped]
            AnyValue,
            KeyValue,
        )

        from klira.sdk.contracts.phi_pipeline import (
            PHI_SCANNABLE_ATTRIBUTES,
            PHI_SCANNABLE_PATTERNS,
        )

        try:
            # Step 1: Serialize spans to mutable protobuf
            pb_request = encode_spans(spans)

            # Step 2: Walk and mutate protobuf
            for resource_spans in pb_request.resource_spans:
                for scope_spans in resource_spans.scope_spans:
                    for pb_span in scope_spans.spans:
                        self._scan_protobuf_span(
                            pb_span,
                            PHI_SCANNABLE_ATTRIBUTES,
                            PHI_SCANNABLE_PATTERNS,
                            KeyValue,
                            AnyValue,
                        )

            # Step 3: Send modified protobuf bytes via the inner exporter's HTTP path
            serialized_bytes = pb_request.SerializePartialToString()
            resp = self._inner._export(serialized_bytes)  # type: ignore[attr-defined]
            if resp.ok:
                return SpanExportResult.SUCCESS
            body_preview = (resp.text or "")[:200]
            logger.error(
                "PHI-scanned export failed: %s %s", resp.status_code, body_preview
            )
            return SpanExportResult.FAILURE

        except Exception:
            # Fail-open: if PHI scanning fails, export unchanged
            logger.warning("PHI scanning failed, exporting unchanged", exc_info=True)
            return self._inner.export(spans)

    @staticmethod
    def _should_scan_attr(
        key: str,
        scannable_attrs: tuple[str, ...],
        scannable_patterns: tuple[str, ...],
    ) -> bool:
        """Check if an attribute key should be scanned for PHI."""
        if key in scannable_attrs:
            return True
        return any(fnmatch.fnmatch(key, p) for p in scannable_patterns)

    def _scan_protobuf_span(
        self,
        pb_span: Any,
        scannable_attrs: tuple[str, ...],
        scannable_patterns: tuple[str, ...],
        KeyValue: type,
        AnyValue: type,
    ) -> None:
        """Scan and mutate a single protobuf span's attributes.

        Scans string values for PHI, de-identifies in-place, and appends
        klira.phi.* attributes to the span.
        """
        total_entities = 0
        all_entity_types: set[str] = set()

        for attr in pb_span.attributes:
            if self._should_scan_attr(attr.key, scannable_attrs, scannable_patterns):
                value = attr.value
                if value.HasField("string_value") and value.string_value:
                    try:
                        scan_result, raw_results = self._scanner.scan(value.string_value)
                        if scan_result.detected:
                            total_entities += scan_result.entity_count
                            all_entity_types.update(scan_result.entity_types)
                            # Overwrite StringValue in-place, reusing analyzer results
                            deidentified = self._scanner.deidentify(
                                value.string_value, analyzer_results=raw_results,
                            )
                            attr.value.string_value = deidentified
                    except Exception:
                        logger.warning(
                            "PHI scan failed for attr %s", attr.key, exc_info=True
                        )

        # Append klira.phi.* attributes
        detected = total_entities > 0
        pb_span.attributes.append(
            KeyValue(
                key="klira.phi.detected",
                value=AnyValue(bool_value=detected),
            )
        )
        if detected and self._phi_export_entity_details:
            pb_span.attributes.append(
                KeyValue(
                    key="klira.phi.entity_count",
                    value=AnyValue(int_value=total_entities),
                )
            )
            types_str = ",".join(sorted(all_entity_types))
            pb_span.attributes.append(
                KeyValue(
                    key="klira.phi.entity_types",
                    value=AnyValue(string_value=types_str),
                )
            )

    def shutdown(self) -> None:
        """Shut down the exporter."""
        self._inner.shutdown()

    def force_flush(self, timeout_millis: int = 30000) -> bool:
        """Force flush pending exports."""
        return self._inner.force_flush(timeout_millis)
