# Copyright (c) Metis. All rights reserved.

"""Mapping engine for JSON push-to-trace.

Pure data transformation -- no OTel dependency. Handles:
- Dot-path resolution through nested dicts
- Convention-based auto-detection (snake_case + camelCase)
- Dict-to-TraceMapping normalization
- Field merge strategy (explicit > mapped > context > defaults)
- Timestamp coercion to nanoseconds
- Mapping validation with structured errors
"""

from __future__ import annotations

import json
import re
from datetime import datetime, timezone
from typing import Any, Dict, List, Optional, Tuple

from .push_models import (
    ChildMapping,
    TraceMapping,
    ValidationIssue,
)

# =============================================================================
# Convention Tables
# =============================================================================

# Each list is checked in order; first match wins.
# Both snake_case and camelCase variants are included.

TRACE_NAME_KEYS = [
    "name",
    "trace_name",
    "traceName",
    "agent_name",
    "agentName",
    "task_name",
    "taskName",
    "workflow_name",
    "workflowName",
]

TRACE_INPUT_KEYS = [
    "input",
    "query",
    "question",
    "prompt",
    "request",
    "user_input",
    "userInput",
    "user_message",
    "userMessage",
]

TRACE_OUTPUT_KEYS = [
    "output",
    "result",
    "answer",
    "response",
    "completion",
    "final_answer",
    "finalAnswer",
    "agent_output",
    "agentOutput",
]

TRACE_METADATA_KEYS = [
    "metadata",
    "config",
    "configuration",
    "params",
    "parameters",
    "context",
]

TRACE_TAGS_KEYS = ["tags", "labels", "categories"]

TRACE_SESSION_ID_KEYS = [
    "session_id",
    "sessionId",
    "session",
    "conversation_id",
    "conversationId",
    "thread_id",
    "threadId",
]

TRACE_USER_ID_KEYS = [
    "user_id",
    "userId",
    "user",
    "author",
]

TRACE_TIMESTAMP_KEYS = [
    "timestamp",
    "created_at",
    "createdAt",
    "start_time",
    "startTime",
    "time",
    "ts",
]

CHILDREN_KEYS = [
    "steps",
    "spans",
    "observations",
    "events",
    "actions",
    "calls",
    "messages",
    "tool_calls",
    "toolCalls",
    "children",
    "traces",
]

# Per-child conventions (in addition to the trace-level ones above)
CHILD_TYPE_KEYS = [
    "type",
    "kind",
    "span_kind",
    "spanKind",
    "observation_type",
    "observationType",
]

CHILD_MODEL_KEYS = [
    "model",
    "model_name",
    "modelName",
    "llm",
    "model_id",
    "modelId",
]

CHILD_START_TIME_KEYS = [
    "start_time",
    "startTime",
    "started_at",
    "startedAt",
    "timestamp",
    "begin",
]

CHILD_END_TIME_KEYS = [
    "end_time",
    "endTime",
    "ended_at",
    "endedAt",
    "finished_at",
    "finishedAt",
    "completed_at",
    "completedAt",
]

CHILD_LEVEL_KEYS = ["level", "severity", "status"]

# Trace-level convention map: target_field -> candidate keys
TRACE_CONVENTIONS = {
    "name": TRACE_NAME_KEYS,
    "input": TRACE_INPUT_KEYS,
    "output": TRACE_OUTPUT_KEYS,
    "metadata": TRACE_METADATA_KEYS,
    "tags": TRACE_TAGS_KEYS,
    "session_id": TRACE_SESSION_ID_KEYS,
    "user_id": TRACE_USER_ID_KEYS,
    "timestamp": TRACE_TIMESTAMP_KEYS,
}

# Child-level convention map (extends trace-level for shared fields)
CHILD_CONVENTIONS = {
    "name": TRACE_NAME_KEYS,
    "input": TRACE_INPUT_KEYS,
    "output": TRACE_OUTPUT_KEYS,
    "metadata": TRACE_METADATA_KEYS,
    "type": CHILD_TYPE_KEYS,
    "model": CHILD_MODEL_KEYS,
    "start_time": CHILD_START_TIME_KEYS,
    "end_time": CHILD_END_TIME_KEYS,
    "level": CHILD_LEVEL_KEYS,
}


