# Copyright (c) Metis. All rights reserved.

"""Public API for JSON push-to-trace.

Provides three functions:
    - push()     -- convert JSON to OTel spans and export
    - validate() -- dry-run preview without exporting
    - inspect()  -- analyze JSON structure and suggest mappings
"""

from __future__ import annotations

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

from .mapping import (
    CHILD_CONVENTIONS,
    CHILDREN_KEYS,
    TRACE_CONVENTIONS,
    apply_conventions,
    describe_type,
    merge_fields,
    normalize_mapping,
    resolve_path,
    suggest_mapping_for_key,
    truncate_sample,
    validate_mapping,
)
from .push_models import (
    ChildMapping,
    FieldInfo,
    InspectResult,
    PushError,
    PushResult,
    SpanPreview,
    TraceMapping,
    TracePreview,
    ValidationIssue,
)
from .span_builder import build_trace_preview, build_trace_spans

logger = logging.getLogger(__name__)


# =============================================================================
# push()
# =============================================================================


def push(
    data: Union[dict, list, str],
    *,
    mapping: Union[TraceMapping, dict, None] = None,
    depth: Optional[int] = 1,
    name: Optional[str] = None,
    session_id: Optional[str] = None,
    user_id: Optional[str] = None,
    run_id: Optional[str] = None,
    tags: Optional[List[str]] = None,
    metadata: Optional[dict] = None,
    environment: Optional[str] = None,
    attributes: Optional[Dict[str, Any]] = None,
    strict: bool = True,
    timestamp_unit: str = "auto",
) -> Union[PushResult, List[PushResult]]:
    """Push structured JSON data as a trace.

    Converts JSON into OTel spans and routes through the existing export
    pipeline. No new server endpoints needed.

    Args:
        data: JSON data as a dict, list of dicts (batch), or JSON string.
            If str, parsed as JSON first (raises PushError on invalid JSON).
        mapping: Field mapping config. Accepts TraceMapping (Pydantic) or dict.
            If None, convention defaults auto-detect fields.
            Pass {} (empty dict) to disable conventions.
        depth: Controls convention auto-detection recursion depth.
            0=flat (root only), 1=one level, None=unlimited.
            Ignored when an explicit mapping with nested children is provided.
        name: Override trace name.
        session_id: Override session ID.
        user_id: Override user ID.
        run_id: Override run ID.
        tags: Tags to merge with mapped tags.
        metadata: Metadata to deep-merge with mapped metadata.
        environment: Override environment.
        attributes: Extra OTel span attributes (escape hatch).
        strict: If True, raise on validation errors. If False, warn and continue.
        timestamp_unit: Unit for numeric timestamp values in the data.
            "s" (seconds), "ms" (milliseconds), "us" (microseconds),
            "ns" (nanoseconds), or "auto" (default, infer from magnitude).

    Returns:
        PushResult for a single dict, List[PushResult] for a batch.

    Raises:
        PushError: On invalid JSON or validation errors (when strict=True).
    """
    parsed = _parse_input(data)

    # Handle batch: resolve mapping once, push each item
    if isinstance(parsed, list):
        return _push_batch(
            parsed,
            mapping=mapping,
            depth=depth,
            name=name,
            session_id=session_id,
            user_id=user_id,
            run_id=run_id,
            tags=tags,
            metadata=metadata,
            environment=environment,
            attributes=attributes,
            strict=strict,
            timestamp_unit=timestamp_unit,
        )

    # Resolve mapping
    resolved_mapping, unrecognized_keys = _resolve_mapping(parsed, mapping, depth)

    # Fail fast on unrecognized mapping keys
    if unrecognized_keys:
        for key in unrecognized_keys:
            issue = ValidationIssue(
                field="mapping",
                path=key,
                issue="unrecognized_key",
                message=f"Mapping key '{key}' not recognized. Keys must start with 'trace.' or 'children.'.",
            )
            if strict:
                raise PushError(
                    f"Unrecognized mapping key: '{key}'. " f"Keys must start with 'trace.' or 'children.'.",
                    issues=[issue],
                )

    # Validate
    errors, warnings = validate_mapping(parsed, resolved_mapping)

    # Surface unrecognized keys as warnings in lenient mode
    for key in unrecognized_keys:
        warnings.append(
            ValidationIssue(
                field="mapping",
                path=key,
                issue="unrecognized_key",
                message=f"Mapping key '{key}' not recognized. Keys must start with 'trace.' or 'children.'.",
            )
        )

    if errors and strict:
        messages = "; ".join(e.message for e in errors)
        raise PushError(f"Validation failed: {messages}", issues=errors)
    if errors and not strict:
        warnings.extend(errors)  # Demote errors to warnings in lenient mode

    # Extract mapped values
    mapped_values = _extract_mapped_values(parsed, resolved_mapping)

    # Build explicit overrides
    explicit = _build_explicit(
        name=name,
        session_id=session_id,
        user_id=user_id,
        run_id=run_id,
        tags=tags,
        metadata=metadata,
        environment=environment,
    )

    # Merge fields
    fields = merge_fields(
        mapped=mapped_values,
        explicit=explicit,
        context={},  # Ambient context injected by TracingContextSpanProcessor
        defaults={},
    )

    # Build OTel spans
    trace_id, observation_count, runtime_warnings = build_trace_spans(
        data=parsed,
        fields=fields,
        child_mapping=resolved_mapping.children,
        extra_attributes=attributes,
        timestamp_unit=timestamp_unit,
    )

    # Strict mode: raise on any runtime issues (non-dict children, bad timestamps)
    if runtime_warnings and strict:
        messages = "; ".join(w.message for w in runtime_warnings)
        raise PushError(f"Runtime validation failed: {messages}", issues=runtime_warnings)

    warnings.extend(runtime_warnings)

    logger.info(
        "Pushed trace %s with %d observations",
        trace_id,
        observation_count,
    )

    return PushResult(
        trace_id=trace_id,
        observation_count=observation_count,
        warnings=warnings,
    )


