from __future__ import annotations

"""Python SDK for interacting with the SmartBlocks gateway services."""

from dataclasses import dataclass, field
from datetime import datetime, timedelta, timezone
import base64
import json
import os
import time
import uuid
from pathlib import Path
from typing import Any, Callable, Mapping, MutableMapping, Sequence

import httpx
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey

__all__ = [
    "GatewayError",
    "GatewayTransportError",
    "MintedToken",
    "RetryConfig",
    "SigningKey",
    "SlotClosure",
    "SlotCreateRequest",
    "SlotHandle",
    "SlotSummary",
    "Receipt",
    "SbnGatewayClient",
]


# ---------------------------------------------------------------------------
# Utility helpers
# ---------------------------------------------------------------------------


def _b64url(data: bytes) -> str:
    return base64.urlsafe_b64encode(data).rstrip(b"=").decode("ascii")


def _canonical_json(data: Mapping[str, Any]) -> bytes:
    return json.dumps(data, separators=(",", ":"), sort_keys=True).encode("utf-8")


def _now_utc() -> datetime:
    return datetime.now(timezone.utc)


def _parse_datetime(value: str | None) -> datetime | None:
    if not value:
        return None
    try:
        if value.endswith("Z"):
            value = value[:-1] + "+00:00"
        return datetime.fromisoformat(value)
    except ValueError:
        return None


# ---------------------------------------------------------------------------
# Error types
# ---------------------------------------------------------------------------


class GatewayError(RuntimeError):
    """Raised when the gateway responds with an error payload."""

    def __init__(
        self,
        message: str,
        *,
        status_code: int,
        code: str | None = None,
        details: Mapping[str, Any] | None = None,
    ) -> None:
        super().__init__(message)
        self.status_code = status_code
        self.code = code
        self.details = dict(details or {})

    def __str__(self) -> str:  # pragma: no cover - human readable only
        base = super().__str__()
        suffix = f" (status={self.status_code}"
        if self.code:
            suffix += f", code={self.code}"
        if self.details:
            suffix += f", details={self.details}"
        return base + suffix + ")"


class GatewayTransportError(RuntimeError):
    """Raised when the SDK cannot reach the gateway after retries."""

    def __init__(self, message: str, *, last_exception: Exception | None = None) -> None:
        super().__init__(message)
        self.last_exception = last_exception


# ---------------------------------------------------------------------------
# Data models
# ---------------------------------------------------------------------------


@dataclass(slots=True)
class SlotCreateRequest:
    """Payload used to create a new slot via the control plane."""

    worker_id: str
    task_type: str
    generation_interval: int | None = None
    metadata: Mapping[str, Any] | None = None

    def to_payload(self) -> dict[str, Any]:
        payload: dict[str, Any] = {
            "worker_id": self.worker_id,
            "task_type": self.task_type,
        }
        if self.generation_interval is not None:
            payload["generation_interval"] = self.generation_interval
        if self.metadata:
            payload["metadata"] = dict(self.metadata)
        return payload


@dataclass(slots=True)
class SlotHandle:
    """Representation of a slot returned by the control plane."""

    id: str
    worker_id: str
    task_type: str
    state: str
    opened_at: datetime | None
    scheduler_key: str | None = None
    receipt_id: str | None = None
    metadata: Mapping[str, Any] = field(default_factory=dict)

    @classmethod
    def from_dict(cls, data: Mapping[str, Any]) -> "SlotHandle":
        return cls(
            id=str(data.get("id")),
            worker_id=str(data.get("worker_id")),
            task_type=str(data.get("task_type")),
            state=str(data.get("state", "unknown")),
            opened_at=_parse_datetime(data.get("opened_at")),
            scheduler_key=data.get("scheduler_key"),
            receipt_id=data.get("receipt_id"),
            metadata=dict(data.get("metadata") or {}),
        )

    def dict(self) -> dict[str, Any]:
        return {
            "id": self.id,
            "worker_id": self.worker_id,
            "task_type": self.task_type,
            "state": self.state,
            "opened_at": self.opened_at.isoformat() if self.opened_at else None,
            "scheduler_key": self.scheduler_key,
            "receipt_id": self.receipt_id,
            "metadata": dict(self.metadata),
        }


