from __future__ import annotations

import json
import time
from datetime import datetime, timezone
from typing import Any, Callable, Dict, Mapping, Optional, TypeVar, Union
from urllib.error import HTTPError, URLError
from urllib.request import Request, urlopen
from uuid import uuid4

SpanPayload = Dict[str, Any]
RequestSender = Callable[[str, str, Mapping[str, Any], float], Dict[str, Any]]
T = TypeVar("T")


class IngestionError(RuntimeError):
    """Raised when Refario ingestion fails."""


class Run:
    def __init__(
        self,
        run_id: str,
        workflow: str,
        environment: Optional[str],
        endpoint: str,
        api_key: str,
        timeout_seconds: float,
        request_sender: RequestSender,
    ) -> None:
        self.run_id = run_id
        self.workflow = workflow
        self.environment = environment
        self._endpoint = endpoint
        self._api_key = api_key
        self._timeout_seconds = timeout_seconds
        self._request_sender = request_sender
        self._started_at = datetime.now(timezone.utc)
        self._spans: list[SpanPayload] = []

    @property
    def spans(self) -> list[SpanPayload]:
        return [dict(span) for span in self._spans]

    def llm_call(self, span: Optional[Mapping[str, Any]] = None, **kwargs: Any) -> None:
        payload = _merge_span_input(span, kwargs)
        self._spans.append(_normalize_span_payload("llm", payload))

    def tool_call(self, span: Optional[Mapping[str, Any]] = None, **kwargs: Any) -> None:
        payload = _merge_span_input(span, kwargs)
        self._spans.append(_normalize_span_payload("tool", payload))

    def chain_call(self, span: Optional[Mapping[str, Any]] = None, **kwargs: Any) -> None:
        payload = _merge_span_input(span, kwargs)
        self._spans.append(_normalize_span_payload("chain", payload))

    def llmCall(self, span: Mapping[str, Any]) -> None:
        self.llm_call(span)

    def toolCall(self, span: Mapping[str, Any]) -> None:
        self.tool_call(span)

    def chainCall(self, span: Mapping[str, Any]) -> None:
        self.chain_call(span)

    def end(self, success: bool) -> Dict[str, Any]:
        event: Dict[str, Any] = {
            "run_id": self.run_id,
            "idempotency_key": self.run_id,
            "workflow": self.workflow,
            "status": "success" if success else "error",
            "started_at": _to_iso8601(self._started_at),
            "ended_at": _to_iso8601(datetime.now(timezone.utc)),
            "spans": self._spans,
        }
        if self.environment:
            event["environment"] = self.environment

        return self._request_sender(self._endpoint, self._api_key, event, self._timeout_seconds)