# =============================================================================
# Path Resolution
# =============================================================================


def resolve_path(data: dict, dot_path: str) -> Tuple[bool, Any]:
    """Navigate a dot-path string through nested dicts.

    Returns (found, value). Does not raise on missing keys.

    Examples::

        resolve_path({"a": {"b": 1}}, "a.b")  -> (True, 1)
        resolve_path({"a": 1}, "a.b.c")       -> (False, None)
    """
    current = data
    for key in dot_path.split("."):
        if isinstance(current, dict) and key in current:
            current = current[key]
        else:
            return False, None
    return True, current


# =============================================================================
# Convention Detection
# =============================================================================


def _find_convention_match(data: dict, candidate_keys: List[str]) -> Optional[str]:
    """Find the first candidate key that exists in data (case-insensitive)."""
    data_keys_lower = {k.lower(): k for k in data.keys()}
    for candidate in candidate_keys:
        actual_key = data_keys_lower.get(candidate.lower())
        if actual_key is not None:
            return actual_key
    return None


def apply_conventions(data: dict, depth: Optional[int] = 1) -> TraceMapping:
    """Scan JSON keys for convention matches and return a TraceMapping.

    Args:
        data: The JSON dict to scan.
        depth: How deep to recurse for children conventions.
            0 = trace-level only, 1 = one level of children, None = unlimited.
    """
    kwargs: Dict[str, Any] = {}

    for target, candidates in TRACE_CONVENTIONS.items():
        match = _find_convention_match(data, candidates)
        if match is not None:
            kwargs[target] = match

    # Detect children array
    if depth is None or depth > 0:
        children_key = _find_convention_match(data, CHILDREN_KEYS)
        if children_key is not None:
            found, children_data = resolve_path(data, children_key)
            if found and isinstance(children_data, list) and children_data:
                child_kwargs: Dict[str, Any] = {"path": children_key}
                sample = children_data[0]
                if isinstance(sample, dict):
                    for target, candidates in CHILD_CONVENTIONS.items():
                        match = _find_convention_match(sample, candidates)
                        if match is not None:
                            child_kwargs[target] = match
                kwargs["children"] = ChildMapping(**child_kwargs)

    return TraceMapping(**kwargs)


# =============================================================================
# Mapping Normalization
# =============================================================================


def normalize_mapping(raw: dict) -> Tuple[TraceMapping, List[str]]:
    """Convert a flat dict-form mapping to a TraceMapping.

    Accepts keys like "trace.name", "trace.input", "children", "children.name".

    Returns:
        (mapping, unrecognized_keys) -- the mapping and any keys that weren't
        recognized. Unrecognized keys likely indicate typos.
    """
    trace_kwargs: Dict[str, Any] = {}
    child_kwargs: Dict[str, Any] = {}
    unrecognized: List[str] = []

    for key, value in raw.items():
        if key.startswith("trace."):
            field = key[len("trace.") :]
            trace_kwargs[field] = value
        elif key == "children":
            child_kwargs["path"] = value
        elif key.startswith("children."):
            field = key[len("children.") :]
            child_kwargs[field] = value
        else:
            unrecognized.append(key)

    if child_kwargs and "path" in child_kwargs:
        trace_kwargs["children"] = ChildMapping(**child_kwargs)

    return TraceMapping(**trace_kwargs), unrecognized


# =============================================================================
# Timestamp Coercion
# =============================================================================

_NANOS_PER_SECOND = 1_000_000_000
_NANOS_PER_MILLI = 1_000_000
_NANOS_PER_MICRO = 1_000

# Thresholds for distinguishing epoch units
_SECONDS_THRESHOLD = 10**12
_MILLIS_THRESHOLD = 10**15
_MICROS_THRESHOLD = 10**18

