"""
Custom Span Exporter for Pendo Agent SDK

Converts OTel ReadableSpan objects into Pendo's agentic event format and
exports them to Pendo's ingestion pipeline.

The exporter is wired into the pipeline by init():
    TracerProvider -> BatchSpanProcessor -> PendoSpanExporter -> HTTP POST

Wire format follows the ingestion spec (APP-146809):
  - Endpoint: POST /data/agentic/<apiKey>?v=<version>&ct=<time>&s=<size>
  - Body: { "events": "<jzb>" }  (JZB = JSON → zlib → URL-safe base64)
  - Event type: "agent_trace" with trace fields in props
  - Source: "api" (set by the endpoint, not the SDK)

OTel span attribute → Pendo event field mapping:
    pendo.span.type              → props.agentTraceType
    llm.model_name               → props.modelUsed
    pendo.conversation_id        → props.conversationId
    message.content              → props.agentTraceContent
    tool.name                    → props.toolsUsed
    OTel trace_id                → props.agentTraceId
    OTel span_id                 → props.agentSpanId
    OTel parent.span_id          → props.agentParentSpanId
"""

import base64
import json
import logging
import re
import time
import uuid
import zlib
from typing import Any, Dict, List, Optional, Sequence

from opentelemetry.sdk.trace import ReadableSpan
from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult
from opentelemetry.trace import StatusCode

from .schema import (
    AgentEvent,
    EventBatch,
    EventProps,
    SpanStatusCode,
    SpanType,
)

logger = logging.getLogger("pendo_sdk")

_SDK_VERSION = "0.1.0"