class McpTelemetrySession:
    def __init__(
        self,
        base: Mapping[str, Any],
        emit_event: Callable[[Mapping[str, Any]], Dict[str, Any]],
    ) -> None:
        self._base = dict(base)
        self._emit_event = emit_event
        self._spans: list[SpanPayload] = []

    def span(self, span: Optional[Mapping[str, Any]] = None, **kwargs: Any) -> None:
        payload = _merge_span_input(span, kwargs)
        self._spans.append(payload)

    def tool_call(
        self,
        *,
        tool_name: str,
        correlation_id: Optional[str] = None,
        input_sample: Optional[str] = None,
        metadata: Optional[Mapping[str, Any]] = None,
        at: Optional[Union[str, datetime]] = None,
    ) -> None:
        span: SpanPayload = {
            "kind": "tool_call",
            "tool_name": tool_name,
            "status": "success",
        }
        if correlation_id is not None:
            span["correlation_id"] = correlation_id
        if input_sample is not None:
            span["input_sample"] = input_sample
        if metadata is not None:
            span["metadata"] = dict(metadata)
        if at is not None:
            span["at"] = at
        self.span(span)

    def tool_result(
        self,
        *,
        tool_name: str,
        correlation_id: Optional[str] = None,
        latency_ms: Optional[int] = None,
        output_sample: Optional[str] = None,
        metadata: Optional[Mapping[str, Any]] = None,
        at: Optional[Union[str, datetime]] = None,
    ) -> None:
        span: SpanPayload = {
            "kind": "tool_result",
            "tool_name": tool_name,
            "status": "success",
        }
        if correlation_id is not None:
            span["correlation_id"] = correlation_id
        if latency_ms is not None:
            span["latency_ms"] = latency_ms
        if output_sample is not None:
            span["output_sample"] = output_sample
        if metadata is not None:
            span["metadata"] = dict(metadata)
        if at is not None:
            span["at"] = at
        self.span(span)

    def tool_error(
        self,
        *,
        tool_name: str,
        error: str,
        correlation_id: Optional[str] = None,
        latency_ms: Optional[int] = None,
        metadata: Optional[Mapping[str, Any]] = None,
        at: Optional[Union[str, datetime]] = None,
    ) -> None:
        span: SpanPayload = {
            "kind": "tool_error",
            "tool_name": tool_name,
            "error": error,
            "status": "error",
        }
        if correlation_id is not None:
            span["correlation_id"] = correlation_id
        if latency_ms is not None:
            span["latency_ms"] = latency_ms
        if metadata is not None:
            span["metadata"] = dict(metadata)
        if at is not None:
            span["at"] = at
        self.span(span)

    def flush(
        self,
        success: bool = True,
        *,
        ended_at: Optional[Union[str, datetime]] = None,
    ) -> Dict[str, Any]:
        payload = dict(self._base)
        payload["status"] = "success" if success else "error"
        payload["ended_at"] = ended_at or datetime.now(timezone.utc)
        payload["spans"] = [dict(span) for span in self._spans]
        return self._emit_event(payload)

    # TS-style aliases
    def toolCall(self, span: Mapping[str, Any]) -> None:
        payload = dict(span)
        if "toolName" in payload and "tool_name" not in payload:
            payload["tool_name"] = payload.pop("toolName")
        if "correlationId" in payload and "correlation_id" not in payload:
            payload["correlation_id"] = payload.pop("correlationId")
        if "inputSample" in payload and "input_sample" not in payload:
            payload["input_sample"] = payload.pop("inputSample")
        self.tool_call(**payload)

    def toolResult(self, span: Mapping[str, Any]) -> None:
        payload = dict(span)
        if "toolName" in payload and "tool_name" not in payload:
            payload["tool_name"] = payload.pop("toolName")
        if "correlationId" in payload and "correlation_id" not in payload:
            payload["correlation_id"] = payload.pop("correlationId")
        if "latencyMs" in payload and "latency_ms" not in payload:
            payload["latency_ms"] = payload.pop("latencyMs")
        if "outputSample" in payload and "output_sample" not in payload:
            payload["output_sample"] = payload.pop("outputSample")
        self.tool_result(**payload)

    def toolError(self, span: Mapping[str, Any]) -> None:
        payload = dict(span)
        if "toolName" in payload and "tool_name" not in payload:
            payload["tool_name"] = payload.pop("toolName")
        if "correlationId" in payload and "correlation_id" not in payload:
            payload["correlation_id"] = payload.pop("correlationId")
        if "latencyMs" in payload and "latency_ms" not in payload:
            payload["latency_ms"] = payload.pop("latencyMs")
        self.tool_error(**payload)


