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

"""Trace schema contract — the single source of truth for all Klira spans.

Every span the SDK produces is defined here. Tests validate captured spans
against these definitions. Platform and SDK share this contract for
attribute naming and hierarchy rules.

Schema version is bumped when span definitions change in a
backward-incompatible way (new required attributes, renamed spans, etc.).
"""

from __future__ import annotations

from dataclasses import dataclass
from enum import Enum
from typing import Any

SCHEMA_VERSION = "0.6.1"


class AttributeType(Enum):
    STRING = "string"
    INT = "int"
    FLOAT = "float"
    BOOL = "bool"
    STRING_LIST = "string_list"
    INT_LIST = "int_list"


@dataclass(frozen=True)
class SpanAttribute:
    """Definition of a single span attribute."""

    name: str
    attr_type: AttributeType
    description: str
    allowed_values: tuple[str, ...] | None = None
    validation_pattern: str | None = None


@dataclass(frozen=True)
class SpanDefinition:
    """Definition of a single span type the SDK can produce."""

    name: str
    description: str
    parent_constraints: tuple[str, ...]
    required_attributes: tuple[SpanAttribute, ...]
    optional_attributes: tuple[SpanAttribute, ...] = ()
    is_async: bool = False
    name_template: str | None = None


# ---------------------------------------------------------------------------
# Attribute Registry — every attribute the SDK sets on any span
# ---------------------------------------------------------------------------

ATTRIBUTE_REGISTRY: dict[str, SpanAttribute] = {}


def _attr(
    name: str,
    attr_type: AttributeType,
    description: str,
    allowed_values: tuple[str, ...] | None = None,
    validation_pattern: str | None = None,
) -> SpanAttribute:
    a = SpanAttribute(
        name=name,
        attr_type=attr_type,
        description=description,
        allowed_values=allowed_values,
        validation_pattern=validation_pattern,
    )
    ATTRIBUTE_REGISTRY[name] = a
    return a


# --- Klira core attributes (present on most spans) ---
ATTR_ENTITY_TYPE = _attr(
    "klira.entity_type",
    AttributeType.STRING,
    "Span entity type",
    allowed_values=(
        "user_message",
        "workflow",
        "agent",
        "task",
        "tool",
        "llm",
        "guardrails",
        "compliance",
        "clinical_decision",
        "human_escalation",
        "agent_handoff",
        "safety_check",
        "rag_retrieval",
        "eval_test_case",
    ),
)
ATTR_ENTITY_NAME = _attr(
    "klira.entity_name", AttributeType.STRING, "Human-readable entity name"
)
ATTR_USER_ID = _attr(
    "klira.user_id", AttributeType.STRING, "User identifier for trace correlation"
)
ATTR_CONVERSATION_ID = _attr(
    "klira.conversation_id", AttributeType.STRING, "Conversation identifier"
)
ATTR_MESSAGE_ID = _attr("klira.message_id", AttributeType.STRING, "Message identifier")
ATTR_DURATION_MS = _attr(
    "klira.duration_ms",
    AttributeType.FLOAT,
    "Span duration in milliseconds (set by processor)",
)
ATTR_FRAMEWORK = _attr(
    "klira.framework", AttributeType.STRING, "Framework adapter name"
)
ATTR_OUTPUT = _attr(
    "klira.output",
    AttributeType.STRING,
    "Decorator output capture (truncated to 5000 chars)",
)
ATTR_INPUT = _attr(
    "klira.input",
    AttributeType.STRING,
    "User input text captured from function arguments (truncated to 10000 chars)",
)
# --- FHIR ---
ATTR_FHIR_RESOURCE_TYPE = _attr(
    "klira.fhir.resource_type", AttributeType.STRING, "FHIR resource type on tool spans"
)

# --- PHI per-span attributes (set by exporter before export) ---
ATTR_PHI_DETECTED = _attr(
    "klira.phi.detected",
    AttributeType.BOOL,
    "Whether PHI was detected in this span's attributes",
)
ATTR_PHI_ENTITY_COUNT = _attr(
    "klira.phi.entity_count", AttributeType.INT, "Number of PHI entities detected"
)
ATTR_PHI_ENTITY_TYPES = _attr(
    "klira.phi.entity_types",
    AttributeType.STRING_LIST,
    "Types of PHI entities detected",
)