def _push_batch(
    items: list,
    *,
    mapping: Union[TraceMapping, dict, None],
    depth: Optional[int],
    name: Optional[str],
    session_id: Optional[str],
    user_id: Optional[str],
    run_id: Optional[str],
    tags: Optional[List[str]],
    metadata: Optional[dict],
    environment: Optional[str],
    attributes: Optional[Dict[str, Any]],
    strict: bool,
    timestamp_unit: str = "auto",
) -> List[PushResult]:
    """Flat batch path: resolve mapping once, build spans for each item."""
    # Check for non-dict items
    skipped = [(i, type(item).__name__) for i, item in enumerate(items) if not isinstance(item, dict)]
    if skipped and strict:
        indices = [i for i, _ in skipped]
        raise PushError(
            f"Batch contains {len(skipped)} non-dict item(s) at indices {indices}. " "All batch items must be dicts."
        )

    dict_items = [item for item in items if isinstance(item, dict)]
    if not dict_items:
        raise PushError("Batch data contained no valid dict items.")

    # Resolve mapping once using first item as reference
    resolved_mapping, unrecognized_keys = _resolve_mapping(dict_items[0], mapping, depth)

    if unrecognized_keys:
        for key in unrecognized_keys:
            issue = ValidationIssue(
                field="mapping",
                path=key,
                issue="unrecognized_key",
                message=f"Mapping key '{key}' not recognized. Keys must start with 'trace.' or 'children.'.",
            )
            if strict:
                raise PushError(
                    f"Unrecognized mapping key: '{key}'. " f"Keys must start with 'trace.' or 'children.'.",
                    issues=[issue],
                )

    # Validate mapping once against first item
    errors, mapping_warnings = validate_mapping(dict_items[0], resolved_mapping)

    for key in unrecognized_keys:
        mapping_warnings.append(
            ValidationIssue(
                field="mapping",
                path=key,
                issue="unrecognized_key",
                message=f"Mapping key '{key}' not recognized. Keys must start with 'trace.' or 'children.'.",
            )
        )

    if errors and strict:
        messages = "; ".join(e.message for e in errors)
        raise PushError(f"Validation failed: {messages}", issues=errors)

    explicit = _build_explicit(
        name=name,
        session_id=session_id,
        user_id=user_id,
        run_id=run_id,
        tags=tags,
        metadata=metadata,
        environment=environment,
    )

    # Build spans for each item (no recursion)
    results: List[PushResult] = []
    for item in dict_items:
        mapped_values = _extract_mapped_values(item, resolved_mapping)
        fields = merge_fields(
            mapped=mapped_values,
            explicit=explicit,
            context={},
            defaults={},
        )

        trace_id, obs_count, runtime_warnings = build_trace_spans(
            data=item,
            fields=fields,
            child_mapping=resolved_mapping.children,
            extra_attributes=attributes,
            timestamp_unit=timestamp_unit,
        )

        if runtime_warnings and strict:
            messages = "; ".join(w.message for w in runtime_warnings)
            raise PushError(f"Runtime validation failed: {messages}", issues=runtime_warnings)

        item_warnings = list(mapping_warnings) + runtime_warnings
        if errors and not strict:
            item_warnings.extend(errors)

        results.append(
            PushResult(
                trace_id=trace_id,
                observation_count=obs_count,
                warnings=item_warnings,
            )
        )

    # Surface skipped items as warnings on the last result (lenient mode)
    if skipped and results:
        indices = [i for i, _ in skipped]
        results[-1].warnings.append(
            ValidationIssue(
                field="batch",
                path="",
                issue="non_dict_items",
                message=f"Skipped {len(skipped)} non-dict item(s) at indices {indices}.",
            )
        )

    logger.info("Pushed batch of %d traces", len(results))
    return results


