import asyncio
import json
import signal
import time
from typing import List, Optional

from grpc.aio import Channel, insecure_channel
from pydantic import Field
from pydantic_settings import BaseSettings, SettingsConfigDict


class GRPCSettings(BaseSettings):
    """Configs are done by env vars because the package is mainly singleton."""

    # Server settings
    HOST: str = "127.0.0.1"
    PORT: int = 50051
    API_KEY: str = ""

    # Shutdown settings
    CHANNEL_SHUTDOWN_GRACE: float = 10.0

    # Keepalive settings
    KEEPALIVE_TIME_MS: int = 60000  # ping every 1 min
    KEEPALIVE_TIMEOUT_MS: int = 5000  # wait 5s for ack
    KEEPALIVE_PERMIT_WITHOUT_CALLS: int = 1  # even without active RPCs
    MAX_PINGS_WITHOUT_DATA: int = 0  # unlimited pings without data

    # Retry policy settings
    # We prefer aiming for only server glitches; restarts are handled by docker swarm zero-downtime
    MAX_ATTEMPTS: int = 5  # 1 original + 4 retries
    INITIAL_BACKOFF: str = "0.1s"  # quick first retry
    MAX_BACKOFF: str = "1.5s"  # max wait 1.5 second
    BACKOFF_MULTIPLIER: float = 2.0  # double wait time each retry
    RETRYABLE_STATUS_CODES: List[str] = Field(
        default=[
            "UNAVAILABLE",  # server temporarily unreachable
            "RESOURCE_EXHAUSTED",  # too many requests (try later)
            "ABORTED",  # conflict or retryable error
        ]
    )

    # RPC level configs
    WAIT_FOR_READY: bool = True
    TIMEOUT: Optional[float] = 5.0  # seconds

    model_config = SettingsConfigDict(env_prefix="ASYNC_ACCOUNT_SDK_")


grpc_settings = GRPCSettings()

SERVICE_NAME = "ADMAC"
TARGET = f"{grpc_settings.HOST}:{grpc_settings.PORT}"
CHANNEL: Channel = None

KEEPALIVE_OPTS = [
    ("grpc.keepalive_time_ms", grpc_settings.KEEPALIVE_TIME_MS),
    ("grpc.keepalive_timeout_ms", grpc_settings.KEEPALIVE_TIMEOUT_MS),
    (
        "grpc.keepalive_permit_without_calls",
        grpc_settings.KEEPALIVE_PERMIT_WITHOUT_CALLS,
    ),
    ("grpc.http2.max_pings_without_data", grpc_settings.MAX_PINGS_WITHOUT_DATA),
]

RETRY_CONFIG = {
    "methodConfig": [
        {
            "name": [{}],  # All methods
            "retryPolicy": {
                "maxAttempts": grpc_settings.MAX_ATTEMPTS,
                "initialBackoff": grpc_settings.INITIAL_BACKOFF,
                "maxBackoff": grpc_settings.MAX_BACKOFF,
                "backoffMultiplier": grpc_settings.BACKOFF_MULTIPLIER,
                "retryableStatusCodes": grpc_settings.RETRYABLE_STATUS_CODES,
            },
        }
    ]
}


async def construct_channel():
    global CHANNEL
    CHANNEL = insecure_channel(
        TARGET,
        options=KEEPALIVE_OPTS + [("grpc.service_config", json.dumps(RETRY_CONFIG))],
    )

    async def graceful_shutdown():
        print(
            f"SIGTERM received, shutting down the {SERVICE_NAME} channel gracefully..."
        )
        start_time = time.time()

        try:
            await CHANNEL.close(grace=grpc_settings.CHANNEL_SHUTDOWN_GRACE)
        except Exception as e:
            print(f"Error during shutdown: {e}")

        elapsed_time = time.time() - start_time
        print(
            f"{SERVICE_NAME} channel gracefully shut down in {elapsed_time:.2f} seconds."
        )

    # We want the currently sent requests to get completed before termination
    loop = asyncio.get_running_loop()
    loop.add_signal_handler(
        signal.SIGTERM, lambda: asyncio.create_task(graceful_shutdown())
    )
    print(f"Graceful shutdown on sigterm singal registered for {SERVICE_NAME} channel")
