from enum import Enum


# NB: These keys are placeholders and subject to change
class TraceMetadataKey:
    INPUTS = "mlflow.traceInputs"
    OUTPUTS = "mlflow.traceOutputs"
    SOURCE_RUN = "mlflow.sourceRun"
    MODEL_ID = "mlflow.modelId"
    # Trace size statistics including total size, number of spans, and max span size
    SIZE_STATS = "mlflow.trace.sizeStats"
    # Aggregated token usage information in a single trace, stored as a dumped JSON string.
    TOKEN_USAGE = "mlflow.trace.tokenUsage"
    # Store the user ID/name of the application request. Do not confuse this with mlflow.user
    # tag, which stores "who created the trace" i.e. developer or system name.
    TRACE_USER = "mlflow.trace.user"
    # Store the session ID of the application request.
    TRACE_SESSION = "mlflow.trace.session"

    # Total size of the trace in bytes. Deprecated, use SIZE_STATS instead.
    SIZE_BYTES = "mlflow.trace.sizeBytes"


class TraceTagKey:
    TRACE_NAME = "mlflow.traceName"
    EVAL_REQUEST_ID = "mlflow.eval.requestId"
    SPANS_LOCATION = "mlflow.trace.spansLocation"
    # Store the source scorer name that generated the trace. This tag is used to determine if a
    # trace is generated by a scorer or prediction during evaluation and filter out the scorer
    # traces in the UI. This supposed to be immutable, but we use tag because we can only set this
    # after the scorer is executed, which is not possible with trace metadata.
    SOURCE_SCORER_NAME = "mlflow.trace.sourceScorer"
    # Store a list of linked prompt versions in JSON format
    # Structure: [{"name": "prompt_name", "version": "version"}]
    LINKED_PROMPTS = "mlflow.linkedPrompts"


class TokenUsageKey:
    """Key for the token usage information in the `mlflow.chat.tokenUsage` span attribute."""

    INPUT_TOKENS = "input_tokens"
    OUTPUT_TOKENS = "output_tokens"
    TOTAL_TOKENS = "total_tokens"

    @classmethod
    def all_keys(cls):
        return [cls.INPUT_TOKENS, cls.OUTPUT_TOKENS, cls.TOTAL_TOKENS]


class TraceSizeStatsKey:
    TOTAL_SIZE_BYTES = "total_size_bytes"
    NUM_SPANS = "num_spans"
    MAX_SPAN_SIZE_BYTES = "max"
    P25_SPAN_SIZE_BYTES = "p25"
    P50_SPAN_SIZE_BYTES = "p50"
    P75_SPAN_SIZE_BYTES = "p75"


# A set of reserved attribute keys
class SpanAttributeKey:
    EXPERIMENT_ID = "mlflow.experimentId"
    REQUEST_ID = "mlflow.traceRequestId"
    INPUTS = "mlflow.spanInputs"
    OUTPUTS = "mlflow.spanOutputs"
    SPAN_TYPE = "mlflow.spanType"
    FUNCTION_NAME = "mlflow.spanFunctionName"
    START_TIME_NS = "mlflow.spanStartTimeNs"
    CHAT_TOOLS = "mlflow.chat.tools"
    # This attribute is used to store token usage information from LLM responses.
    # Stored in {"input_tokens": int, "output_tokens": int, "total_tokens": int} format.
    CHAT_USAGE = "mlflow.chat.tokenUsage"
    # This attribute indicates which flavor/format generated the LLM span. This is
    # used by downstream (e.g., UI) to determine the message format for parsing.
    MESSAGE_FORMAT = "mlflow.message.format"
    # This attribute is used to populate `intermediate_outputs` property of a trace data
    # representing intermediate outputs of the trace. This attribute is not empty only on
    # the root span of a trace created by the `mlflow.log_trace` API. The `intermediate_outputs`
    # property of the normal trace is generated by the outputs of non-root spans.
    INTERMEDIATE_OUTPUTS = "mlflow.trace.intermediate_outputs"
    # This attribute is used to store prompt version information when load_prompt is called
    # within an active span. Stored as a JSON list of {"name": "...", "version": "..."} objects,
    # same format as LINKED_PROMPTS_TAG_KEY in traces.
    LINKED_PROMPTS = "mlflow.linkedPrompts"


class AssessmentMetadataKey:
    # When the assessment is generated by an eval run, log the run ID here.
    SOURCE_RUN_ID = "mlflow.assessment.sourceRunId"
    # Total LLM cost spent for generating the feedback (llm-as-a-judge).
    JUDGE_COST = "mlflow.assessment.judgeCost"
    # When the scorer generates a trace for assessment scoring, log the trace ID here.
    SCORER_TRACE_ID = "mlflow.assessment.scorerTraceId"


# All storage backends are guaranteed to support request_metadata key/value up to 250 characters
MAX_CHARS_IN_TRACE_INFO_METADATA = 250
# All storage backends are guaranteed to support tag keys up to 250 characters,
# values up to 4096 characters
MAX_CHARS_IN_TRACE_INFO_TAGS_KEY = 250
MAX_CHARS_IN_TRACE_INFO_TAGS_VALUE = 4096
TRUNCATION_SUFFIX = "..."

TRACE_REQUEST_RESPONSE_PREVIEW_MAX_LENGTH_DBX = 10000
TRACE_REQUEST_RESPONSE_PREVIEW_MAX_LENGTH_OSS = 1000

# Trace request ID must have the prefix "tr-" appended to the OpenTelemetry trace ID
TRACE_REQUEST_ID_PREFIX = "tr-"
# Trace ID V4 format starts with "trace:/" in the format of "trace:/<location>/<trace_id>"
TRACE_ID_V4_PREFIX = "trace:/"

# Schema version of traces and spans.
TRACE_SCHEMA_VERSION = 3

# Key for the trace schema version in the trace. This key is also used in
# Databricks model serving to be careful when modifying it.
TRACE_SCHEMA_VERSION_KEY = "mlflow.trace_schema.version"


STREAM_CHUNK_EVENT_NAME_FORMAT = "mlflow.chunk.item.{index}"
STREAM_CHUNK_EVENT_VALUE_KEY = "mlflow.chunk.value"


# Key for Databricks model serving options to return the trace in the response
DATABRICKS_OPTIONS_KEY = "databricks_options"
RETURN_TRACE_OPTION_KEY = "return_trace"
DATABRICKS_OUTPUT_KEY = "databricks_output"

# Assessment constants
ASSESSMENT_ID_PREFIX = "a-"


# Maximum number of seconds to retry getting a trace from the v4 endpoint.
# V4 traces have some delay in propagation after the log_spans call returns success response.
# To make sure get_trace API does not fail due to this delay, we retry up to a reasonable timeout.
# Setting 15 seconds because the initial version of the backend is known to have 1~5 seconds delay.
GET_TRACE_V4_RETRY_TIMEOUT_SECONDS = 15


# The location of the spans in the trace.
# This is used to determine where the spans are stored when exporting.
class SpansLocation(str, Enum):
    TRACKING_STORE = "TRACKING_STORE"
    ARTIFACT_REPO = "ARTIFACT_REPO"


# Path to the notebook trace renderer directory
TRACE_RENDERER_ASSET_PATH = "/static-files/lib/notebook-trace-renderer"