# =============================================================================
# validate()
# =============================================================================


def validate(
    data: Union[dict, str],
    *,
    mapping: Union[TraceMapping, dict, None] = None,
    depth: Optional[int] = 1,
    name: Optional[str] = None,
    session_id: Optional[str] = None,
    user_id: Optional[str] = None,
    run_id: Optional[str] = None,
    tags: Optional[List[str]] = None,
    metadata: Optional[dict] = None,
    environment: Optional[str] = None,
) -> TracePreview:
    """Validate data + mapping without exporting.

    Returns a preview of the trace tree that would be created,
    along with any warnings and errors.
    """
    parsed = _parse_input(data)
    if isinstance(parsed, list):
        raise PushError("validate() does not accept batch data. Pass a single dict.")

    resolved_mapping, unrecognized_keys = _resolve_mapping(parsed, mapping, depth)
    errors, warnings = validate_mapping(parsed, resolved_mapping)

    for key in unrecognized_keys:
        errors.append(
            ValidationIssue(
                field="mapping",
                path=key,
                issue="unrecognized_key",
                message=f"Mapping key '{key}' not recognized. Keys must start with 'trace.' or 'children.'.",
            )
        )

    mapped_values = _extract_mapped_values(parsed, resolved_mapping)

    explicit = _build_explicit(
        name=name,
        session_id=session_id,
        user_id=user_id,
        run_id=run_id,
        tags=tags,
        metadata=metadata,
        environment=environment,
    )

    fields = merge_fields(
        mapped=mapped_values,
        explicit=explicit,
        context={},
        defaults={},
    )

    children_preview = build_trace_preview(
        parsed,
        fields,
        resolved_mapping.children,
    )

    return TracePreview(
        name=fields.get("name"),
        input=fields.get("input"),
        output=fields.get("output"),
        metadata=fields.get("metadata"),
        tags=fields.get("tags", []),
        session_id=fields.get("session_id"),
        user_id=fields.get("user_id"),
        run_id=fields.get("run_id"),
        environment=fields.get("environment"),
        children=children_preview,
        warnings=warnings,
        errors=errors,
    )


# =============================================================================
# inspect()
# =============================================================================