class McpTelemetryEmitter:
    def __init__(
        self,
        *,
        endpoint: str,
        api_key: str,
        default_workflow: Optional[str] = None,
        default_environment: Optional[str] = None,
        default_transport: Optional[str] = None,
        timeout_seconds: float = 10.0,
        request_sender: Optional[RequestSender] = None,
    ) -> None:
        endpoint = endpoint.strip().rstrip("/")
        if not endpoint:
            raise ValueError("endpoint is required")
        if not api_key:
            raise ValueError("api_key is required")
        if timeout_seconds <= 0:
            raise ValueError("timeout_seconds must be > 0")

        self.endpoint = endpoint
        self.api_key = api_key
        self.default_workflow = default_workflow
        self.default_environment = default_environment
        self.default_transport = default_transport
        self.timeout_seconds = timeout_seconds
        self._request_sender = request_sender or _send_mcp_event

    def emit_event(self, event: Optional[Mapping[str, Any]] = None, **kwargs: Any) -> Dict[str, Any]:
        payload = _merge_span_input(event, kwargs)
        session_id = _get_value(payload, "session_id", "sessionId", "mcp_session_id")
        if not isinstance(session_id, str) or not session_id.strip():
            raise ValueError("session_id is required")

        workflow = _get_value(payload, "workflow")
        if workflow is None:
            workflow = self.default_workflow or "mcp-session"

        spans = payload.get("spans") or []
        if not isinstance(spans, list):
            raise ValueError("spans must be a list")

        wire_event: Dict[str, Any] = {
            "mcp_session_id": session_id,
            "workflow": workflow,
            "spans": [_to_mcp_wire_span(span) for span in spans],
        }

        optional_mappings = [
            (("request_id", "requestId"), "request_id"),
            (("run_id", "runId"), "run_id"),
            (("idempotency_key", "idempotencyKey"), "idempotency_key"),
            (("environment",), "environment"),
            (("transport",), "transport"),
            (("status",), "status"),
        ]
        for source_keys, target_key in optional_mappings:
            value = _get_value(payload, *source_keys)
            if value is not None:
                wire_event[target_key] = value

        if "environment" not in wire_event and self.default_environment is not None:
            wire_event["environment"] = self.default_environment
        if "transport" not in wire_event and self.default_transport is not None:
            wire_event["transport"] = self.default_transport

        started_at = _get_value(payload, "started_at", "startedAt")
        ended_at = _get_value(payload, "ended_at", "endedAt")
        if started_at is not None:
            wire_event["started_at"] = _to_time_value(started_at)
        if ended_at is not None:
            wire_event["ended_at"] = _to_time_value(ended_at)

        return self._request_sender(self.endpoint, self.api_key, wire_event, self.timeout_seconds)

    def start_session(
        self,
        *,
        session_id: Optional[str] = None,
        request_id: Optional[str] = None,
        run_id: Optional[str] = None,
        idempotency_key: Optional[str] = None,
        workflow: Optional[str] = None,
        environment: Optional[str] = None,
        transport: Optional[str] = None,
        started_at: Optional[Union[str, datetime]] = None,
    ) -> McpTelemetrySession:
        base: Dict[str, Any] = {
            "session_id": session_id or str(uuid4()),
            "workflow": workflow or self.default_workflow or "mcp-session",
            "started_at": started_at or datetime.now(timezone.utc),
        }
        if request_id is not None:
            base["request_id"] = request_id
        if run_id is not None:
            base["run_id"] = run_id
        if idempotency_key is not None:
            base["idempotency_key"] = idempotency_key
        env_value = environment if environment is not None else self.default_environment
        if env_value is not None:
            base["environment"] = env_value
        transport_value = transport if transport is not None else self.default_transport
        if transport_value is not None:
            base["transport"] = transport_value

        return McpTelemetrySession(base=base, emit_event=self.emit_event)

    # TS-style aliases
    def emitEvent(self, event: Mapping[str, Any]) -> Dict[str, Any]:
        return self.emit_event(event)

    def startSession(self, event: Mapping[str, Any]) -> McpTelemetrySession:
        payload = dict(event)
        if "sessionId" in payload and "session_id" not in payload:
            payload["session_id"] = payload.pop("sessionId")
        if "requestId" in payload and "request_id" not in payload:
            payload["request_id"] = payload.pop("requestId")
        if "runId" in payload and "run_id" not in payload:
            payload["run_id"] = payload.pop("runId")
        if "idempotencyKey" in payload and "idempotency_key" not in payload:
            payload["idempotency_key"] = payload.pop("idempotencyKey")
        if "startedAt" in payload and "started_at" not in payload:
            payload["started_at"] = payload.pop("startedAt")
        return self.start_session(**payload)


