# 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.

"""Clinical context functions for healthcare applications."""

from __future__ import annotations

from opentelemetry import trace


def set_patient_context(
    patient_id: str,
    encounter_id: str | None = None,
) -> None:
    """Set patient context on the current span.

    Args:
        patient_id: Patient identifier.
        encounter_id: Optional encounter identifier.
    """
    span = trace.get_current_span()
    span.set_attribute("klira.healthcare.patient_id", patient_id)
    if encounter_id is not None:
        span.set_attribute("klira.healthcare.encounter_id", encounter_id)


def set_clinical_context(
    department: str | None = None,
    specialty: str | None = None,
    clinical_domain: str | None = None,
) -> None:
    """Set clinical context on the current span.

    Args:
        department: Hospital department.
        specialty: Medical specialty.
        clinical_domain: Clinical domain (e.g., "cardiology").
    """
    span = trace.get_current_span()
    if department is not None:
        span.set_attribute("klira.healthcare.department", department)
    if specialty is not None:
        span.set_attribute("klira.healthcare.specialty", specialty)
    if clinical_domain is not None:
        span.set_attribute("klira.healthcare.clinical_domain", clinical_domain)


def set_interaction_modality(modality: str) -> None:
    """Set the interaction modality on the current span.

    Args:
        modality: Interaction type (e.g., "chat", "voice", "portal").
    """
    span = trace.get_current_span()
    span.set_attribute("klira.healthcare.interaction_modality", modality)