# --- GenAI semantic convention attributes (on LLM spans) ---
ATTR_GEN_AI_SYSTEM = _attr(
    "gen_ai.system", AttributeType.STRING, "LLM provider (lowercase)"
)
ATTR_GEN_AI_REQUEST_MODEL = _attr(
    "gen_ai.request.model",
    AttributeType.STRING,
    "Requested model name (no provider prefix)",
)
ATTR_GEN_AI_RESPONSE_MODEL = _attr(
    "gen_ai.response.model", AttributeType.STRING, "Actual model used in response"
)
ATTR_GEN_AI_INPUT_TOKENS = _attr(
    "gen_ai.usage.input_tokens", AttributeType.INT, "Input token count"
)
ATTR_GEN_AI_OUTPUT_TOKENS = _attr(
    "gen_ai.usage.output_tokens", AttributeType.INT, "Output token count"
)
ATTR_GEN_AI_FINISH_REASONS = _attr(
    "gen_ai.response.finish_reasons",
    AttributeType.STRING_LIST,
    "Finish reasons from LLM",
)
ATTR_GEN_AI_PROMPT = _attr(
    "gen_ai.prompt", AttributeType.STRING, "LLM prompt content (truncated)"
)
ATTR_GEN_AI_PROMPT_ORIGINAL = _attr(
    "gen_ai.prompt.original",
    AttributeType.STRING,
    "Original prompt before augmentation",
)

# --- Guardrails attributes ---
ATTR_GUARDRAILS_DIRECTION = _attr(
    "klira.compliance.direction",
    AttributeType.STRING,
    "Guardrails check direction",
    allowed_values=("inbound", "outbound"),
)
ATTR_GUARDRAILS_DECISION_ALLOWED = _attr(
    "klira.compliance.decision.allowed",
    AttributeType.BOOL,
    "Whether the guardrails decision allowed the content",
)
ATTR_GUARDRAILS_DECISION_ACTION = _attr(
    "klira.compliance.decision.action",
    AttributeType.STRING,
    "Guardrails decision action",
    allowed_values=("block", "allow"),
)
ATTR_GUARDRAILS_DECISION_LAYER = _attr(
    "klira.compliance.decision.layer",
    AttributeType.STRING,
    "Evaluation layer that made the decision",
)
ATTR_GUARDRAILS_AUGMENTATION_APPLIED = _attr(
    "klira.guardrails.augmentation_applied",
    AttributeType.BOOL,
    "Whether prompt augmentation was applied",
)
ATTR_GUARDRAILS_GUIDELINES_COUNT = _attr(
    "klira.guardrails.guidelines_count", AttributeType.INT, "Number of guidelines injected"
)
ATTR_GUARDRAILS_POLICY_COUNT = _attr(
    "klira.guardrails.policy_count", AttributeType.INT, "Number of policies evaluated"
)
ATTR_GUARDRAILS_LATENCY_MS = _attr(
    "klira.guardrails.latency_ms",
    AttributeType.FLOAT,
    "Guardrails processing latency in ms",
)
ATTR_GUARDRAILS_POLICIES_INJECTED = _attr(
    "klira.llm.guardrails.policies_injected", AttributeType.INT, "Number of policies injected"
)
ATTR_GUARDRAILS_AUGMENTED = _attr(
    "klira.llm.guardrails.augmented", AttributeType.BOOL, "Whether LLM call was augmented"
)
# --- Guardrails evaluate-level attributes (on klira.guardrails.evaluate spans) ---
ATTR_GUARDRAILS_EVALUATE_DIRECTION = _attr(
    "klira.guardrails.direction",
    AttributeType.STRING,
    "Guardrails evaluation direction (input or output)",
    allowed_values=("input", "output"),
)
ATTR_GUARDRAILS_EVALUATE_DECISION = _attr(
    "klira.guardrails.decision",
    AttributeType.STRING,
    "Guardrails evaluation decision action",
    allowed_values=("allow", "block"),
)
ATTR_GUARDRAILS_EVALUATE_ALLOWED = _attr(
    "klira.guardrails.allowed",
    AttributeType.BOOL,
    "Whether guardrails evaluation allowed the content",
)

