# Copyright (c) Metis. All rights reserved.

"""OTel span builder for JSON push-to-trace.

Creates OpenTelemetry spans from resolved mapping data. Isolated from
the mapping engine -- takes already-validated, already-resolved field
dicts and produces real spans with proper parent-child linking.
"""

from __future__ import annotations

import json
import logging
import time
from typing import Any, Dict, List, Optional

from opentelemetry import context as otel_context
from opentelemetry import trace as otel_trace

from . import semconv
from .mapping import ChildMapping, coerce_timestamp, resolve_path
from .push_models import SpanPreview, ValidationIssue

logger = logging.getLogger(__name__)

# Default timestamp: "now" in nanoseconds
_NANOS_PER_SECOND = 1_000_000_000


def _now_ns() -> int:
    return int(time.time() * _NANOS_PER_SECOND)


def _safe_json(value: Any) -> str:
    """Serialize a value to JSON string for span attributes."""
    if isinstance(value, str):
        return value
    try:
        return json.dumps(value)
    except (TypeError, ValueError):
        logger.warning(
            "Value of type %s is not JSON-serializable, using str() fallback: %s",
            type(value).__name__,
            repr(value)[:100],
        )
        return str(value)


# =============================================================================
# Span Building
# =============================================================================


def build_trace_spans(
    data: dict,
    fields: Dict[str, Any],
    child_mapping: Optional[ChildMapping],
    tracer: Optional[otel_trace.Tracer] = None,
    extra_attributes: Optional[Dict[str, Any]] = None,
    timestamp_unit: str = "auto",
) -> tuple[str, int, List[ValidationIssue]]:
    """Create OTel spans for a trace and its children.

    Args:
        data: The original JSON data dict (for resolving child fields).
        fields: Merged trace-level fields (name, input, output, etc.).
        child_mapping: Optional mapping for children observations.
        tracer: OTel tracer to use. Defaults to the global mantisdk tracer.
        extra_attributes: Additional OTel span attributes to set on the root span.

    Returns:
        (trace_id, observation_count, warnings) -- the hex trace ID, number of
        child observations created, and any runtime warnings.
    """
    if tracer is None:
        tracer = otel_trace.get_tracer("mantisdk.tracing.push")

    start_ns = coerce_timestamp(fields.get("timestamp"), unit=timestamp_unit) or _now_ns()

    # Build root span attributes
    root_attrs = _build_trace_attributes(fields)
    if extra_attributes:
        root_attrs.update(extra_attributes)

    root_span = tracer.start_span(
        name=fields.get("name") or "push",
        start_time=start_ns,
        attributes=root_attrs,
    )

    observation_count = 0
    runtime_warnings: List[ValidationIssue] = []

    # Build children
    if child_mapping is not None:
        found, children_data = resolve_path(data, child_mapping.path)
        if found and isinstance(children_data, list):
            root_ctx = otel_trace.set_span_in_context(root_span)
            observation_count, child_warnings = _build_children(
                children_data,
                child_mapping,
                tracer,
                root_ctx,
                start_ns,
                timestamp_unit=timestamp_unit,
            )
            runtime_warnings.extend(child_warnings)

    # End root span
    end_ns = _now_ns()
    root_span.end(end_time=end_ns)

    # Extract trace ID
    span_context = root_span.get_span_context()
    trace_id = format(span_context.trace_id, "032x")

    return trace_id, observation_count, runtime_warnings


def _build_trace_attributes(fields: Dict[str, Any]) -> Dict[str, Any]:
    """Build OTel span attributes for the root trace span."""
    attrs: Dict[str, Any] = {}

    if fields.get("name"):
        attrs[semconv.TRACE_NAME] = fields["name"]
    if fields.get("input") is not None:
        attrs[semconv.TRACE_INPUT] = _safe_json(fields["input"])
    if fields.get("output") is not None:
        attrs[semconv.TRACE_OUTPUT] = _safe_json(fields["output"])
    if fields.get("metadata") is not None:
        attrs[semconv.TRACE_METADATA] = _safe_json(fields["metadata"])
    if fields.get("tags"):
        attrs[semconv.TRACE_TAGS] = _safe_json(fields["tags"])
    if fields.get("user_id"):
        attrs[semconv.USER_ID] = fields["user_id"]
    if fields.get("session_id"):
        attrs[semconv.SESSION_ID] = fields["session_id"]
    if fields.get("run_id"):
        attrs[semconv.RUN_ID] = fields["run_id"]
    if fields.get("environment"):
        attrs[semconv.ENVIRONMENT] = fields["environment"]

    return attrs


def _build_children(
    children_data: List[dict],
    child_mapping: ChildMapping,
    tracer: otel_trace.Tracer,
    parent_ctx: otel_context.Context,
    parent_start_ns: int,
    timestamp_unit: str = "auto",
) -> tuple[int, List[ValidationIssue]]:
    """Create child observation spans under the parent context.

    Returns (observation_count, warnings).
    """
    count = 0
    warnings: List[ValidationIssue] = []

    for i, child_dict in enumerate(children_data):
        if not isinstance(child_dict, dict):
            warnings.append(
                ValidationIssue(
                    field="children",
                    path=f"{child_mapping.path}[{i}]",
                    issue="invalid_child",
                    message=f"Children item at index {i} is {type(child_dict).__name__}, expected dict. Skipped.",
                )
            )
            continue

        child_fields, field_warnings = _extract_child_fields(child_dict, child_mapping, timestamp_unit=timestamp_unit)
        warnings.extend(field_warnings)

        child_start = child_fields.get("start_time_ns") or (parent_start_ns + i)
        child_end = child_fields.get("end_time_ns")

        child_attrs = _build_observation_attributes(child_fields)

        child_span = tracer.start_span(
            name=child_fields.get("name") or f"observation-{i}",
            context=parent_ctx,
            start_time=child_start,
            attributes=child_attrs,
        )

        # Recurse into nested children if the mapping defines them
        nested_count = 0
        if child_mapping.children is not None:
            found, nested_data = resolve_path(child_dict, child_mapping.children.path)
            if found and isinstance(nested_data, list):
                child_ctx = otel_trace.set_span_in_context(child_span)
                nested_count, nested_warnings = _build_children(
                    nested_data,
                    child_mapping.children,
                    tracer,
                    child_ctx,
                    child_start,
                    timestamp_unit=timestamp_unit,
                )
                warnings.extend(nested_warnings)

        child_span.end(end_time=child_end or _now_ns())
        count += 1 + nested_count

    return count, warnings


