"""Async Omni client."""

from __future__ import annotations

from contextlib import asynccontextmanager
from pathlib import Path
from typing import Any

from omni.api import OIDCAPI, AuthAPI, ManagementAPI, ResourceAPI
from omni.auth import SideroSigner, decode_service_account, load_service_account_from_env
from omni.config import CacheConfig, ConfigInput, OmniSDKConfig, RetryConfig, load_config
from omni.http import RpcMethod
from omni.http.transport import OmniTransport
from omni.state import StateStore, default_state_path


class AsyncOmniClient:
    """Async-first Omni API client."""

    def __init__(
        self,
        config: ConfigInput | None = None,
        *,
        config_path: str | Path | None = None,
        instance: str | None = None,
        cluster: str | None = None,
        state_path: str | Path | None = None,
        timeout: float = 30.0,
        verify_tls: bool | str = True,
        retries: RetryConfig | None = None,
        cache: CacheConfig | None = None,
    ) -> None:
        resolved = load_config(config, config_path=config_path)
        self._resolved_config = resolved
        self.config: OmniSDKConfig = resolved.data

        instance_name, instance_cfg = self.config.pick_instance(instance)
        self.instance_name = instance_name
        self.instance_config = instance_cfg

        effective_state = state_path
        if effective_state is None:
            effective_state = default_state_path(resolved.path)
        self._state = StateStore(effective_state)

        self.cluster = (
            cluster
            or self._state.get_preference("default_cluster")
            or self.config.preferences.instance_clusters.get(instance_name)
            or self.config.preferences.default_cluster
        )

        account = None
        if instance_cfg.auth.service_account_key:
            account = decode_service_account(instance_cfg.auth.service_account_key)
        else:
            account = load_service_account_from_env(instance_cfg.auth.service_account_env_aliases)

        signer = SideroSigner(account) if account is not None else None

        effective_retry = retries or instance_cfg.retry
        effective_cache = cache or instance_cfg.cache
        effective_timeout = timeout if timeout != 30.0 else instance_cfg.timeout
        effective_verify = verify_tls if verify_tls is not True else instance_cfg.verify_tls

        self._transport = OmniTransport(
            instance_name=instance_name,
            instance=instance_cfg,
            state=self._state,
            signer=signer,
            timeout=effective_timeout,
            verify_tls=effective_verify,
            retry_config=effective_retry,
            cache_config=effective_cache,
        )

        self.resources = ResourceAPI(self._transport)
        self.management = ManagementAPI(self._transport)
        self.oidc = OIDCAPI(self._transport)
        self.auth = AuthAPI(self._transport)

    @property
    def raw(self) -> AsyncOmniClient:
        return self

    @property
    def state(self) -> StateStore:
        return self._state

    async def call(
        self,
        method: RpcMethod,
        request: dict[str, Any],
        *,
        metadata: dict[str, str] | None = None,
    ) -> dict[str, Any]:
        return await self._transport.call(method, request, metadata=metadata)

    async def stream(
        self,
        method: RpcMethod,
        request: dict[str, Any],
        *,
        metadata: dict[str, str] | None = None,
    ):
        async for event in self._transport.stream(method, request, metadata=metadata):
            yield event

    async def ensure_talosconfig(
        self,
        *,
        cluster: str | None = None,
        break_glass: bool = False,
        force_refresh: bool = False,
        ttl_seconds: int = 300,
        path: str | Path | None = None,
    ) -> Path:
        from omni.cli.talos_manager import TalosConfigManager

        manager = TalosConfigManager(self)
        resolved = await manager.ensure(
            cluster=cluster,
            break_glass=break_glass,
            force_refresh=force_refresh,
            ttl_seconds=ttl_seconds,
            path=path,
        )
        return resolved.path

    async def aclose(self) -> None:
        await self._transport.aclose()
        self._state.close()

    async def __aenter__(self) -> AsyncOmniClient:
        return self

    async def __aexit__(self, exc_type, exc, tb) -> None:
        await self.aclose()


@asynccontextmanager
async def connect(*args: Any, **kwargs: Any):
    client = AsyncOmniClient(*args, **kwargs)
    try:
        yield client
    finally:
        await client.aclose()