# --- Compliance attributes ---
ATTR_COMPLIANCE_POLICY_IDS = _attr(
    "klira.compliance.policy_ids",
    AttributeType.STRING,
    "Comma-separated matched policy IDs",
)
ATTR_COMPLIANCE_REGULATION = _attr(
    "klira.compliance.regulation", AttributeType.STRING, "Regulation tag from policy metadata"
)
ATTR_COMPLIANCE_AUDIT_REQUIRED = _attr(
    "klira.compliance.audit_required", AttributeType.BOOL, "Whether audit is required"
)
ATTR_COMPLIANCE_HUMAN_REVIEW_REQUIRED = _attr(
    "klira.compliance.human_review_required",
    AttributeType.BOOL,
    "Whether human review is required",
)


# --- Healthcare context attributes (set by healthcare/context.py) ---
ATTR_HEALTHCARE_PATIENT_ID = _attr(
    "klira.healthcare.patient_id", AttributeType.STRING, "Patient identifier"
)
ATTR_HEALTHCARE_ENCOUNTER_ID = _attr(
    "klira.healthcare.encounter_id", AttributeType.STRING, "Encounter identifier"
)
ATTR_HEALTHCARE_DEPARTMENT = _attr(
    "klira.healthcare.department", AttributeType.STRING, "Hospital department"
)
ATTR_HEALTHCARE_SPECIALTY = _attr(
    "klira.healthcare.specialty", AttributeType.STRING, "Medical specialty"
)
ATTR_HEALTHCARE_CLINICAL_DOMAIN = _attr(
    "klira.healthcare.clinical_domain", AttributeType.STRING, "Clinical domain"
)
ATTR_HEALTHCARE_INTERACTION_MODALITY = _attr(
    "klira.healthcare.interaction_modality",
    AttributeType.STRING,
    "Interaction modality (chat, voice, etc.)",
)

# --- Clinical logging attributes (set by healthcare/logging.py) ---
ATTR_CLINICAL_DECISION = _attr(
    "klira.clinical.decision", AttributeType.STRING, "Clinical decision text"
)
ATTR_CLINICAL_REASONING = _attr(
    "klira.clinical.reasoning", AttributeType.STRING, "Clinical decision reasoning"
)
ATTR_CLINICAL_CONFIDENCE = _attr(
    "klira.clinical.confidence", AttributeType.FLOAT, "Clinical decision confidence"
)
ATTR_CLINICAL_URGENCY = _attr(
    "klira.clinical.urgency", AttributeType.STRING, "Clinical urgency level"
)
ATTR_CLINICAL_ESCALATION_REASON = _attr(
    "klira.clinical.escalation_reason", AttributeType.STRING, "Escalation reason"
)
ATTR_CLINICAL_FROM_AGENT = _attr(
    "klira.clinical.from_agent", AttributeType.STRING, "Source agent in handoff"
)
ATTR_CLINICAL_TO_AGENT = _attr(
    "klira.clinical.to_agent", AttributeType.STRING, "Target agent in handoff"
)
ATTR_CLINICAL_HANDOFF_REASON = _attr(
    "klira.clinical.handoff_reason", AttributeType.STRING, "Handoff reason"
)
ATTR_CLINICAL_CHECK_TYPE = _attr(
    "klira.clinical.check_type", AttributeType.STRING, "Safety check type"
)
ATTR_CLINICAL_CHECK_PASSED = _attr(
    "klira.clinical.check_passed", AttributeType.BOOL, "Whether safety check passed"
)
ATTR_CLINICAL_CHECK_DETAILS = _attr(
    "klira.clinical.check_details", AttributeType.STRING, "Safety check details"
)
ATTR_CLINICAL_RAG_SOURCE = _attr(
    "klira.clinical.rag_source", AttributeType.STRING, "RAG retrieval source"
)
ATTR_CLINICAL_RAG_QUERY = _attr(
    "klira.clinical.rag_query", AttributeType.STRING, "RAG retrieval query"
)
ATTR_CLINICAL_RAG_RESULT_COUNT = _attr(
    "klira.clinical.rag_result_count", AttributeType.INT, "RAG retrieval result count"
)

# --- Evals attributes ---
ATTR_EVALS_RUN_ID = _attr(
    "klira.evals.evals_run", AttributeType.STRING, "Eval run identifier"
)
ATTR_EVALS_DATASET_ID = _attr(
    "klira.evals.dataset_id", AttributeType.STRING, "Eval dataset identifier"
)
ATTR_EVALS_TEST_CASE_ID = _attr(
    "klira.evals.test_case_id", AttributeType.STRING, "Eval test case identifier"
)
ATTR_EVALS_TEST_CASE_INPUT = _attr(
    "klira.evals.test_case_input", AttributeType.STRING, "Eval test case input"
)
ATTR_EVALS_EXPECTED_OUTPUT = _attr(
    "klira.evals.expected_output", AttributeType.STRING, "Expected output for the test case"
)