_UNIT_MULTIPLIERS = {
    "s": _NANOS_PER_SECOND,
    "ms": _NANOS_PER_MILLI,
    "us": _NANOS_PER_MICRO,
    "ns": 1,
}


def coerce_timestamp(value: Any, unit: str = "auto") -> Optional[int]:
    """Convert a timestamp value to nanoseconds since epoch.

    Args:
        value: The timestamp value to convert.
        unit: Explicit unit for numeric values. One of "s", "ms", "us", "ns",
            or "auto" (default). When "auto", the unit is inferred from
            the magnitude of the value. Use an explicit unit to avoid
            magnitude-based guessing.

    Accepts:
        - int/float: epoch seconds, milliseconds, microseconds, or nanoseconds
        - str: ISO 8601 format
        - datetime: Python datetime object

    Returns None if the value cannot be coerced.
    """
    if value is None:
        return None

    if isinstance(value, datetime):
        if value.tzinfo is None:
            value = value.replace(tzinfo=timezone.utc)
        return int(value.timestamp() * _NANOS_PER_SECOND)

    if isinstance(value, (int, float)):
        if unit != "auto":
            multiplier = _UNIT_MULTIPLIERS.get(unit)
            if multiplier is None:
                return None
            return int(value * multiplier)

        # Auto: infer unit from magnitude.
        # Compare as int when possible to avoid float64 precision loss
        # at the microsecond/nanosecond boundary (~1e18).
        magnitude = value if isinstance(value, int) else float(value)
        if magnitude < _SECONDS_THRESHOLD:
            return int(value * _NANOS_PER_SECOND)
        if magnitude < _MILLIS_THRESHOLD:
            return int(value * _NANOS_PER_MILLI)
        if magnitude < _MICROS_THRESHOLD:
            return int(value * _NANOS_PER_MICRO)
        return int(value)

    if isinstance(value, str):
        try:
            dt = datetime.fromisoformat(value.replace("Z", "+00:00"))
            if dt.tzinfo is None:
                dt = dt.replace(tzinfo=timezone.utc)
            return int(dt.timestamp() * _NANOS_PER_SECOND)
        except (ValueError, TypeError):
            return None

    return None


# =============================================================================
# Field Merge
# =============================================================================


def merge_fields(
    mapped: Dict[str, Any],
    explicit: Dict[str, Any],
    context: Dict[str, Any],
    defaults: Dict[str, Any],
) -> Dict[str, Any]:
    """Apply the merge strategy across priority levels.

    Priority (highest to lowest):
        1. explicit -- push() keyword arguments
        2. mapped -- values extracted from JSON via mapping
        3. context -- ambient tracing_context() / run context
        4. defaults -- from tracing.init()

    Special rules:
        - tags: union, deduplicated
        - metadata: deep merge, higher priority wins on key conflict
    """
    result: Dict[str, Any] = {}
    all_sources = [defaults, context, mapped, explicit]

    # Scalar fields: highest priority wins
    scalar_fields = [
        "name",
        "session_id",
        "user_id",
        "run_id",
        "environment",
        "input",
        "output",
        "timestamp",
    ]
    for field in scalar_fields:
        for source in reversed(all_sources):
            if field in source and source[field] is not None:
                result[field] = source[field]
                break

    # Tags: union across all sources, deduplicated
    all_tags: List[str] = []
    for source in all_sources:
        tags = source.get("tags")
        if isinstance(tags, list):
            all_tags.extend(tags)
    if all_tags:
        result["tags"] = list(dict.fromkeys(all_tags))  # preserves order

    # Metadata: deep merge, higher priority wins on conflict
    merged_meta: Dict[str, Any] = {}
    for source in all_sources:
        meta = source.get("metadata")
        if isinstance(meta, dict):
            merged_meta.update(meta)
    if merged_meta:
        result["metadata"] = merged_meta

    return result


# =============================================================================
# Mapping Validation
# =============================================================================

# Expected types for each mappable field
_EXPECTED_TYPES = {
    "name": "str",
    "input": "json_serializable",
    "output": "json_serializable",
    "metadata": "dict",
    "tags": "list",
    "session_id": "str",
    "user_id": "str",
    "run_id": "str",
    "timestamp": "timestamp",
    "environment": "str",
}