# PII redaction patterns applied before network export.
_PII_PATTERNS = [
    (re.compile(r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b"), "[EMAIL_REDACTED]"),
    (re.compile(r"\b\d{3}[-.]?\d{3}[-.]?\d{4}\b"), "[PHONE_REDACTED]"),
    (re.compile(r"\b\d{3}-\d{2}-\d{4}\b"), "[SSN_REDACTED]"),
]

# Maps OTel span attribute keys → EventProps field names.
# Attributes not in this map go into props.extra.
_ATTR_TO_PROP: Dict[str, str] = {
    "llm.model_name": "model",
    "llm.token_count.prompt": "prompt_tokens",
    "llm.token_count.completion": "completion_tokens",
    "llm.token_count.total": "total_tokens",
    "llm.cost_usd": "cost_usd",
    "tool.name": "tool_name",
    "tool.input": "tool_input",
    "tool.output": "tool_output",
    "tool.status": "tool_status",
    "message.role": "role",
    "message.content": "content",
    "message.suggested_prompt": "suggested_prompt",
    "message.file_uploaded": "file_uploaded",
    "signal.name": "signal_name",
    "reaction.type": "reaction_type",
    "reaction.feedback_comment": "feedback_comment",
    "reaction.message_span_id": "message_span_id",
}

# Attribute keys consumed at the event level (not forwarded to props).
_EVENT_LEVEL_ATTRS = {"pendo.span.type", "pendo.conversation_id", "pendo.event_type", "pendo.tags", "message.trace_content"}

# All "known" attribute keys — anything else goes to props.extra.
_KNOWN_ATTRS = set(_ATTR_TO_PROP.keys()) | _EVENT_LEVEL_ATTRS

# Map our SpanType to the agentTraceType values the backend expects.
_SPAN_TYPE_TO_TRACE_TYPE: Dict[str, str] = {
    "GENERATION": "llm_request",
    "TOOL": "tool_request",
    "AGENT": "llm_request",
    "CHAIN": "llm_request",
    "USER_MESSAGE": "llm_request",
    "REACTION": "llm_response",
    "RETRIEVAL": "tool_request",
}


# ---------------------------------------------------------------------------
# JZB encoding (JSON → zlib → URL-safe base64, strip padding)
# ---------------------------------------------------------------------------

def _to_jzb(obj: Any) -> str:
    """Encode a Python object as a JZB string.

    JZB = JSON → zlib compress → URL-safe base64 (padding stripped).
    This is Pendo's standard wire encoding for event payloads.
    """
    json_bytes = json.dumps(obj, separators=(",", ":"), default=str).encode("utf-8")
    compressed = zlib.compress(json_bytes)
    b64 = base64.urlsafe_b64encode(compressed).decode("ascii")
    # Strip trailing '=' padding (decoder restores it)
    return b64.rstrip("=")


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _redact(value: str) -> str:
    """Apply all PII redaction patterns to a string."""
    for pattern, replacement in _PII_PATTERNS:
        value = pattern.sub(replacement, value)
    return value


def _redact_attributes(attrs: dict) -> dict:
    """Redact PII from all string values in an attributes dict."""
    redacted = {}
    for key, value in attrs.items():
        if isinstance(value, str):
            redacted[key] = _redact(value)
        else:
            redacted[key] = value
    return redacted


def _ns_to_ms(ns: Optional[int]) -> Optional[int]:
    """Convert nanosecond OTel timestamp to epoch milliseconds."""
    if ns is None:
        return None
    return ns // 1_000_000


def _otel_status_to_schema(status_code: StatusCode) -> SpanStatusCode:
    """Map OTel StatusCode to our schema's SpanStatusCode."""
    mapping = {
        StatusCode.UNSET: SpanStatusCode.UNSET,
        StatusCode.OK: SpanStatusCode.OK,
        StatusCode.ERROR: SpanStatusCode.ERROR,
    }
    return mapping.get(status_code, SpanStatusCode.UNSET)


def _resolve_span_type(raw: Any) -> SpanType:
    """Resolve a raw pendo.span.type attribute value to a SpanType enum."""
    if raw is None:
        return SpanType.UNKNOWN
    try:
        return SpanType(str(raw))
    except ValueError:
        return SpanType.UNKNOWN


class PendoSpanExporter(SpanExporter):
    """
    Exports spans to Pendo's ingestion endpoint as agent_trace events.

    Converts OTel ReadableSpan objects into Pendo's agentic event format,
    JZB-encodes them, and POSTs to /data/agentic/<apiKey>.

    Called by BatchSpanProcessor when its buffer flushes. This runs on
    the BatchSpanProcessor's background thread, so the agent's main
    thread never blocks.

    Args:
        api_key: Pendo integration key (also used in the URL path).
        endpoint: Base URL for the ingestion endpoint.
        redact_pii: Whether to run PII regex redaction before export.
        visitor_id: Pendo visitor ID stamped on every event.
        account_id: Pendo account ID stamped on every event.
    """

    def __init__(
        self,
        api_key: str,
        endpoint: str,
        redact_pii: bool = True,
        visitor_id: Optional[str] = None,
        account_id: Optional[str] = None,
        agent_id: Optional[str] = None,
        url: Optional[str] = None,
    ):
        self._api_key = api_key
        self._endpoint = endpoint.rstrip("/")
        self._redact_pii = redact_pii
        self._visitor_id = visitor_id
        self._account_id = account_id
        self._agent_id = agent_id
        self._url = url

    def set_identity(
        self,
        visitor_id: Optional[str] = None,
        account_id: Optional[str] = None,
    ) -> None:
        """Update identity after construction (e.g. when user logs in)."""
        if visitor_id is not None:
            self._visitor_id = visitor_id
        if account_id is not None:
            self._account_id = account_id

    def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
        """
        Called by OTel BatchSpanProcessor when the buffer flushes.

        Flow: OTel spans → Pendo wire events → JZB encode → HTTP GET
        """
        try:
            logger.debug(
                "[OTel → Pendo] BatchSpanProcessor flushed %d span(s)", len(spans)
            )
            wire_events = []
            for seq, span in enumerate(spans, start=1):
                attrs = dict(span.attributes or {})
                span_type = attrs.get("pendo.span.type", "UNKNOWN")
                logger.debug(
                    "  Span %d: %s (%s) trace=%s",
                    seq, span.name, span_type,
                    format(span.context.trace_id, "032x")[:12] + "...",
                )
                if self._redact_pii:
                    attrs = _redact_attributes(attrs)
                wire_events.extend(self._build_wire_events(span, attrs, sequence=seq))

            if wire_events:
                self._send_jzb(wire_events)

            return SpanExportResult.SUCCESS
        except Exception:
            logger.exception("Failed to export %d spans", len(spans))
            return SpanExportResult.FAILURE

    def _build_event(self, span: ReadableSpan, attrs: dict) -> AgentEvent:
        """Convert an OTel span + redacted attributes into an AgentEvent.

        Kept for backward compatibility (dry-run exporter uses this).
        """
        prop_fields: Dict[str, Any] = {}
        extra: Dict[str, Any] = {}

        for key, value in attrs.items():
            if key in _EVENT_LEVEL_ATTRS:
                continue
            elif key in _ATTR_TO_PROP:
                mapped_key = _ATTR_TO_PROP[key]
                # model needs to be a list now
                if mapped_key == "model" and isinstance(value, str):
                    value = [value]
                prop_fields[mapped_key] = value
            else:
                extra[key] = value

        if extra:
            prop_fields["extra"] = extra

        props = EventProps(**prop_fields)

        start_ms = _ns_to_ms(span.start_time)
        end_ms = _ns_to_ms(span.end_time)
        duration_ms = None
        if start_ms is not None and end_ms is not None:
            duration_ms = round(end_ms - start_ms, 1)

        error_message = None
        if span.status.status_code == StatusCode.ERROR:
            error_message = span.status.description

        # Extract tags from span attributes
        tags = None
        raw_tags = attrs.get("pendo.tags")
        if raw_tags and isinstance(raw_tags, str):
            try:
                tags = json.loads(raw_tags)
            except (json.JSONDecodeError, TypeError):
                pass
        elif raw_tags and isinstance(raw_tags, dict):
            tags = raw_tags

        return AgentEvent(
            visitor_id=self._visitor_id or "anonymous",
            account_id=self._account_id,
            trace_id=format(span.context.trace_id, "032x"),
            span_id=format(span.context.span_id, "016x"),
            parent_span_id=(
                format(span.parent.span_id, "016x")
                if span.parent else None
            ),
            conversation_id=attrs.get("pendo.conversation_id"),
            event_type=attrs.get("pendo.event_type"),
            type=_resolve_span_type(attrs.get("pendo.span.type")),
            name=span.name,
            status=_otel_status_to_schema(span.status.status_code),
            error_message=error_message,
            start_time=start_ms,
            end_time=end_ms,
            duration_ms=duration_ms,
            props=props,
            tags=tags,
        )

    def _build_wire_events(self, span: ReadableSpan, attrs: dict, sequence: int) -> List[dict]:
        """Convert an OTel span into Pendo wire events.

        Returns a list of events:
        - For USER_MESSAGE/REACTION spans: one conversation event (prompt/user_reaction)
        - For all spans: one agent_trace event with full trace fields
        - Conversation events only carry agentTraceId for linking
        - Trace events carry all trace fields + agentTraceContent
        """
        span_type_raw = str(attrs.get("pendo.span.type", "UNKNOWN"))
        conversation_id = attrs.get("pendo.conversation_id", "")
        model_name = attrs.get("llm.model_name", "")
        content = attrs.get("message.content", "")
        trace_content = attrs.get("message.trace_content", "")  # Full LLM prompt for agentTraceContent
        role = attrs.get("message.role", "")
        suggested_prompt = bool(attrs.get("message.suggested_prompt", False))
        file_uploaded = bool(attrs.get("message.file_uploaded", False))
        feedback_comment = attrs.get("reaction.feedback_comment", "")
        reaction_type = attrs.get("reaction.type", "")

        tool_names = []
        tool_name = attrs.get("tool.name")
        if tool_name:
            tool_names.append(tool_name)

        trace_type = _SPAN_TYPE_TO_TRACE_TYPE.get(span_type_raw, "llm_request")

        trace_id = format(span.context.trace_id, "032x")
        span_id = format(span.context.span_id, "016x")
        parent_span_id = (
            format(span.parent.span_id, "016x")
            if span.parent else ""
        )

        browser_time = _ns_to_ms(span.start_time) or int(time.time() * 1000)
        message_id = str(uuid.uuid4())
        visitor_id = self._visitor_id or "anonymous"
        account_id = self._account_id or ""
        url = self._url or ""

        events = []

        # --- Conversation event (prompt / agent_response / user_reaction) ---
        # Only emit for user messages and reactions — these feed agentConversations pipeline
        is_user_message = role == "user" or span_type_raw == "USER_MESSAGE"
        is_reaction = span_type_raw == "REACTION" or reaction_type

        if is_user_message or is_reaction:
            conv_type = "user_reaction" if is_reaction else "prompt"
            conv_props: Dict[str, Any] = {
                "agentId": self._agent_id or "",
                "agentType": "conversation",
                "conversationId": conversation_id or "",
                "messageId": message_id,
                "content": reaction_type if is_reaction else (str(content)[:500] if content else ""),
                "modelUsed": str(model_name) if model_name else "",
                "suggestedPrompt": suggested_prompt,
                "toolsUsed": tool_names,
                "fileUploaded": file_uploaded,
                # Only agentTraceId for linking to trace — no other trace fields
                "agentTraceId": trace_id,
                "visitorId": visitor_id,
                "accountId": account_id,
            }
            if feedback_comment:
                conv_props["feedbackComment"] = feedback_comment

            events.append({
                "type": conv_type,
                "visitor_id": visitor_id,
                "account_id": account_id,
                "url": url,
                "timestamp": browser_time,
                "props": conv_props,
            })

        # --- Trace event (agent_trace) — always emitted ---
        trace_props: Dict[str, Any] = {
            "agentId": self._agent_id or "",
            "agentType": "conversation",
            "agentSpanId": span_id,
            "agentParentSpanId": parent_span_id,
            "agentTraceId": trace_id,
            "agentTraceType": trace_type,
            "agentTraceContent": str(trace_content)[:500] if trace_content else (str(content)[:500] if content else span.name),
            "conversationId": conversation_id or "",
            "messageId": message_id,
            "modelUsed": str(model_name) if model_name else "",
            "suggestedPrompt": suggested_prompt,
            "toolsUsed": tool_names,
            "fileUploaded": file_uploaded,
            "visitorId": visitor_id,
            "accountId": account_id,
        }

        raw_tags = attrs.get("pendo.tags")
        if raw_tags:
            if isinstance(raw_tags, str):
                try:
                    trace_props["tags"] = json.loads(raw_tags)
                except (json.JSONDecodeError, TypeError):
                    pass
            elif isinstance(raw_tags, dict):
                trace_props["tags"] = raw_tags

        events.append({
            "type": "agent_trace",
            "visitor_id": visitor_id,
            "account_id": account_id,
            "url": url,
            "timestamp": browser_time,
            "props": trace_props,
        })

        return events

    def _send_jzb(self, events: List[dict]) -> None:
        """
        JZB-encode events and send to the Pendo ingestion endpoint.

        Uses GET with JZB as query param (matching the Go agentic-send tool):
        GET /data/agentic/<apiKey>?v=<version>&ct=<time>&jzb=<jzb>

        Falls back to POST for large payloads (JZB > 8000 chars).
        """
        import urllib.error
        import urllib.request

        # Populate bytes field on each event
        for evt in events:
            evt["bytes"] = len(json.dumps(evt, separators=(",", ":")).encode("utf-8"))

        jzb = _to_jzb(events)
        ct = int(time.time() * 1000)

        base_path = f"{self._endpoint}/data/agentic/{self._api_key}"

        # GET is preferred; POST for large payloads
        if len(jzb) <= 8000:
            url = f"{base_path}?v={_SDK_VERSION}&ct={ct}&jzb={jzb}"
            req = urllib.request.Request(url, method="GET")
        else:
            url = f"{base_path}?v={_SDK_VERSION}&ct={ct}&s={len(jzb)}"
            body = jzb.encode("utf-8")
            req = urllib.request.Request(
                url,
                data=body,
                headers={"Content-Type": "text/plain"},
                method="POST",
            )

        num_events = len(events)
        log_url = base_path
        try:
            with urllib.request.urlopen(req, timeout=10) as resp:
                if resp.status >= 400:
                    logger.warning(
                        "Ingestion returned %d for %d event(s)",
                        resp.status,
                        num_events,
                    )
                else:
                    logger.info(
                        "Exported %d event(s) to %s [%d]",
                        num_events,
                        log_url,
                        resp.status,
                    )
        except urllib.error.HTTPError as exc:
            logger.warning(
                "Ingestion returned HTTP %d for %d event(s) to %s: %s",
                exc.code,
                num_events,
                log_url,
                exc.reason,
            )
        except Exception as exc:
            logger.warning(
                "Failed to send %d event(s) to %s: %s",
                num_events,
                log_url,
                exc,
            )

    def shutdown(self) -> None:
        """Clean up resources. Called when the TracerProvider shuts down."""
        pass

    def force_flush(self, timeout_millis: int = 30000) -> bool:
        """Force-flush is handled by BatchSpanProcessor, not the exporter."""
        return True
