"""
Lightweight Pub/Sub helper for g-gremlin.

Uses the Pub/Sub REST API with BaseIntegration retries/backoff and
Google auth for access tokens (ADC or service account JSON).
"""

from __future__ import annotations

import base64
import json
import os
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, Iterable, List, Optional

from g_gremlin.base import AuthenticationError, BaseIntegration, IntegrationError, RateLimitError, RetryConfig
from g_gremlin.credentials import resolve_credentials
from g_gremlin.credentials.provider import CredentialProvider, SecretNotFoundError

PUBSUB_BASE_URL = "https://pubsub.googleapis.com/v1"
PUBSUB_SCOPE = "https://www.googleapis.com/auth/pubsub"
DEFAULT_RETRY_CONFIG = RetryConfig(max_attempts=4, base_delay=1.0, max_delay=12.0)


class PubSubError(RuntimeError):
    """Raised when Pub/Sub operations fail."""


class PubSubAuthError(PubSubError):
    """Raised when Pub/Sub authentication fails."""


@dataclass
class PubSubConfig:
    project_id: str
    credentials_path: Optional[str] = None


class PubSubClient(BaseIntegration):
    """Pub/Sub REST client with retry/backoff and auth helpers."""

    SERVICE_NAME = "pubsub"

    def __init__(
        self,
        project_id: Optional[str] = None,
        credentials_path: Optional[str] = None,
        credential_provider: CredentialProvider | None = None,
        retry_config: RetryConfig | None = None,
        access_token: Optional[str] = None,
    ) -> None:
        super().__init__(credential_provider=credential_provider, default_retry_config=retry_config or DEFAULT_RETRY_CONFIG)
        resolved_project, resolved_creds = self._resolve_config(project_id, credentials_path)
        self.credentials_path = resolved_creds
        self._access_token = access_token
        self._credentials = None
        adc_project = None
        if not access_token:
            self._credentials, adc_project = self._load_credentials(resolved_creds)
        self.project_id = resolved_project or adc_project
        if not self.project_id:
            raise PubSubAuthError(
                "Pub/Sub project_id is required. Set G_GREMLIN_PUBSUB_PROJECT_ID, "
                "GOOGLE_CLOUD_PROJECT, or pass project_id explicitly."
            )

    def _resolve_config(
        self,
        project_id: Optional[str],
        credentials_path: Optional[str],
    ) -> tuple[Optional[str], Optional[str]]:
        resolved_project = None
        resolved_creds = None
        try:
            creds = resolve_credentials(
                "pubsub",
                required_keys=[],
                optional_keys=["project_id", "credentials_path"],
                credential_provider=self.credential_provider,
            )
            resolved_project = creds.get("project_id")
            resolved_creds = creds.get("credentials_path")
        except SecretNotFoundError:
            pass

        env_project = (
            os.environ.get("G_GREMLIN_PUBSUB_PROJECT_ID")
            or os.environ.get("GOOGLE_CLOUD_PROJECT")
            or os.environ.get("GCP_PROJECT")
        )
        env_creds = (
            os.environ.get("G_GREMLIN_PUBSUB_CREDENTIALS_PATH")
            or os.environ.get("G_GREMLIN_PUBSUB_CREDENTIALS")
            or os.environ.get("GOOGLE_APPLICATION_CREDENTIALS")
        )

        project = project_id or resolved_project or env_project
        creds_path = credentials_path or resolved_creds or env_creds
        if not project and creds_path:
            project = self._project_from_service_account(creds_path)
        return project, creds_path

    def _project_from_service_account(self, credentials_path: str) -> Optional[str]:
        try:
            data = json.loads(Path(credentials_path).read_text(encoding="utf-8"))
            return data.get("project_id")
        except Exception:
            return None

    def _load_credentials(self, credentials_path: Optional[str]):
        try:
            if credentials_path:
                from google.oauth2 import service_account

                info = json.loads(Path(credentials_path).read_text(encoding="utf-8"))
                credentials = service_account.Credentials.from_service_account_info(
                    info,
                    scopes=[PUBSUB_SCOPE],
                )
                return credentials, info.get("project_id")

            import google.auth

            credentials, project_id = google.auth.default(scopes=[PUBSUB_SCOPE])
            return credentials, project_id
        except Exception as exc:
            raise PubSubAuthError(f"Failed to load Pub/Sub credentials: {exc}") from exc

    def _get_access_token(self) -> str:
        if self._access_token:
            return self._access_token

        if not self._credentials:
            raise PubSubAuthError("Pub/Sub credentials not configured.")

        try:
            from google.auth.transport.requests import Request

            if not self._credentials.valid:
                self._credentials.refresh(Request())
            token = self._credentials.token
        except Exception as exc:
            raise PubSubAuthError(f"Failed to refresh Pub/Sub token: {exc}") from exc

        if not token:
            raise PubSubAuthError("Failed to acquire Pub/Sub access token.")
        return token

    def _headers(self) -> Dict[str, str]:
        return {
            "Authorization": f"Bearer {self._get_access_token()}",
            "Content-Type": "application/json",
        }

    def _topic_path(self, topic: str) -> str:
        if topic.startswith("projects/"):
            return topic
        return f"projects/{self.project_id}/topics/{topic}"

    def _subscription_path(self, subscription: str) -> str:
        if subscription.startswith("projects/"):
            return subscription
        return f"projects/{self.project_id}/subscriptions/{subscription}"

    def _encode_message(
        self,
        message: Any,
        attributes: Optional[Dict[str, str]],
        ordering_key: Optional[str],
    ) -> Dict[str, Any]:
        if isinstance(message, bytes):
            raw = message
        elif isinstance(message, str):
            raw = message.encode("utf-8")
        else:
            raw = json.dumps(message, ensure_ascii=True, separators=(",", ":")).encode("utf-8")

        payload: Dict[str, Any] = {"data": base64.b64encode(raw).decode("ascii")}
        if attributes:
            payload["attributes"] = attributes
        if ordering_key:
            payload["orderingKey"] = ordering_key
        return payload

    def publish(
        self,
        topic: str,
        messages: Iterable[Any],
        attributes: Optional[Dict[str, str]] = None,
        ordering_key: Optional[str] = None,
    ) -> List[str]:
        topic_path = self._topic_path(topic)
        url = f"{PUBSUB_BASE_URL}/{topic_path}:publish"
        payload = {
            "messages": [
                self._encode_message(msg, attributes=attributes, ordering_key=ordering_key)
                for msg in messages
            ]
        }
        try:
            resp = self._post(url, headers=self._headers(), json=payload)
        except AuthenticationError as exc:
            raise PubSubAuthError(str(exc)) from exc
        except (RateLimitError, IntegrationError) as exc:
            raise PubSubError(str(exc)) from exc

        data = resp.json() if resp.content else {}
        return data.get("messageIds", []) if isinstance(data, dict) else []

    def pull(
        self,
        subscription: str,
        max_messages: int = 10,
        return_immediately: bool = True,
        decode: bool = True,
    ) -> List[Dict[str, Any]]:
        subscription_path = self._subscription_path(subscription)
        url = f"{PUBSUB_BASE_URL}/{subscription_path}:pull"
        payload = {
            "maxMessages": max_messages,
            "returnImmediately": return_immediately,
        }
        try:
            resp = self._post(url, headers=self._headers(), json=payload)
        except AuthenticationError as exc:
            raise PubSubAuthError(str(exc)) from exc
        except (RateLimitError, IntegrationError) as exc:
            raise PubSubError(str(exc)) from exc

        data = resp.json() if resp.content else {}
        received = data.get("receivedMessages", []) if isinstance(data, dict) else []
        results: List[Dict[str, Any]] = []
        for item in received:
            message = item.get("message", {}) if isinstance(item, dict) else {}
            data_b64 = message.get("data", "")
            decoded_bytes = base64.b64decode(data_b64) if decode and data_b64 else b""
            decoded_text = None
            decoded_json = None
            if decode and decoded_bytes:
                decoded_text = decoded_bytes.decode("utf-8", errors="replace")
                try:
                    decoded_json = json.loads(decoded_text)
                except Exception:
                    decoded_json = None

            results.append(
                {
                    "ack_id": item.get("ackId"),
                    "message_id": message.get("messageId"),
                    "publish_time": message.get("publishTime"),
                    "attributes": message.get("attributes") or {},
                    "data": decoded_text if decode else data_b64,
                    "data_bytes": decoded_bytes if decode else b"",
                    "data_json": decoded_json,
                    "data_raw": data_b64,
                }
            )
        return results

    def ack(self, subscription: str, ack_ids: Iterable[str]) -> None:
        subscription_path = self._subscription_path(subscription)
        url = f"{PUBSUB_BASE_URL}/{subscription_path}:acknowledge"
        payload = {"ackIds": list(ack_ids)}
        try:
            self._post(url, headers=self._headers(), json=payload)
        except AuthenticationError as exc:
            raise PubSubAuthError(str(exc)) from exc
        except (RateLimitError, IntegrationError) as exc:
            raise PubSubError(str(exc)) from exc

    def pull_and_ack(
        self,
        subscription: str,
        max_messages: int = 10,
        return_immediately: bool = True,
        decode: bool = True,
    ) -> List[Dict[str, Any]]:
        messages = self.pull(
            subscription,
            max_messages=max_messages,
            return_immediately=return_immediately,
            decode=decode,
        )
        ack_ids = [m.get("ack_id") for m in messages if m.get("ack_id")]
        if ack_ids:
            self.ack(subscription, ack_ids)
        return messages

    def ensure_topic(self, topic: str) -> None:
        topic_path = self._topic_path(topic)
        url = f"{PUBSUB_BASE_URL}/{topic_path}"
        try:
            self._put(url, headers=self._headers(), allow_statuses=(409,))
        except AuthenticationError as exc:
            raise PubSubAuthError(str(exc)) from exc
        except (RateLimitError, IntegrationError) as exc:
            raise PubSubError(str(exc)) from exc

    def ensure_subscription(self, subscription: str, topic: str) -> None:
        subscription_path = self._subscription_path(subscription)
        url = f"{PUBSUB_BASE_URL}/{subscription_path}"
        payload = {"topic": self._topic_path(topic)}
        try:
            self._put(url, headers=self._headers(), json=payload, allow_statuses=(409,))
        except AuthenticationError as exc:
            raise PubSubAuthError(str(exc)) from exc
        except (RateLimitError, IntegrationError) as exc:
            raise PubSubError(str(exc)) from exc


__all__ = [
    "PubSubClient",
    "PubSubConfig",
    "PubSubError",
    "PubSubAuthError",
]