_CHILD_EXPECTED_TYPES = {
    "name": "str",
    "input": "json_serializable",
    "output": "json_serializable",
    "metadata": "dict",
    "type": "str",
    "model": "str",
    "level": "str",
    "start_time": "timestamp",
    "end_time": "timestamp",
    "usage": "dict",
    "cost": "dict",
    "model_parameters": "dict",
    "status_message": "str",
    "prompt_name": "str",
    "prompt_version": "number",
    "attributes": "dict",
}

_VALID_OBSERVATION_TYPES = {
    "SPAN",
    "GENERATION",
    "TOOL",
    "AGENT",
    "CHAIN",
    "RETRIEVER",
    "EMBEDDING",
    "EVALUATOR",
    "GUARDRAIL",
    "EVENT",
    # Also accept lowercase
    "span",
    "generation",
    "tool",
    "agent",
    "chain",
    "retriever",
    "embedding",
    "evaluator",
    "guardrail",
    "event",
}

_VALID_LEVELS = {"DEBUG", "DEFAULT", "WARNING", "ERROR"}


def _is_json_serializable(value: Any) -> bool:
    """Check if a value is JSON-serializable (str, dict, list, number, bool, None)."""
    return isinstance(value, (str, int, float, bool, dict, list, type(None)))


def _check_type(value: Any, expected: str) -> bool:
    """Check if a value matches the expected type category."""
    if expected == "str":
        return isinstance(value, str)
    if expected == "dict":
        return isinstance(value, dict)
    if expected == "list":
        return isinstance(value, list)
    if expected == "number":
        return isinstance(value, (int, float))
    if expected == "json_serializable":
        return _is_json_serializable(value)
    if expected == "timestamp":
        return coerce_timestamp(value) is not None
    return True


def validate_mapping(
    data: dict,
    mapping: TraceMapping,
) -> Tuple[List[ValidationIssue], List[ValidationIssue]]:
    """Validate that all mapping paths resolve and types match.

    Returns (errors, warnings).
    """
    errors: List[ValidationIssue] = []
    warnings: List[ValidationIssue] = []

    def _available_keys(d: dict) -> List[str]:
        return list(d.keys())[:20]

    # Validate trace-level fields
    for field, expected in _EXPECTED_TYPES.items():
        dot_path = getattr(mapping, field, None)
        if dot_path is None:
            continue

        found, value = resolve_path(data, dot_path)
        if not found:
            issue = ValidationIssue(
                field=f"trace.{field}",
                path=dot_path,
                issue="path_not_found",
                message=f"Mapping path '{dot_path}' not found in data.",
                available_keys=_available_keys(data),
            )
            if field in ("name", "input", "output"):
                errors.append(issue)
            else:
                warnings.append(issue)
            continue

        if not _check_type(value, expected):
            errors.append(
                ValidationIssue(
                    field=f"trace.{field}",
                    path=dot_path,
                    issue="type_mismatch",
                    message=f"Field 'trace.{field}' resolved to type {type(value).__name__}, expected {expected}.",
                    expected_type=expected,
                    actual_type=type(value).__name__,
                )
            )

    # Validate children
    if mapping.children is not None:
        _validate_children(data, mapping.children, errors, warnings)

    return errors, warnings