# ---------------------------------------------------------------------------
# Span Definitions — every span the SDK can produce
# ---------------------------------------------------------------------------

# Helper constants for parent constraints
_ROOT_OR_USER_MESSAGE = ("root", "klira.user.message")
_CHILD_OF_USER_MESSAGE = ("klira.user.message",)
_CHILD_OF_ANY_KLIRA = (
    "klira.user.message",
    "klira.workflow.*",
    "klira.agent.*",
    "klira.task.*",
    "klira.tool.*",
)
_CHILD_OF_GUARDRAILS = (
    "klira.guardrails.input",
    "klira.guardrails.output",
)

SPAN_USER_MESSAGE = SpanDefinition(
    name="klira.user.message",
    description="Root span for a user interaction. One per user message.",
    parent_constraints=("root",),
    required_attributes=(
        ATTR_ENTITY_TYPE,
        ATTR_USER_ID,
        ATTR_CONVERSATION_ID,
        ATTR_MESSAGE_ID,
    ),
    optional_attributes=(ATTR_EVALS_RUN_ID, ATTR_ENTITY_NAME, ATTR_DURATION_MS, ATTR_INPUT),
)

SPAN_WORKFLOW = SpanDefinition(
    name="klira.workflow.{name}",
    name_template="klira.workflow.{name}",
    description="Workflow execution span.",
    parent_constraints=_ROOT_OR_USER_MESSAGE,
    required_attributes=(ATTR_ENTITY_TYPE, ATTR_ENTITY_NAME),
    optional_attributes=(
        ATTR_OUTPUT,
        ATTR_DURATION_MS,
        ATTR_FRAMEWORK,
        ATTR_INPUT,
        ATTR_USER_ID,
        ATTR_CONVERSATION_ID,
    ),
)

SPAN_AGENT = SpanDefinition(
    name="klira.agent.{name}",
    name_template="klira.agent.{name}",
    description="Agent execution span.",
    parent_constraints=_CHILD_OF_ANY_KLIRA,
    required_attributes=(ATTR_ENTITY_TYPE, ATTR_ENTITY_NAME),
    optional_attributes=(
        ATTR_OUTPUT,
        ATTR_DURATION_MS,
        ATTR_FRAMEWORK,
        ATTR_USER_ID,
        ATTR_CONVERSATION_ID,
    ),
)

SPAN_TASK = SpanDefinition(
    name="klira.task.{name}",
    name_template="klira.task.{name}",
    description="Task execution span.",
    parent_constraints=_CHILD_OF_ANY_KLIRA,
    required_attributes=(ATTR_ENTITY_TYPE, ATTR_ENTITY_NAME),
    optional_attributes=(
        ATTR_OUTPUT,
        ATTR_DURATION_MS,
        ATTR_FRAMEWORK,
        ATTR_USER_ID,
        ATTR_CONVERSATION_ID,
    ),
)

SPAN_TOOL = SpanDefinition(
    name="klira.tool.{name}",
    name_template="klira.tool.{name}",
    description="Tool execution span.",
    parent_constraints=_CHILD_OF_ANY_KLIRA,
    required_attributes=(ATTR_ENTITY_TYPE, ATTR_ENTITY_NAME),
    optional_attributes=(
        ATTR_OUTPUT,
        ATTR_DURATION_MS,
        ATTR_FRAMEWORK,
        ATTR_FHIR_RESOURCE_TYPE,
        ATTR_USER_ID,
        ATTR_CONVERSATION_ID,
    ),
)

