"""Synchronous HTTP span sender for Claude Code hooks.

Each hook invocation is a short-lived process, so we use a simple synchronous
send with retry rather than the SDK's background BatchTransport thread.
"""
from __future__ import annotations

import logging
import os
import time
from collections import defaultdict
from typing import Any

import httpx

logger = logging.getLogger("staso.claude_code")

_MAX_ATTEMPTS = 3
_RETRY_DELAYS = [0.5, 1.0]  # delays before attempt 2 and 3; no delay after final attempt


class ClaudeCodeSender:
    """Sends span dicts synchronously to the Staso ingest endpoint."""

    def __init__(
        self,
        base_url: str,
        api_key: str,
        environment: str = "default",
        workspace_slug: str = "",
    ) -> None:
        self._url = f"{base_url.rstrip('/')}/v1/traces/ingest"
        self._api_key = api_key
        self._environment = environment
        self._workspace_slug = workspace_slug

    def send(self, spans: list[dict[str, Any]]) -> None:
        if not spans:
            return
        self._send_with_retry(self._build_payload(spans))

    def _build_payload(self, spans: list[dict[str, Any]]) -> dict[str, Any]:
        groups: dict[str, list[dict[str, Any]]] = defaultdict(list)
        for span in spans:
            groups[span.get("trace_id", "")].append(span)

        traces = []
        for trace_id, group in groups.items():
            earliest_ts = min(
                (s.get("timestamp", "") for s in group if s.get("timestamp")), default=""
            )
            trace_spans = [
                {
                    "span_id": s.get("span_id"),
                    "parent_span_id": s.get("parent_span_id"),
                    "name": s.get("name", ""),
                    "kind": s.get("kind", "agent"),
                    "timestamp": s.get("timestamp", ""),
                    "duration_ms": s.get("duration_ms", 0.0),
                    "status": s.get("status", "ok"),
                    "model": s.get("model", ""),
                    "input_tokens": s.get("input_tokens", 0),
                    "output_tokens": s.get("output_tokens", 0),
                    "total_tokens": s.get("total_tokens", 0),
                    "input": s.get("input", ""),
                    "output": s.get("output", ""),
                    "attributes": s.get("attributes", "{}"),
                    "error_message": s.get("error_message"),
                    "session_id": s.get("session_id", ""),
                    "agent_id": s.get("agent_id", ""),
                    "user_id": s.get("user_id", ""),
                }
                for s in group
            ]
            traces.append(
                {
                    "trace_id": trace_id,
                    "timestamp": earliest_ts,
                    "environment": self._environment,
                    "workspace_slug": self._workspace_slug,
                    "spans": trace_spans,
                }
            )
        return {"traces": traces}

    def _send_with_retry(self, payload: dict[str, Any]) -> None:
        with httpx.Client(timeout=10.0) as client:
            for attempt in range(_MAX_ATTEMPTS):
                try:
                    resp = client.post(
                        self._url,
                        json=payload,
                        headers={
                            "X-API-Key": self._api_key,
                            "Content-Type": "application/json",
                        },
                    )
                    if resp.status_code < 500:
                        if resp.status_code >= 400:
                            logger.warning(
                                "staso: ingest returned %d – %s",
                                resp.status_code,
                                resp.text[:300],
                            )
                        return
                    logger.warning(
                        "staso: ingest returned %d, attempt %d/%d",
                        resp.status_code,
                        attempt + 1,
                        _MAX_ATTEMPTS,
                    )
                except Exception as exc:
                    logger.warning(
                        "staso: send error attempt %d/%d: %s", attempt + 1, _MAX_ATTEMPTS, exc
                    )
                if attempt < _MAX_ATTEMPTS - 1:
                    time.sleep(_RETRY_DELAYS[attempt])


def sender_from_env() -> ClaudeCodeSender | None:
    """Create a sender from environment variables. Returns None if STASO_API_KEY is unset."""
    api_key = os.environ.get("STASO_API_KEY", "")
    if not api_key:
        logger.warning("staso: STASO_API_KEY not set, hook is a no-op")
        return None
    return ClaudeCodeSender(
        base_url=os.environ.get("STASO_BASE_URL", "https://api.staso.ai"),
        api_key=api_key,
        environment=os.environ.get("STASO_ENVIRONMENT", "default"),
        workspace_slug=os.environ.get("STASO_WORKSPACE_SLUG", ""),
    )