@dataclass(slots=True)
class SlotClosure:
    """Response returned when a slot is closed."""

    slot_id: str
    receipt_id: str
    closed_at: datetime | None

    @classmethod
    def from_dict(cls, data: Mapping[str, Any]) -> "SlotClosure":
        return cls(
            slot_id=str(data.get("slot_id")),
            receipt_id=str(data.get("receipt_id")),
            closed_at=_parse_datetime(data.get("closed_at")),
        )

    def dict(self) -> dict[str, Any]:
        return {
            "slot_id": self.slot_id,
            "receipt_id": self.receipt_id,
            "closed_at": self.closed_at.isoformat() if self.closed_at else None,
        }


@dataclass(slots=True)
class SlotSummary:
    """Slot summary payload required when requesting attestations."""

    slot_id: str
    snap_hash: str
    snap_version: str = "1"
    metadata: Mapping[str, Any] = field(default_factory=dict)
    metrics: Mapping[str, Any] = field(default_factory=dict)

    def to_payload(self) -> dict[str, Any]:
        payload: dict[str, Any] = {
            "slot_id": self.slot_id,
            "snap_hash": self.snap_hash,
            "snap_version": self.snap_version,
        }
        if self.metadata:
            payload["metadata"] = dict(self.metadata)
        if self.metrics:
            payload["metrics"] = dict(self.metrics)
        return payload


@dataclass(slots=True)
class Receipt:
    """Receipt payload fetched from the read API."""

    id: str
    tenant_id: str | None
    subject_id: str
    subject_type: str
    signer_id: str | None
    created_at: datetime | None
    payload: Mapping[str, Any]
    metadata: Mapping[str, Any]
    snapchore_hash: str | None = None

    @classmethod
    def from_dict(cls, data: Mapping[str, Any]) -> "Receipt":
        return cls(
            id=str(data.get("id")),
            tenant_id=data.get("tenant_id"),
            subject_id=str(data.get("subject_id")),
            subject_type=str(data.get("subject_type")),
            signer_id=data.get("signer_id"),
            created_at=_parse_datetime(data.get("created_at")),
            payload=dict(data.get("payload") or {}),
            metadata=dict(data.get("metadata") or {}),
            snapchore_hash=data.get("snapchore_hash"),
        )

    def dict(self) -> dict[str, Any]:
        return {
            "id": self.id,
            "tenant_id": self.tenant_id,
            "subject_id": self.subject_id,
            "subject_type": self.subject_type,
            "signer_id": self.signer_id,
            "created_at": self.created_at.isoformat() if self.created_at else None,
            "payload": dict(self.payload),
            "metadata": dict(self.metadata),
            "snapchore_hash": self.snapchore_hash,
        }


@dataclass(slots=True)
class MintedToken:
    """Represents a freshly minted bearer token and its expiry."""

    token: str
    expires_at: datetime

    @classmethod
    def from_ttl(cls, token: str, ttl_seconds: int) -> "MintedToken":
        return cls(token=token, expires_at=_now_utc() + timedelta(seconds=ttl_seconds))


@dataclass(slots=True)
class RetryConfig:
    """Retry configuration for transient failures."""

    attempts: int = 3
    backoff_factor: float = 0.5
    status_forcelist: Sequence[int] = (500, 502, 503, 504)
    retry_on_methods: Sequence[str] = ("GET", "POST", "PUT", "DELETE", "PATCH")


# ---------------------------------------------------------------------------
# Signing helpers
# ---------------------------------------------------------------------------