SPAN_LLM = SpanDefinition(
    name="klira.llm.{provider}",
    name_template="klira.llm.{provider}",
    description="LLM API call wrapper span.",
    # Note: ATTR_FRAMEWORK is intentionally absent here. klira.framework is set
    # at the decorator level (workflow/agent/task/tool spans), not on LLM spans.
    # LLM adapters never set this attribute. See PROD-464 / PR #20.
    parent_constraints=_CHILD_OF_ANY_KLIRA,
    required_attributes=(
        ATTR_ENTITY_TYPE,
        ATTR_GEN_AI_SYSTEM,
        ATTR_GEN_AI_REQUEST_MODEL,
    ),
    optional_attributes=(
        ATTR_GEN_AI_RESPONSE_MODEL,
        ATTR_GEN_AI_INPUT_TOKENS,
        ATTR_GEN_AI_OUTPUT_TOKENS,
        ATTR_GEN_AI_FINISH_REASONS,
        ATTR_GEN_AI_PROMPT,
        ATTR_GEN_AI_PROMPT_ORIGINAL,
        ATTR_GUARDRAILS_AUGMENTED,
        ATTR_GUARDRAILS_POLICIES_INJECTED,
        ATTR_DURATION_MS,
        ATTR_ENTITY_NAME,
    ),
)

SPAN_GUARDRAILS_INPUT = SpanDefinition(
    name="klira.guardrails.input",
    description="Inbound guardrails check.",
    parent_constraints=_CHILD_OF_ANY_KLIRA,
    required_attributes=(ATTR_ENTITY_TYPE, ATTR_GUARDRAILS_DIRECTION),
    optional_attributes=(
        ATTR_ENTITY_NAME,
        ATTR_GUARDRAILS_LATENCY_MS,
        ATTR_GUARDRAILS_POLICY_COUNT,
        ATTR_GUARDRAILS_EVALUATE_DECISION,
        ATTR_GUARDRAILS_EVALUATE_ALLOWED,
        ATTR_DURATION_MS,
    ),
)

SPAN_GUARDRAILS_OUTPUT = SpanDefinition(
    name="klira.guardrails.output",
    description="Outbound guardrails check.",
    parent_constraints=_CHILD_OF_GUARDRAILS + _CHILD_OF_ANY_KLIRA,
    required_attributes=(ATTR_ENTITY_TYPE, ATTR_GUARDRAILS_DIRECTION),
    optional_attributes=(
        ATTR_ENTITY_NAME,
        ATTR_GUARDRAILS_LATENCY_MS,
        ATTR_GUARDRAILS_POLICY_COUNT,
        ATTR_GUARDRAILS_EVALUATE_DECISION,
        ATTR_GUARDRAILS_EVALUATE_ALLOWED,
        ATTR_DURATION_MS,
    ),
)


SPAN_COMPLIANCE_ALLOWED = SpanDefinition(
    name="klira.compliance.allowed",
    description="Compliance decision — content allowed.",
    parent_constraints=_CHILD_OF_GUARDRAILS + _CHILD_OF_ANY_KLIRA,
    required_attributes=(
        ATTR_ENTITY_TYPE,
        ATTR_GUARDRAILS_DIRECTION,
        ATTR_GUARDRAILS_DECISION_ALLOWED,
        ATTR_GUARDRAILS_DECISION_ACTION,
        ATTR_GUARDRAILS_DECISION_LAYER,
        ATTR_GUARDRAILS_AUGMENTATION_APPLIED,
        ATTR_GUARDRAILS_GUIDELINES_COUNT,
        ATTR_COMPLIANCE_AUDIT_REQUIRED,
        ATTR_COMPLIANCE_HUMAN_REVIEW_REQUIRED,
    ),
    optional_attributes=(
        ATTR_COMPLIANCE_POLICY_IDS,
        ATTR_COMPLIANCE_REGULATION,
        ATTR_DURATION_MS,
    ),
    is_async=True,
)

SPAN_COMPLIANCE_BLOCKED = SpanDefinition(
    name="klira.compliance.blocked",
    description="Compliance decision — content blocked.",
    parent_constraints=_CHILD_OF_GUARDRAILS + _CHILD_OF_ANY_KLIRA,
    required_attributes=(
        ATTR_ENTITY_TYPE,
        ATTR_GUARDRAILS_DIRECTION,
        ATTR_GUARDRAILS_DECISION_ALLOWED,
        ATTR_GUARDRAILS_DECISION_ACTION,
        ATTR_GUARDRAILS_DECISION_LAYER,
        ATTR_GUARDRAILS_AUGMENTATION_APPLIED,
        ATTR_GUARDRAILS_GUIDELINES_COUNT,
        ATTR_COMPLIANCE_AUDIT_REQUIRED,
        ATTR_COMPLIANCE_HUMAN_REVIEW_REQUIRED,
    ),
    optional_attributes=(
        ATTR_COMPLIANCE_POLICY_IDS,
        ATTR_COMPLIANCE_REGULATION,
        ATTR_DURATION_MS,
    ),
    is_async=True,
)