def _extract_child_fields(
    child_dict: dict,
    mapping: ChildMapping,
    timestamp_unit: str = "auto",
) -> tuple[Dict[str, Any], List[ValidationIssue]]:
    """Extract mapped fields from a single child dict.

    Returns (fields, warnings).
    """
    fields: Dict[str, Any] = {}
    warnings: List[ValidationIssue] = []

    simple_fields = [
        "name",
        "input",
        "output",
        "type",
        "model",
        "metadata",
        "level",
        "usage",
        "cost",
        "model_parameters",
        "status_message",
        "prompt_name",
        "prompt_version",
        "attributes",
    ]
    for field in simple_fields:
        dot_path = getattr(mapping, field, None)
        if dot_path is not None:
            found, value = resolve_path(child_dict, dot_path)
            if found:
                fields[field] = value

    # Timestamps need coercion
    for ts_field in ("start_time", "end_time"):
        dot_path = getattr(mapping, ts_field, None)
        if dot_path is not None:
            found, value = resolve_path(child_dict, dot_path)
            if found:
                coerced = coerce_timestamp(value, unit=timestamp_unit)
                if coerced is None and value is not None:
                    warnings.append(
                        ValidationIssue(
                            field=f"children.{ts_field}",
                            path=dot_path,
                            issue="timestamp_coercion_failed",
                            message=f"Could not parse timestamp value: {value!r}. Using fallback.",
                        )
                    )
                fields[f"{ts_field}_ns"] = coerced

    return fields, warnings


def _build_observation_attributes(fields: Dict[str, Any]) -> Dict[str, Any]:
    """Build OTel span attributes for an observation span."""
    attrs: Dict[str, Any] = {}

    if fields.get("input") is not None:
        attrs[semconv.OBSERVATION_INPUT] = _safe_json(fields["input"])
    if fields.get("output") is not None:
        attrs[semconv.OBSERVATION_OUTPUT] = _safe_json(fields["output"])
    if fields.get("type"):
        attrs[semconv.OBSERVATION_TYPE] = fields["type"].upper()
    if fields.get("model"):
        attrs[semconv.MODEL_NAME] = fields["model"]
        # Presence of model implies GENERATION type (unless explicitly set)
        if not fields.get("type"):
            attrs[semconv.OBSERVATION_TYPE] = "GENERATION"
    if fields.get("metadata") is not None:
        attrs[semconv.OBSERVATION_METADATA] = _safe_json(fields["metadata"])
    if fields.get("level"):
        attrs[semconv.OBSERVATION_LEVEL] = fields["level"].upper()
    if fields.get("status_message"):
        attrs[semconv.OBSERVATION_STATUS_MESSAGE] = fields["status_message"]
    if fields.get("model_parameters") is not None:
        attrs[semconv.MODEL_PARAMETERS] = _safe_json(fields["model_parameters"])

    # Usage and cost use the langfuse-compatible attributes
    if fields.get("usage") is not None:
        attrs["insight.observation.usage_details"] = _safe_json(fields["usage"])
    if fields.get("cost") is not None:
        attrs["insight.observation.cost_details"] = _safe_json(fields["cost"])

    # Prompt reference
    if fields.get("prompt_name"):
        attrs["insight.observation.prompt.name"] = fields["prompt_name"]
    if fields.get("prompt_version") is not None:
        attrs["insight.observation.prompt.version"] = fields["prompt_version"]

    # Extra attributes escape hatch
    extra = fields.get("attributes")
    if isinstance(extra, dict):
        attrs.update(extra)

    return attrs


# =============================================================================
# Preview Building (no OTel dependency)
# =============================================================================


def build_trace_preview(
    data: dict,
    fields: Dict[str, Any],
    child_mapping: Optional[ChildMapping],
) -> List[SpanPreview]:
    """Build SpanPreview objects for children (no OTel spans created).

    Returns list of SpanPreview for the children.
    """
    if child_mapping is None:
        return []

    found, children_data = resolve_path(data, child_mapping.path)
    if not found or not isinstance(children_data, list):
        return []

    previews = []
    for i, child_dict in enumerate(children_data):
        if not isinstance(child_dict, dict):
            continue

        child_fields, _ = _extract_child_fields(child_dict, child_mapping)

        nested = []
        if child_mapping.children is not None:
            nested = build_trace_preview(child_dict, child_fields, child_mapping.children)

        previews.append(
            SpanPreview(
                name=child_fields.get("name") or f"observation-{i}",
                type=child_fields.get("type"),
                input=child_fields.get("input"),
                output=child_fields.get("output"),
                model=child_fields.get("model"),
                metadata=child_fields.get("metadata"),
                start_time=str(child_fields.get("start_time_ns")) if child_fields.get("start_time_ns") else None,
                end_time=str(child_fields.get("end_time_ns")) if child_fields.get("end_time_ns") else None,
                children=nested,
            )
        )

    return previews
