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 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 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]:
    url = f"{endpoint.rstrip('/')}/api/v1/events"
    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 _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 _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 ""