SPAN_COMPLIANCE_AUGMENTED = SpanDefinition(
    name="klira.compliance.augmented",
    description="Compliance decision — content augmented with guidelines.",
    parent_constraints=_CHILD_OF_GUARDRAILS + _CHILD_OF_ANY_KLIRA,
    required_attributes=(
        ATTR_ENTITY_TYPE,
        ATTR_GUARDRAILS_DIRECTION,
        ATTR_GUARDRAILS_DECISION_ALLOWED,
        ATTR_GUARDRAILS_DECISION_ACTION,
        ATTR_GUARDRAILS_DECISION_LAYER,
        ATTR_GUARDRAILS_AUGMENTATION_APPLIED,
        ATTR_GUARDRAILS_GUIDELINES_COUNT,
        ATTR_COMPLIANCE_AUDIT_REQUIRED,
        ATTR_COMPLIANCE_HUMAN_REVIEW_REQUIRED,
    ),
    optional_attributes=(
        ATTR_COMPLIANCE_POLICY_IDS,
        ATTR_COMPLIANCE_REGULATION,
        ATTR_DURATION_MS,
    ),
    is_async=True,
)

SPAN_COMPLIANCE_LLM_FALLBACK = SpanDefinition(
    name="klira.compliance.llm_fallback",
    description="Compliance decision — LLM fallback evaluation.",
    parent_constraints=_CHILD_OF_GUARDRAILS + _CHILD_OF_ANY_KLIRA,
    required_attributes=(
        ATTR_ENTITY_TYPE,
        ATTR_GUARDRAILS_DIRECTION,
        ATTR_GUARDRAILS_DECISION_ALLOWED,
        ATTR_GUARDRAILS_DECISION_ACTION,
        ATTR_GUARDRAILS_DECISION_LAYER,
        ATTR_GUARDRAILS_AUGMENTATION_APPLIED,
        ATTR_GUARDRAILS_GUIDELINES_COUNT,
        ATTR_COMPLIANCE_AUDIT_REQUIRED,
        ATTR_COMPLIANCE_HUMAN_REVIEW_REQUIRED,
    ),
    optional_attributes=(
        ATTR_COMPLIANCE_POLICY_IDS,
        ATTR_COMPLIANCE_REGULATION,
        ATTR_DURATION_MS,
    ),
    is_async=True,
)

# --- Clinical / Healthcare spans ---
SPAN_CLINICAL_DECISION = SpanDefinition(
    name="klira.clinical.decision",
    description="Clinical decision logging span.",
    parent_constraints=_CHILD_OF_ANY_KLIRA,
    required_attributes=(ATTR_ENTITY_TYPE, ATTR_ENTITY_NAME),
    optional_attributes=(
        ATTR_CLINICAL_DECISION,
        ATTR_CLINICAL_REASONING,
        ATTR_CLINICAL_CONFIDENCE,
        ATTR_HEALTHCARE_SPECIALTY,
        ATTR_CLINICAL_URGENCY,
        ATTR_DURATION_MS,
    ),
)

SPAN_CLINICAL_ESCALATION = SpanDefinition(
    name="klira.clinical.escalation",
    description="Human escalation logging span.",
    parent_constraints=_CHILD_OF_ANY_KLIRA,
    required_attributes=(ATTR_ENTITY_TYPE, ATTR_ENTITY_NAME),
    optional_attributes=(
        ATTR_CLINICAL_ESCALATION_REASON,
        ATTR_CLINICAL_URGENCY,
        ATTR_DURATION_MS,
    ),
)

SPAN_CLINICAL_HANDOFF = SpanDefinition(
    name="klira.clinical.handoff",
    description="Agent handoff logging span.",
    parent_constraints=_CHILD_OF_ANY_KLIRA,
    required_attributes=(ATTR_ENTITY_TYPE, ATTR_ENTITY_NAME),
    optional_attributes=(
        ATTR_CLINICAL_FROM_AGENT,
        ATTR_CLINICAL_TO_AGENT,
        ATTR_CLINICAL_HANDOFF_REASON,
        ATTR_DURATION_MS,
    ),
)