class RefarioClient:
    def __init__(
        self,
        endpoint: str,
        api_key: str,
        timeout_seconds: float = 10.0,
        request_sender: Optional[RequestSender] = None,
    ) -> None:
        endpoint = endpoint.strip().rstrip("/")
        if not endpoint:
            raise ValueError("endpoint is required")
        if not api_key:
            raise ValueError("api_key is required")
        if timeout_seconds <= 0:
            raise ValueError("timeout_seconds must be > 0")

        self.endpoint = endpoint
        self.api_key = api_key
        self.timeout_seconds = timeout_seconds
        self._request_sender = request_sender or _send_event

    def start_run(
        self,
        workflow: Union[str, Mapping[str, Any]],
        environment: Optional[str] = None,
    ) -> Run:
        if isinstance(workflow, Mapping):
            run_args = dict(workflow)
            workflow_name = run_args.get("workflow")
            environment = run_args.get("environment")
        else:
            workflow_name = workflow

        if not isinstance(workflow_name, str) or not workflow_name.strip():
            raise ValueError("workflow is required")

        return Run(
            run_id=str(uuid4()),
            workflow=workflow_name.strip(),
            environment=environment,
            endpoint=self.endpoint,
            api_key=self.api_key,
            timeout_seconds=self.timeout_seconds,
            request_sender=self._request_sender,
        )

    def startRun(
        self,
        workflow: Union[str, Mapping[str, Any]],
        environment: Optional[str] = None,
    ) -> Run:
        return self.start_run(workflow, environment=environment)


class _RefarioNamespace:
    def init(
        self,
        config: Optional[Mapping[str, Any]] = None,
        *,
        endpoint: Optional[str] = None,
        api_key: Optional[str] = None,
        timeout_seconds: float = 10.0,
    ) -> RefarioClient:
        if config is not None:
            endpoint = config.get("endpoint", endpoint)
            api_key = config.get("apiKey", config.get("api_key", api_key))
            timeout_seconds = config.get("timeout_seconds", timeout_seconds)

        if endpoint is None:
            raise ValueError("endpoint is required")
        if api_key is None:
            raise ValueError("api_key is required")

        return RefarioClient(
            endpoint=endpoint,
            api_key=api_key,
            timeout_seconds=float(timeout_seconds),
        )


refario = _RefarioNamespace()


def create_mcp_telemetry_emitter(
    config: Optional[Mapping[str, Any]] = None,
    *,
    endpoint: Optional[str] = None,
    api_key: Optional[str] = None,
    default_workflow: Optional[str] = None,
    default_environment: Optional[str] = None,
    default_transport: Optional[str] = None,
    timeout_seconds: float = 10.0,
) -> McpTelemetryEmitter:
    if config is not None:
        endpoint = config.get("endpoint", endpoint)
        api_key = config.get("apiKey", config.get("api_key", api_key))
        default_workflow = config.get("defaultWorkflow", config.get("default_workflow", default_workflow))
        default_environment = config.get(
            "defaultEnvironment",
            config.get("default_environment", default_environment),
        )
        default_transport = config.get("defaultTransport", config.get("default_transport", default_transport))
        timeout_seconds = config.get("timeout_seconds", timeout_seconds)

    if endpoint is None:
        raise ValueError("endpoint is required")
    if api_key is None:
        raise ValueError("api_key is required")

    return McpTelemetryEmitter(
        endpoint=endpoint,
        api_key=api_key,
        default_workflow=default_workflow,
        default_environment=default_environment,
        default_transport=default_transport,
        timeout_seconds=float(timeout_seconds),
    )