@dataclass(slots=True)
class SigningKey:
    """Ed25519 signing helper used for JWT minting and payload signatures."""

    private_key: Ed25519PrivateKey
    issuer: str
    audience: str
    kid: str = "sbn-ed25519"

    @classmethod
    def from_pem(
        cls,
        source: str | bytes | os.PathLike[str],
        *,
        issuer: str,
        audience: str,
        kid: str = "sbn-ed25519",
    ) -> "SigningKey":
        if isinstance(source, (str, os.PathLike)):
            text = str(source)
            path = Path(text)
            if path.exists():
                pem_data = path.read_bytes()
            else:
                pem_data = text.encode("utf-8")
        else:
            pem_data = bytes(source)
        key = serialization.load_pem_private_key(pem_data, password=None)
        if not isinstance(key, Ed25519PrivateKey):  # pragma: no cover - defensive
            raise TypeError("Expected an Ed25519 private key")
        return cls(private_key=key, issuer=issuer, audience=audience, kid=kid)

    @classmethod
    def generate(
        cls,
        *,
        issuer: str,
        audience: str,
        kid: str = "sbn-ed25519",
    ) -> "SigningKey":
        return cls(private_key=Ed25519PrivateKey.generate(), issuer=issuer, audience=audience, kid=kid)

    def issue_token(
        self,
        *,
        subject: str,
        scopes: Sequence[str],
        cdna: str,
        ttl_seconds: int,
        tenant_id: str | None = None,
        role: str = "agent",
    ) -> MintedToken:
        if ttl_seconds <= 0:
            raise ValueError("ttl_seconds must be positive")

        issued_at = _now_utc()
        scope_list = sorted({scope.strip() for scope in scopes if scope})
        header = {"alg": "EdDSA", "typ": "JWT", "kid": self.kid}
        payload: dict[str, Any] = {
            "iss": self.issuer,
            "aud": self.audience,
            "sub": subject,
            "role": role,
            "scope": scope_list,
            "cdna": cdna,
            "iat": int(issued_at.timestamp()),
            "exp": int((issued_at + timedelta(seconds=ttl_seconds)).timestamp()),
            "jti": uuid.uuid4().hex,
        }
        if tenant_id:
            payload["tenant_id"] = tenant_id

        signing_input = f"{_b64url(_canonical_json(header))}.{_b64url(_canonical_json(payload))}"
        signature = self.private_key.sign(signing_input.encode("utf-8"))
        token = f"{signing_input}.{_b64url(signature)}"
        return MintedToken.from_ttl(token, ttl_seconds)

    def sign_summary(self, summary: SlotSummary) -> str:
        payload = summary.to_payload()
        canonical = _canonical_json(payload)
        signature = self.private_key.sign(canonical)
        return _b64url(signature)

    def private_key_pem(self) -> bytes:
        return self.private_key.private_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PrivateFormat.PKCS8,
            encryption_algorithm=serialization.NoEncryption(),
        )


# ---------------------------------------------------------------------------
# Client implementation
# ---------------------------------------------------------------------------


