# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

"""Process wiring: build the store, bind the servicers, serve."""

from __future__ import annotations

import asyncio
import logging
import signal
from dataclasses import dataclass

import grpc

from .auth import BearerAuthInterceptor
from .catalog.loader import CatalogLoader
from .health import add_health_service, enter_graceful_shutdown
from .hf.client import HFClient
from .registration import planned_bindings, register
from .services.capabilities import CapabilitiesService
from .services.filefolder import FileFolderService
from .services.fileobject import FileObjectService
from .services.metadata_service import MetadataService
from .settings import Settings
from .store import AssetStore

logger = logging.getLogger(__name__)


@dataclass
class Service:
    """A wired-up, not-yet-serving service. Split out so tests can drive it over
    an in-process channel without a port."""

    settings: Settings
    hf: HFClient
    store: AssetStore
    loader: CatalogLoader
    server: grpc.aio.Server
    health: object

    async def aclose(self) -> None:
        await self.loader.stop()
        await self.hf.aclose()


async def build(settings: Settings, *, hf: HFClient | None = None) -> Service:
    """Load the catalog and wire everything up.

    The catalog is the only thing loaded here. No HuggingFace asset data is
    fetched — folders load on first request (plan §5.2) — so startup does not
    depend on the Hub being reachable.
    """
    hf = hf or HFClient(
        endpoint=settings.hf_endpoint,
        token=settings.resolve_hf_token(),
        timeout=settings.http_timeout,
        max_connections=settings.http_max_connections,
        max_retries=settings.http_max_retries,
        backoff_base=settings.http_backoff_base,
    )

    loader = CatalogLoader(settings, hf)
    catalog = await loader.load()
    store = AssetStore(catalog, hf, settings)

    interceptors = []
    if not settings.allow_anonymous:
        interceptors.append(BearerAuthInterceptor(settings.resolve_auth_token()))
    else:
        logger.warning("inbound auth is DISABLED (HF_ADAPTER_ALLOW_ANONYMOUS=true)")

    server = grpc.aio.server(
        interceptors=interceptors,
        maximum_concurrent_rpcs=settings.max_concurrent_rpcs,
    )

    # One plan drives both what gets bound and what ListServices reports.
    bindings = planned_bindings(
        advertise_filefolder=settings.advertise_filefolder,
        advertise_v1beta=settings.advertise_v1beta,
    )
    capabilities = CapabilitiesService(store, settings, bindings)
    register(
        server,
        bindings=bindings,
        capabilities=capabilities,
        fileobject=FileObjectService(store, settings),
        metadata=MetadataService(store, settings),
        filefolder=FileFolderService(store, settings) if settings.advertise_filefolder else None,
    )
    logger.info(
        "advertising: %s",
        ", ".join(f"{name} ({'+'.join(v)})" for name, v in bindings.items()),
    )

    # Registered last, and only once everything above succeeded: the catalog is
    # loaded and every servicer is bound, which is exactly what SERVING claims.
    health = await add_health_service(server, bindings)

    await loader.start_refresh(store.swap_catalog)
    return Service(settings=settings, hf=hf, store=store, loader=loader, server=server, health=health)


async def serve(settings: Settings | None = None) -> None:
    settings = settings or Settings()
    logging.basicConfig(
        level=settings.log_level.upper(),
        format="%(asctime)s %(levelname)s %(name)s %(message)s",
    )

    service = await build(settings)
    address = f"{settings.grpc_host}:{settings.grpc_port}"
    service.server.add_insecure_port(address)
    await service.server.start()
    logger.info("listening on %s", address)

    stop = asyncio.Event()
    loop = asyncio.get_running_loop()
    for sig in (signal.SIGINT, signal.SIGTERM):
        loop.add_signal_handler(sig, stop.set)

    try:
        await stop.wait()
    finally:
        # NOT_SERVING first, so Kubernetes pulls the pod from the Service
        # endpoints before we start refusing connections.
        logger.info("draining...")
        await enter_graceful_shutdown(service.health)
        await service.server.stop(grace=10)
        await service.aclose()