SPAN_CLINICAL_SAFETY_CHECK = SpanDefinition(
    name="klira.clinical.safety_check",
    description="Safety check logging span.",
    parent_constraints=_CHILD_OF_ANY_KLIRA,
    required_attributes=(ATTR_ENTITY_TYPE, ATTR_ENTITY_NAME),
    optional_attributes=(
        ATTR_CLINICAL_CHECK_TYPE,
        ATTR_CLINICAL_CHECK_PASSED,
        ATTR_CLINICAL_CHECK_DETAILS,
        ATTR_DURATION_MS,
    ),
)

SPAN_CLINICAL_RAG_RETRIEVAL = SpanDefinition(
    name="klira.clinical.rag_retrieval",
    description="RAG retrieval logging span.",
    parent_constraints=_CHILD_OF_ANY_KLIRA,
    required_attributes=(ATTR_ENTITY_TYPE, ATTR_ENTITY_NAME),
    optional_attributes=(
        ATTR_CLINICAL_RAG_SOURCE,
        ATTR_CLINICAL_RAG_QUERY,
        ATTR_CLINICAL_RAG_RESULT_COUNT,
        ATTR_DURATION_MS,
    ),
)

SPAN_AGENT_HANDOFF = SpanDefinition(
    name="klira.agent.handoff",
    description="Agent handoff logging span.",
    parent_constraints=_CHILD_OF_ANY_KLIRA,
    required_attributes=(ATTR_ENTITY_TYPE, ATTR_ENTITY_NAME),
    optional_attributes=(ATTR_DURATION_MS,),
)

SPAN_RAG_RETRIEVAL = SpanDefinition(
    name="klira.rag.retrieval",
    description="RAG retrieval logging span.",
    parent_constraints=_CHILD_OF_ANY_KLIRA,
    required_attributes=(ATTR_ENTITY_TYPE, ATTR_ENTITY_NAME),
    optional_attributes=(ATTR_DURATION_MS,),
)

SPAN_AUDIT_REGULATORY = SpanDefinition(
    name="klira.audit.regulatory",
    description="Regulatory audit logging span.",
    parent_constraints=_CHILD_OF_ANY_KLIRA,
    required_attributes=(ATTR_ENTITY_TYPE, ATTR_ENTITY_NAME),
    optional_attributes=(ATTR_COMPLIANCE_REGULATION, ATTR_DURATION_MS),
)

SPAN_EVALS_TEST_CASE = SpanDefinition(
    name="klira.evals.test_case",
    description="Eval test case execution span.",
    parent_constraints=("root", "klira.user.message"),
    required_attributes=(
        ATTR_ENTITY_TYPE,
        ATTR_EVALS_RUN_ID,
        ATTR_EVALS_TEST_CASE_ID,
    ),
    optional_attributes=(
        ATTR_EVALS_DATASET_ID,
        ATTR_EVALS_TEST_CASE_INPUT,
        ATTR_EVALS_EXPECTED_OUTPUT,
        ATTR_DURATION_MS,
    ),
)


# ---------------------------------------------------------------------------
# Span definition registry — for programmatic lookup
# ---------------------------------------------------------------------------

SPAN_DEFINITIONS: dict[str, SpanDefinition] = {
    sd.name: sd
    for sd in [
        SPAN_USER_MESSAGE,
        SPAN_WORKFLOW,
        SPAN_AGENT,
        SPAN_TASK,
        SPAN_TOOL,
        SPAN_LLM,
        SPAN_GUARDRAILS_INPUT,
        SPAN_GUARDRAILS_OUTPUT,
        SPAN_COMPLIANCE_ALLOWED,
        SPAN_COMPLIANCE_BLOCKED,
        SPAN_COMPLIANCE_AUGMENTED,
        SPAN_COMPLIANCE_LLM_FALLBACK,
        SPAN_CLINICAL_DECISION,
        SPAN_CLINICAL_ESCALATION,
        SPAN_CLINICAL_HANDOFF,
        SPAN_CLINICAL_SAFETY_CHECK,
        SPAN_CLINICAL_RAG_RETRIEVAL,
        SPAN_AGENT_HANDOFF,
        SPAN_RAG_RETRIEVAL,
        SPAN_AUDIT_REGULATORY,
        SPAN_EVALS_TEST_CASE,
    ]
}


# ---------------------------------------------------------------------------
# Validation helpers — used by contract tests
# ---------------------------------------------------------------------------