class SbnGatewayClient:
    """HTTP client for the SmartBlocks gateway control plane and read APIs."""

    def __init__(
        self,
        *,
        base_url: str,
        read_base_url: str | None = None,
        tenant_id: str | None = None,
        signing_key: SigningKey | None = None,
        token_subject: str | None = None,
        token_scopes: Sequence[str] = ("attest.write",),
        token_cdna: str = "sdk",
        token_ttl_seconds: int = 300,
        token_provider: Callable[[], MintedToken] | None = None,
        retry_config: RetryConfig | None = None,
        timeout: float = 10.0,
        user_agent: str = "sbn-gateway-sdk/0.1",
        transport: httpx.BaseTransport | None = None,
    ) -> None:
        self._base_url = base_url.rstrip("/")
        self._read_base_url = (read_base_url or base_url).rstrip("/")
        self._tenant_id = tenant_id
        self._retry = retry_config or RetryConfig()
        self._client = httpx.Client(
            base_url=self._base_url,
            timeout=timeout,
            headers={"User-Agent": user_agent},
            transport=transport,
        )
        self._token_cache: MintedToken | None = None
        if token_provider is not None:
            self._token_provider = token_provider
        elif signing_key is not None:
            subject = token_subject or "sbn-sdk"
            scopes = token_scopes

            def _provider() -> MintedToken:
                return signing_key.issue_token(
                    subject=subject,
                    scopes=scopes,
                    cdna=token_cdna,
                    ttl_seconds=token_ttl_seconds,
                    tenant_id=tenant_id,
                )

            self._token_provider = _provider
        else:
            self._token_provider = None

    # ------------------------------ context mgmt ---------------------------
    def close(self) -> None:
        self._client.close()

    def __enter__(self) -> "SbnGatewayClient":
        return self

    def __exit__(self, exc_type, exc, tb) -> None:  # pragma: no cover - trivial
        self.close()

    # ------------------------------ auth helpers --------------------------
    def _ensure_token(self) -> str:
        if self._token_provider is None:
            return ""
        now = _now_utc()
        if self._token_cache is None or (self._token_cache.expires_at - now).total_seconds() < 30:
            self._token_cache = self._token_provider()
        return self._token_cache.token

    def _headers(self, extra: Mapping[str, str] | None = None) -> dict[str, str]:
        headers: dict[str, str] = {}
        if self._tenant_id:
            headers["X-Tenant-ID"] = self._tenant_id
        token = self._ensure_token()
        if token:
            headers["Authorization"] = f"Bearer {token}"
        if extra:
            headers.update(extra)
        return headers

    # ------------------------------ request core --------------------------
    def _request(self, method: str, url: str, *, headers: Mapping[str, str] | None = None, **kwargs: Any) -> httpx.Response:
        attempts = max(1, int(self._retry.attempts))
        backoff = max(0.0, float(self._retry.backoff_factor))
        method_upper = method.upper()
        for attempt in range(1, attempts + 1):
            try:
                response = self._client.request(
                    method_upper,
                    url,
                    headers=self._headers(headers),
                    **kwargs,
                )
            except httpx.RequestError as exc:
                if attempt >= attempts:
                    raise GatewayTransportError("Failed to reach gateway", last_exception=exc) from exc
                time.sleep(backoff)
                backoff *= 2
                continue

            if (
                method_upper in self._retry.retry_on_methods
                and response.status_code in self._retry.status_forcelist
                and attempt < attempts
            ):
                response.close()
                time.sleep(backoff)
                backoff *= 2
                continue

            if response.status_code >= 400:
                raise self._error_from_response(response)
            return response
        raise GatewayTransportError("Failed to reach gateway")

    def _error_from_response(self, response: httpx.Response) -> GatewayError:
        try:
            payload = response.json()
        except ValueError:
            payload = None

        message = response.reason_phrase or "Gateway request failed"
        code: str | None = None
        details: MutableMapping[str, Any] = {}

        if isinstance(payload, Mapping):
            if isinstance(payload.get("error"), Mapping):
                err = payload["error"]
                message = str(err.get("message") or message)
                code = err.get("code")
                if isinstance(err.get("details"), Mapping):
                    details.update(err["details"])  # type: ignore[index]
                if err.get("trace_id"):
                    details.setdefault("trace_id", err["trace_id"])
            else:
                message = str(payload.get("message") or message)
                if isinstance(payload.get("details"), Mapping):
                    details.update(payload["details"])  # type: ignore[index]
        elif payload is not None:
            details["response"] = payload

        return GatewayError(message, status_code=response.status_code, code=code, details=details)

    # ------------------------------ public API ----------------------------
    def create_slot(self, request: SlotCreateRequest) -> SlotHandle:
        path = self._slot_collection_path()
        response = self._request("POST", path, json=request.to_payload())
        body = response.json()
        # Write gateway wraps slot responses in {"data": {...}}
        return SlotHandle.from_dict(body.get("data") or body)

    def list_slot_blocks(
        self,
        slot_id: str,
        *,
        limit: int | None = None,
    ) -> dict[str, Any]:
        """List all SmartBlocks belonging to a slot.

        Returns ``{"slot_id": "...", "slot_state": "...", "blocks": [...], "count": N}``.
        """
        params: dict[str, str] = {}
        if limit is not None:
            params["limit"] = str(limit)
        path = f"{self._slot_collection_path()}/{slot_id}/blocks"
        response = self._request("GET", path, params=params)
        body = response.json()
        return body.get("data") or body

    def close_slot(self, slot_id: str) -> SlotClosure:
        path = f"{self._slot_collection_path()}/{slot_id}/close"
        response = self._request("POST", path, json={})
        body = response.json()
        return SlotClosure.from_dict(body.get("data") or body)

    def fetch_receipt(self, receipt_id: str) -> Receipt:
        url = f"{self._read_base_url}/receipts/{receipt_id}"
        response = self._request("GET", url)
        return Receipt.from_dict(response.json())

    def list_flags(self) -> Mapping[str, Any]:
        response = self._request("GET", "/flags")
        data = response.json()
        if isinstance(data, Mapping):
            return data
        return {}

    def request_attestation(self, summary: SlotSummary) -> Mapping[str, Any]:
        response = self._request("POST", "/attest", json=summary.to_payload())
        data = response.json()
        if isinstance(data, Mapping):
            return data
        raise GatewayError("Unexpected attestation response", status_code=response.status_code)

    def _slot_collection_path(self) -> str:
        if self._tenant_id:
            return f"/tenants/{self._tenant_id}/slots"
        return "/slots"


__all__.sort()
