# Copyright (c) Metis. All rights reserved.

"""Pydantic models for the JSON push-to-trace feature.

Defines the data contracts for mapping configuration, validation output,
inspection results, and push results. All models are JSON-serializable
and export clean JSON Schema via model_json_schema().
"""

from __future__ import annotations

from typing import Any, Dict, List, Optional

from pydantic import BaseModel

# =============================================================================
# Mapping Configuration
# =============================================================================


class ChildMapping(BaseModel):
    """Maps fields from each item in a children array to observation attributes.

    Each field is a dot-path string that resolves against child dict items.
    Supports recursive nesting via the ``children`` field.
    """

    path: str
    """Dot-path to the children array in the parent dict."""

    name: Optional[str] = None
    input: Optional[str] = None
    output: Optional[str] = None
    type: Optional[str] = None
    model: Optional[str] = None
    metadata: Optional[str] = None
    level: Optional[str] = None
    start_time: Optional[str] = None
    end_time: Optional[str] = None
    usage: Optional[str] = None
    cost: Optional[str] = None
    model_parameters: Optional[str] = None
    status_message: Optional[str] = None
    prompt_name: Optional[str] = None
    prompt_version: Optional[str] = None
    attributes: Optional[str] = None
    children: Optional[ChildMapping] = None


class TraceMapping(BaseModel):
    """Maps fields from a JSON dict to trace and observation attributes.

    Accepts dot-path strings that resolve against the input data.
    Two formats are supported -- this Pydantic model for type-checked pipelines,
    and a plain dict for quick scripts. Both are first-class.

    Example (Pydantic)::

        mapping = TraceMapping(
            name="agent_name",
            input="request.query",
            output="response.answer",
            children=ChildMapping(
                path="steps",
                name="action",
                input="action_input",
                output="action_result",
            ),
        )

    Example (dict, equivalent)::

        mapping = {
            "trace.name": "agent_name",
            "trace.input": "request.query",
            "trace.output": "response.answer",
            "children": "steps",
            "children.name": "action",
            "children.input": "action_input",
            "children.output": "action_result",
        }
    """

    name: Optional[str] = None
    input: Optional[str] = None
    output: Optional[str] = None
    metadata: Optional[str] = None
    tags: Optional[str] = None
    session_id: Optional[str] = None
    user_id: Optional[str] = None
    run_id: Optional[str] = None
    timestamp: Optional[str] = None
    environment: Optional[str] = None
    children: Optional[ChildMapping] = None

    @classmethod
    def from_dict(cls, raw: dict) -> TraceMapping:
        """Convert a flat dict-form mapping to a TraceMapping.

        The dict uses dotted keys like ``"trace.name"`` and ``"children.input"``.
        Raises PushError if any keys are unrecognized.
        """
        from .mapping import normalize_mapping

        mapping, unrecognized = normalize_mapping(raw)
        if unrecognized:
            raise PushError(
                f"Unrecognized mapping key(s): {unrecognized}. " "Keys must start with 'trace.' or 'children.'."
            )
        return mapping


# =============================================================================
# Validation
# =============================================================================


class ValidationIssue(BaseModel):
    """A structured validation error or warning.

    Machine-parseable for LLM-assisted workflows and programmatic handling.
    """

    field: str
    """Target field, e.g. "trace.input" or "children.model"."""

    path: str
    """Dot-path that was attempted, e.g. "request.query"."""

    issue: str
    """Machine-readable issue type: "path_not_found", "type_mismatch",
    "invalid_json", "invalid_type_value"."""

    message: str
    """Human-readable explanation."""

    available_keys: Optional[List[str]] = None
    """For path_not_found: top-level keys available in the data at that level."""

    expected_type: Optional[str] = None
    """For type_mismatch: what type was expected."""

    actual_type: Optional[str] = None
    """For type_mismatch: what type was found."""


class PushError(Exception):
    """Raised when push() encounters an unrecoverable error.

    Carries structured issues when available.
    """

    def __init__(self, message: str, issues: Optional[List[ValidationIssue]] = None):
        super().__init__(message)
        self.issues = issues or []


# =============================================================================
# Preview / Results
# =============================================================================


class SpanPreview(BaseModel):
    """Preview of a single observation that would be created."""

    name: Optional[str] = None
    type: Optional[str] = None
    input: Optional[Any] = None
    output: Optional[Any] = None
    model: Optional[str] = None
    metadata: Optional[Dict[str, Any]] = None
    start_time: Optional[str] = None
    end_time: Optional[str] = None
    children: List[SpanPreview] = []


class TracePreview(BaseModel):
    """Preview of the full trace tree that would be created.

    Returned by ``tracing.validate()`` for dry-run inspection.
    """

    name: Optional[str] = None
    input: Optional[Any] = None
    output: Optional[Any] = None
    metadata: Optional[Dict[str, Any]] = None
    tags: List[str] = []
    session_id: Optional[str] = None
    user_id: Optional[str] = None
    run_id: Optional[str] = None
    environment: Optional[str] = None
    children: List[SpanPreview] = []
    warnings: List[ValidationIssue] = []
    errors: List[ValidationIssue] = []


class PushResult(BaseModel):
    """Result returned by ``tracing.push()`` after successful span creation."""

    trace_id: str
    observation_count: int
    warnings: List[ValidationIssue] = []


# =============================================================================
# Inspection
# =============================================================================


class FieldInfo(BaseModel):
    """Description of a single field found during JSON inspection."""

    path: str
    """Dot-path to the field, e.g. "steps[0].step_input"."""

    type: str
    """Python type name, e.g. "str", "dict", "list[dict]"."""

    sample: str
    """Truncated sample value for display."""

    matched_as: Optional[str] = None
    """Convention target this matched, e.g. "trace.input". None if unmatched."""


class InspectResult(BaseModel):
    """Result of ``tracing.inspect()`` -- shows JSON structure and convention matches."""

    structure: List[FieldInfo]
    """All fields found in the JSON, with types and sample values."""

    convention_matches: Dict[str, str]
    """Fields that matched conventions: {"trace.input": "query", ...}."""

    unmatched_keys: List[FieldInfo]
    """Fields that didn't match any convention."""

    suggested_mapping: Dict[str, str]
    """Heuristic suggestions for unmatched keys only."""

    def to_mapping(self) -> Dict[str, str]:
        """Return the complete merged mapping (conventions + suggestions).

        Combines convention_matches with suggested_mapping into a single
        dict ready to pass to ``push()``. Tweak the result before using it.
        """
        return {**self.convention_matches, **self.suggested_mapping}