@dataclass
class SpanValidationError:
    """A single validation error for a span."""

    span_name: str
    error: str


def _match_span_definition(span_name: str) -> SpanDefinition | None:
    """Find the SpanDefinition matching a concrete span name.

    Handles template matching: 'klira.workflow.my_flow' matches 'klira.workflow.{name}'.
    """
    if span_name in SPAN_DEFINITIONS:
        return SPAN_DEFINITIONS[span_name]

    for defn in SPAN_DEFINITIONS.values():
        if defn.name_template is None:
            continue
        prefix = defn.name_template.split("{")[0]
        if span_name.startswith(prefix):
            return defn

    return None


def validate_span(
    span_name: str,
    attributes: dict[str, Any],
    parent_span_name: str | None = None,
) -> list[SpanValidationError]:
    """Validate a span against the trace schema contract.

    Args:
        span_name: The concrete span name (e.g., 'klira.workflow.my_flow').
        attributes: The span's attributes as a dict.
        parent_span_name: The parent span's name, or None for root spans.

    Returns:
        List of validation errors. Empty list means the span is valid.
    """
    errors: list[SpanValidationError] = []

    defn = _match_span_definition(span_name)
    if defn is None:
        errors.append(SpanValidationError(span_name, f"Unknown span name: {span_name}"))
        return errors

    # Check required attributes
    for attr in defn.required_attributes:
        if attr.name not in attributes:
            errors.append(
                SpanValidationError(
                    span_name,
                    f"Missing required attribute: {attr.name}",
                )
            )
            continue

        value = attributes[attr.name]

        # Type check
        if attr.attr_type == AttributeType.STRING and not isinstance(value, str):
            errors.append(
                SpanValidationError(
                    span_name,
                    f"Attribute {attr.name} must be string, got {type(value).__name__}",
                )
            )
        elif attr.attr_type == AttributeType.INT and not isinstance(value, int):
            errors.append(
                SpanValidationError(
                    span_name,
                    f"Attribute {attr.name} must be int, got {type(value).__name__}",
                )
            )
        elif attr.attr_type == AttributeType.FLOAT and not isinstance(
            value, (int, float)
        ):
            errors.append(
                SpanValidationError(
                    span_name,
                    f"Attribute {attr.name} must be float, got {type(value).__name__}",
                )
            )
        elif attr.attr_type == AttributeType.BOOL and not isinstance(value, bool):
            errors.append(
                SpanValidationError(
                    span_name,
                    f"Attribute {attr.name} must be bool, got {type(value).__name__}",
                )
            )
        elif attr.attr_type == AttributeType.STRING_LIST:
            if not isinstance(value, (list, tuple)):
                errors.append(
                    SpanValidationError(
                        span_name,
                        f"Attribute {attr.name} must be list, got {type(value).__name__}",
                    )
                )

        # Allowed values check
        if attr.allowed_values is not None and isinstance(value, str):
            if value not in attr.allowed_values:
                errors.append(
                    SpanValidationError(
                        span_name,
                        f"Attribute {attr.name} value '{value}' not in allowed values: {attr.allowed_values}",
                    )
                )

    # Check parent constraints
    non_root_constraints = tuple(c for c in defn.parent_constraints if c != "root")
    allows_root = "root" in defn.parent_constraints

    if parent_span_name is not None:
        if not non_root_constraints:
            # Only "root" is allowed — having a parent is invalid
            errors.append(
                SpanValidationError(
                    span_name,
                    f"Invalid parent '{parent_span_name}'. Allowed: {defn.parent_constraints}",
                )
            )
        else:
            parent_matches = False
            for constraint in non_root_constraints:
                if constraint.endswith(".*"):
                    prefix = constraint[:-2]
                    if parent_span_name.startswith(prefix + "."):
                        parent_matches = True
                        break
                elif parent_span_name == constraint or parent_span_name.startswith(
                    constraint.split("{")[0]
                ):
                    parent_matches = True
                    break

            if not parent_matches:
                errors.append(
                    SpanValidationError(
                        span_name,
                        f"Invalid parent '{parent_span_name}'. Allowed: {defn.parent_constraints}",
                    )
                )
    elif not allows_root:
        errors.append(
            SpanValidationError(
                span_name,
                f"Span requires a parent. Allowed: {defn.parent_constraints}",
            )
        )

    return errors
