"""Health service handles health check operations."""

from __future__ import annotations

from typing import TYPE_CHECKING, Callable, Generator, Optional

import grpc

if TYPE_CHECKING:
    from ..types import HealthStatus
    from .service_context import ServiceContext


class HealthService:
    """Health service for checking service health."""

    def __init__(self, ctx: ServiceContext):
        self.ctx = ctx
        self._stub = None

    def _get_stub(self):
        if self._stub is None:
            from ..proto import gqldb_pb2_grpc
            channel = self.ctx.get_channel()
            self._stub = gqldb_pb2_grpc.HealthStub(channel)
        return self._stub

    def health_check(self, service: str) -> HealthStatus:
        from ..errors import GqldbError
        from ..proto import gqldb_pb2
        from .converters import convert_health_status

        try:
            request = gqldb_pb2.HealthCheckRequest(service=service)
            metadata = self.ctx.get_session_metadata()
            response = self._get_stub().Check(
                request, metadata=metadata, timeout=self.ctx.config.timeout
            )
            return convert_health_status(response.status)
        except grpc.RpcError as e:
            raise GqldbError(f"Health check failed: {e.details()}") from e

    def watch(
        self,
        service: str,
        callback: Optional[Callable[[HealthStatus], None]]
    ) -> None:
        from ..errors import GqldbError
        from ..proto import gqldb_pb2
        from .converters import convert_health_status

        try:
            request = gqldb_pb2.HealthCheckRequest(service=service)
            metadata = self.ctx.get_session_metadata()
            stream = self._get_stub().Watch(request, metadata=metadata)

            for response in stream:
                status = convert_health_status(response.status)
                if callback:
                    callback(status)
        except grpc.RpcError as e:
            raise GqldbError(f"Health watch failed: {e.details()}") from e

    def watch_iter(self, service: str) -> Generator[HealthStatus, None, None]:
        from ..errors import GqldbError
        from ..proto import gqldb_pb2
        from .converters import convert_health_status

        try:
            request = gqldb_pb2.HealthCheckRequest(service=service)
            metadata = self.ctx.get_session_metadata()
            stream = self._get_stub().Watch(request, metadata=metadata)

            for response in stream:
                yield convert_health_status(response.status)
        except grpc.RpcError as e:
            raise GqldbError(f"Health watch failed: {e.details()}") from e