def inspect(
    data: Union[dict, str],
    *,
    depth: Optional[int] = None,
) -> InspectResult:
    """Inspect a JSON structure and show what conventions match.

    Does not require tracing.init() -- purely local analysis.

    Args:
        data: JSON data as a dict or JSON string.
        depth: How deep to recurse into children for inspection.
            None = unlimited.

    Returns:
        InspectResult with structure, convention matches, unmatched keys,
        and a suggested mapping.
    """
    parsed = _parse_input(data)
    if isinstance(parsed, list):
        raise PushError("inspect() does not accept batch data. Pass a single dict.")

    structure: List[FieldInfo] = []
    convention_matches: Dict[str, str] = {}
    unmatched: List[FieldInfo] = []
    suggestions: Dict[str, str] = {}

    # Scan trace-level fields
    for target, candidates in TRACE_CONVENTIONS.items():
        from .mapping import _find_convention_match

        match = _find_convention_match(parsed, candidates)
        if match is not None:
            convention_matches[f"trace.{target}"] = match

    # Detect children
    from .mapping import _find_convention_match

    children_key = _find_convention_match(parsed, CHILDREN_KEYS)
    if children_key is not None:
        convention_matches["children"] = children_key

    # Build structure and find unmatched keys
    matched_root_keys = set(convention_matches.values())

    for key, value in parsed.items():
        info = FieldInfo(
            path=key,
            type=describe_type(value),
            sample=truncate_sample(value),
            matched_as=None,
        )

        # Check if this key was matched
        for target, matched_key in convention_matches.items():
            if matched_key == key:
                info.matched_as = target
                break

        structure.append(info)

        if key not in matched_root_keys:
            suggestion = suggest_mapping_for_key(key, "root")
            if suggestion:
                suggestions[suggestion] = key
            unmatched.append(info)

    # Scan children if present
    if children_key is not None:
        found, children_data = resolve_path(parsed, children_key)
        if found and isinstance(children_data, list) and children_data:
            sample_child = children_data[0]
            if isinstance(sample_child, dict):
                _inspect_children(
                    sample_child,
                    children_key,
                    convention_matches,
                    structure,
                    unmatched,
                    suggestions,
                )

    return InspectResult(
        structure=structure,
        convention_matches=convention_matches,
        unmatched_keys=unmatched,
        suggested_mapping=suggestions,
    )


# =============================================================================
# Internal Helpers
# =============================================================================


def _parse_input(data: Union[dict, list, str]) -> Union[dict, list]:
    """Parse input data, accepting dict, list, or JSON string."""
    if isinstance(data, (dict, list)):
        return data
    if isinstance(data, str):
        try:
            parsed = json.loads(data)
        except json.JSONDecodeError as e:
            raise PushError(f"Invalid JSON string: {e}. Pass a dict instead, or fix the JSON.") from e
        if not isinstance(parsed, (dict, list)):
            raise PushError(f"JSON parsed to {type(parsed).__name__}, expected dict or list.")
        return parsed
    raise PushError(f"Unsupported data type: {type(data).__name__}. Expected dict, list, or str.")


def _resolve_mapping(
    data: dict,
    mapping: Union[TraceMapping, dict, None],
    depth: Optional[int],
) -> tuple[TraceMapping, List[str]]:
    """Resolve the mapping parameter to a TraceMapping.

    Returns (mapping, unrecognized_keys).
    """
    if mapping is None:
        return apply_conventions(data, depth=depth), []
    if isinstance(mapping, dict):
        if not mapping:
            return TraceMapping(), []  # Empty = no conventions
        return normalize_mapping(mapping)
    return mapping, []


def _extract_mapped_values(data: dict, mapping: TraceMapping) -> Dict[str, Any]:
    """Extract values from data using the resolved mapping."""
    values: Dict[str, Any] = {}

    fields = [
        "name",
        "input",
        "output",
        "metadata",
        "tags",
        "session_id",
        "user_id",
        "run_id",
        "timestamp",
        "environment",
    ]
    for field in fields:
        dot_path = getattr(mapping, field, None)
        if dot_path is not None:
            found, value = resolve_path(data, dot_path)
            if found:
                values[field] = value

    return values


def _build_explicit(**kwargs: Any) -> Dict[str, Any]:
    """Build the explicit overrides dict, filtering out None values."""
    return {k: v for k, v in kwargs.items() if v is not None}


def _inspect_children(
    sample_child: dict,
    children_key: str,
    convention_matches: Dict[str, str],
    structure: List[FieldInfo],
    unmatched: List[FieldInfo],
    suggestions: Dict[str, str],
) -> None:
    """Inspect child fields and update conventions/suggestions."""
    from .mapping import _find_convention_match

    matched_child_keys: set = set()

    for target, candidates in CHILD_CONVENTIONS.items():
        match = _find_convention_match(sample_child, candidates)
        if match is not None:
            convention_matches[f"children.{target}"] = match
            matched_child_keys.add(match)

    for key, value in sample_child.items():
        path = f"{children_key}[*].{key}"
        info = FieldInfo(
            path=path,
            type=describe_type(value),
            sample=truncate_sample(value),
            matched_as=None,
        )

        if key in matched_child_keys:
            for target, matched_key in convention_matches.items():
                if target.startswith("children.") and matched_key == key:
                    info.matched_as = target
                    break
            structure.append(info)
        else:
            suggestion = suggest_mapping_for_key(key, "child")
            if suggestion:
                suggestions[suggestion] = key
            structure.append(info)
            unmatched.append(info)