def parse_usage(provider: Optional[str], usage: Any) -> Dict[str, int]:
    if not provider or usage is None:
        return {}

    if provider == "openai":
        prompt = _safe_int(_usage_get(usage, "prompt_tokens", "input_tokens"), default=0)
        completion = _safe_int(_usage_get(usage, "completion_tokens", "output_tokens"), default=0)
        total = _safe_int(_usage_get(usage, "total_tokens"), default=prompt + completion)
        return {
            "prompt_tokens": prompt,
            "completion_tokens": completion,
            "total_tokens": total,
        }

    if provider == "anthropic":
        prompt = _safe_int(_usage_get(usage, "input_tokens"), default=0)
        completion = _safe_int(_usage_get(usage, "output_tokens"), default=0)
        return {
            "prompt_tokens": prompt,
            "completion_tokens": completion,
            "total_tokens": prompt + completion,
        }

    return {}


def traced_llm(
    run: Any,
    *,
    provider: str,
    model: str,
    name: str,
    fn: Callable[[], T],
) -> T:
    start = time.perf_counter()
    try:
        result = fn()
        tokens = parse_usage(provider, _extract_usage(result))
        run.llm_call(
            name=name,
            provider=provider,
            model=model,
            status="success",
            latency_ms=int((time.perf_counter() - start) * 1000),
            **tokens,
        )
        return result
    except Exception as error:
        run.llm_call(
            name=name,
            provider=provider,
            model=model,
            status="error",
            latency_ms=int((time.perf_counter() - start) * 1000),
            error=str(error),
        )
        raise


def _send_event(
    endpoint: str,
    api_key: str,
    payload: Mapping[str, Any],
    timeout_seconds: float,
) -> Dict[str, Any]:
    return _send_event_to_path(endpoint, api_key, payload, timeout_seconds, "/api/v1/events")


def _send_mcp_event(
    endpoint: str,
    api_key: str,
    payload: Mapping[str, Any],
    timeout_seconds: float,
) -> Dict[str, Any]:
    return _send_event_to_path(endpoint, api_key, payload, timeout_seconds, "/api/v1/events/mcp")


def _send_event_to_path(
    endpoint: str,
    api_key: str,
    payload: Mapping[str, Any],
    timeout_seconds: float,
    path: str,
) -> Dict[str, Any]:
    url = f"{endpoint.rstrip('/')}{path}"
    body = json.dumps(payload).encode("utf-8")
    request = Request(url=url, data=body, method="POST")
    request.add_header("Content-Type", "application/json")
    request.add_header("x-refario-key", api_key)

    try:
        with urlopen(request, timeout=timeout_seconds) as response:
            response_body = response.read()
            status_code = response.status if hasattr(response, "status") else response.getcode()
    except HTTPError as error:
        detail = _decode_body(error.read())
        raise IngestionError(
            f"Ingestion failed with status {error.code}: {detail or error.reason}"
        ) from error
    except URLError as error:
        raise IngestionError(f"Ingestion request failed: {error.reason}") from error

    decoded = _decode_body(response_body)
    if not decoded:
        return {"status": status_code}
    try:
        parsed = json.loads(decoded)
        if isinstance(parsed, dict):
            return parsed
        return {"status": status_code, "data": parsed}
    except json.JSONDecodeError:
        return {"status": status_code, "body": decoded}


def _merge_span_input(
    span: Optional[Mapping[str, Any]],
    kwargs: Mapping[str, Any],
) -> SpanPayload:
    if span is None:
        payload: SpanPayload = {}
    else:
        payload = dict(span)
    payload.update(kwargs)
    return payload


def _normalize_span_payload(span_type: str, payload: SpanPayload) -> SpanPayload:
    if not isinstance(payload.get("name"), str) or not payload["name"].strip():
        raise ValueError("span name is required")

    normalized = dict(payload)
    normalized["type"] = span_type

    prompt_tokens = _coerce_non_negative_int(normalized.get("prompt_tokens"), "prompt_tokens")
    completion_tokens = _coerce_non_negative_int(
        normalized.get("completion_tokens"),
        "completion_tokens",
    )
    total_tokens = _coerce_non_negative_int(normalized.get("total_tokens"), "total_tokens")
    if total_tokens is None and prompt_tokens is not None and completion_tokens is not None:
        total_tokens = prompt_tokens + completion_tokens

    if prompt_tokens is not None:
        normalized["prompt_tokens"] = prompt_tokens
    if completion_tokens is not None:
        normalized["completion_tokens"] = completion_tokens
    if total_tokens is not None:
        normalized["total_tokens"] = total_tokens

    return normalized