def _validate_children(
    data: dict,
    child_mapping: ChildMapping,
    errors: List[ValidationIssue],
    warnings: List[ValidationIssue],
) -> None:
    """Validate children mapping paths against the data."""
    found, children_data = resolve_path(data, child_mapping.path)

    if not found:
        errors.append(
            ValidationIssue(
                field="children",
                path=child_mapping.path,
                issue="path_not_found",
                message=f"Children path '{child_mapping.path}' not found in data.",
                available_keys=list(data.keys())[:20],
            )
        )
        return

    if not isinstance(children_data, list):
        errors.append(
            ValidationIssue(
                field="children",
                path=child_mapping.path,
                issue="type_mismatch",
                message=f"Children path '{child_mapping.path}' resolved to {type(children_data).__name__}, expected list.",
                expected_type="list",
                actual_type=type(children_data).__name__,
            )
        )
        return

    if not children_data:
        return

    # Validate against first child as sample
    sample = children_data[0]
    if not isinstance(sample, dict):
        return

    for field, expected in _CHILD_EXPECTED_TYPES.items():
        dot_path = getattr(child_mapping, field, None)
        if dot_path is None:
            continue

        found, value = resolve_path(sample, dot_path)
        if not found:
            warnings.append(
                ValidationIssue(
                    field=f"children.{field}",
                    path=dot_path,
                    issue="path_not_found",
                    message=f"Mapping path '{dot_path}' not found in children[0].",
                    available_keys=list(sample.keys())[:20],
                )
            )
            continue

        if not _check_type(value, expected):
            warnings.append(
                ValidationIssue(
                    field=f"children.{field}",
                    path=dot_path,
                    issue="type_mismatch",
                    message=f"Field 'children.{field}' resolved to type {type(value).__name__}, expected {expected}.",
                    expected_type=expected,
                    actual_type=type(value).__name__,
                )
            )

        # Extra validation for constrained string fields
        if field == "type" and isinstance(value, str) and value not in _VALID_OBSERVATION_TYPES:
            warnings.append(
                ValidationIssue(
                    field="children.type",
                    path=dot_path,
                    issue="invalid_type_value",
                    message=f"Observation type '{value}' not recognized. Valid: SPAN, GENERATION, TOOL, AGENT, etc.",
                )
            )

        if field == "level" and isinstance(value, str) and value.upper() not in _VALID_LEVELS:
            warnings.append(
                ValidationIssue(
                    field="children.level",
                    path=dot_path,
                    issue="invalid_type_value",
                    message=f"Level '{value}' not recognized. Valid: DEBUG, DEFAULT, WARNING, ERROR.",
                )
            )


# =============================================================================
# Inspection Helpers
# =============================================================================

# Heuristic suggestions for unmatched keys
_SUGGESTION_PATTERNS: List[Tuple[str, str, re.Pattern]] = [
    ("trace.output", "root", re.compile(r"(final|agent|llm).*(answer|output|result|response)", re.I)),
    ("trace.input", "root", re.compile(r"(user|agent|llm).*(input|query|question|prompt)", re.I)),
    ("children.input", "child", re.compile(r"(step|action|tool).*(input|query)", re.I)),
    ("children.output", "child", re.compile(r"(step|action|tool).*(output|result|response)", re.I)),
    ("children.model", "child", re.compile(r"(model|llm).*(used|name|id)", re.I)),
]


def suggest_mapping_for_key(key: str, level: str) -> Optional[str]:
    """Try to suggest a mapping target for an unrecognized key using heuristics."""
    for target, target_level, pattern in _SUGGESTION_PATTERNS:
        if target_level == level and pattern.search(key):
            return target

    # Simple suffix heuristics
    lower = key.lower()
    if level == "child":
        if lower.endswith("_input") or lower.endswith("input"):
            return "children.input"
        if lower.endswith("_output") or lower.endswith("output"):
            return "children.output"
        if lower.endswith("_result") or lower.endswith("result"):
            return "children.output"

    return None


def describe_type(value: Any) -> str:
    """Return a human-readable type description for a value."""
    if isinstance(value, list):
        if value and isinstance(value[0], dict):
            return "list[dict]"
        if value and isinstance(value[0], str):
            return "list[str]"
        return "list"
    return type(value).__name__


def truncate_sample(value: Any, max_len: int = 60) -> str:
    """Return a truncated string representation of a value."""
    if isinstance(value, str):
        if len(value) > max_len:
            return json.dumps(value[:max_len] + "...")
        return json.dumps(value)
    if isinstance(value, dict):
        s = json.dumps(value)
        if len(s) > max_len:
            return s[:max_len] + "..."
        return s
    if isinstance(value, list):
        return f"({len(value)} items)"
    return str(value)