def _get_value(payload: Mapping[str, Any], *keys: str) -> Any:
    for key in keys:
        if key in payload:
            return payload[key]
    return None


def _to_mcp_wire_span(span: Mapping[str, Any]) -> Dict[str, Any]:
    mapped_fields = [
        (("index",), "index"),
        (("kind",), "kind"),
        (("name",), "name"),
        (("method",), "method"),
        (("subtype",), "subtype"),
        (("provider",), "provider"),
        (("model",), "model"),
        (("cost_cents", "costCents"), "cost_cents"),
        (("latency_ms", "latencyMs"), "latency_ms"),
        (("status",), "status"),
        (("error",), "error"),
        (("input_sample", "inputSample"), "input_sample"),
        (("output_sample", "outputSample"), "output_sample"),
        (("prompt_tokens", "promptTokens"), "prompt_tokens"),
        (("completion_tokens", "completionTokens"), "completion_tokens"),
        (("total_tokens", "totalTokens"), "total_tokens"),
        (("prompt_cost_cents", "promptCostCents"), "prompt_cost_cents"),
        (("completion_cost_cents", "completionCostCents"), "completion_cost_cents"),
        (("metadata",), "metadata"),
    ]
    out: Dict[str, Any] = {}
    for source_keys, target_key in mapped_fields:
        value = _get_value(span, *source_keys)
        if value is not None:
            out[target_key] = value

    at_value = _get_value(span, "at")
    if at_value is not None:
        out["at"] = _to_time_value(at_value)

    tool_name = _get_value(span, "tool_name", "toolName")
    correlation_id = _get_value(span, "correlation_id", "correlationId")
    parent_span_id = _get_value(span, "parent_span_id", "parentSpanId")
    if tool_name is not None:
        out["tool_name"] = tool_name
    if correlation_id is not None:
        out["correlation_id"] = correlation_id
    if parent_span_id is not None:
        out["parent_span_id"] = parent_span_id
    return out


def _coerce_non_negative_int(value: Any, field_name: str) -> Optional[int]:
    if value is None:
        return None
    if isinstance(value, bool) or not isinstance(value, int):
        raise ValueError(f"{field_name} must be an integer")
    if value < 0:
        raise ValueError(f"{field_name} must be non-negative")
    return value


def _usage_get(usage: Any, *keys: str) -> Any:
    if usage is None:
        return None
    if isinstance(usage, Mapping):
        for key in keys:
            if key in usage:
                return usage[key]
        return None

    for key in keys:
        if hasattr(usage, key):
            return getattr(usage, key)
    return None


def _extract_usage(result: Any) -> Any:
    if result is None:
        return None
    if isinstance(result, Mapping):
        return result.get("usage")
    return getattr(result, "usage", None)


def _safe_int(value: Any, default: int) -> int:
    if value is None:
        return default
    if isinstance(value, bool):
        return default
    if isinstance(value, int):
        return value
    try:
        return int(value)
    except (TypeError, ValueError):
        return default


def _to_iso8601(value: datetime) -> str:
    return value.astimezone(timezone.utc).isoformat(timespec="milliseconds").replace("+00:00", "Z")


def _to_time_value(value: Union[str, datetime]) -> str:
    if isinstance(value, datetime):
        return _to_iso8601(value)
    if isinstance(value, str):
        return value
    raise ValueError("time value must be a datetime or ISO string")


def _decode_body(body: Any) -> str:
    if isinstance(body, (bytes, bytearray)):
        return body.decode("utf-8", errors="replace").strip()
    if isinstance(body, str):
        return body.strip()
    return ""
