"""Shared Azure Table Storage helpers for service job stores."""

from __future__ import annotations

import logging
import os
import weakref
from collections import OrderedDict
from collections.abc import Callable, Mapping
from dataclasses import dataclass
from threading import Lock
from types import TracebackType
from typing import TYPE_CHECKING, Any, Generic, NamedTuple, Protocol, TypeVar, cast

from .exceptions import AzureTableError

if TYPE_CHECKING:  # pragma: no cover - for type checkers
    from azure.core.exceptions import ResourceNotFoundError
else:  # pragma: no cover - runtime fallback
    try:
        from azure.core.exceptions import ResourceNotFoundError  # type: ignore
    except Exception:

        class ResourceNotFoundError(Exception):
            """Fallback when Azure SDK is unavailable."""


logger = logging.getLogger(__name__)

RecordT = TypeVar("RecordT", bound="AzureTableRecordProtocol")
_ACCOUNT_ENV = "AZURE_STORAGE_ACCOUNT"
_TABLE_ENDPOINT_ENV = "AZURE_STORAGE_TABLE_ENDPOINT"
_TABLE_CLIENT_CACHE_MAXSIZE = 16
_TABLE_CLIENT_CACHE_LOCK = Lock()


def _escape_odata_string(value: str) -> str:
    """Escape an OData string literal by doubling embedded single quotes."""
    return value.replace("'", "''")


class AzureTableRecordProtocol(Protocol):
    """Protocol describing the record interface required by the job store."""

    batch_id: str
    job_id: str
    updated_at: str | None

    def to_entity(self) -> Mapping[str, Any]:  # pragma: no cover - Protocol stub
        raise NotImplementedError

    @classmethod
    def from_entity(cls, entity: Mapping[str, Any]) -> AzureTableRecordProtocol:  # pragma: no cover - Protocol stub
        raise NotImplementedError


@dataclass(frozen=True, init=False)
class TableConfig:
    """Resolved configuration for connecting to an Azure Table."""

    table_name: str
    connection_string: str | None = None
    account_name: str | None = None
    table_endpoint: str | None = None

    # Preserve the legacy positional constructor order while allowing
    # connection-less auth via account name / endpoint.
    def __init__(
        self,
        connection_string: str | None = None,
        table_name: str | None = None,
        *,
        account_name: str | None = None,
        table_endpoint: str | None = None,
    ) -> None:
        if table_name is None:
            raise TypeError("TableConfig.__init__() missing 1 required positional argument: 'table_name'")
        object.__setattr__(self, "connection_string", connection_string)
        object.__setattr__(self, "table_name", table_name)
        object.__setattr__(self, "account_name", account_name)
        object.__setattr__(self, "table_endpoint", table_endpoint)

    @property
    def resolved_table_endpoint(self) -> str | None:
        if self.table_endpoint:
            return self.table_endpoint
        if self.account_name:
            return f"https://{self.account_name}.table.core.windows.net"
        return None

    @classmethod
    def from_env(
        cls,
        *,
        connection_var: str,
        table_var: str,
        fallback_connection_var: str = "AZURE_STORAGE_CONNECTION_STRING",
        env: Mapping[str, str] | None = None,
    ) -> TableConfig:
        source = env if env is not None else os.environ
        table_name = (source.get(table_var) or "").strip()
        if not table_name:
            raise AzureTableError(f"{table_var} environment variable must be set for Azure job tracking")

        connection = (source.get(connection_var) or source.get(fallback_connection_var) or "").strip()
        account_name = (source.get(_ACCOUNT_ENV) or "").strip() or None
        table_endpoint = (source.get(_TABLE_ENDPOINT_ENV) or "").strip() or None
        if not table_endpoint and account_name:
            table_endpoint = f"https://{account_name}.table.core.windows.net"

        if not connection and not (account_name or table_endpoint):
            raise AzureTableError(
                f"Azure job tracking requires {connection_var} (or {fallback_connection_var}) "
                f"or {_ACCOUNT_ENV}/{_TABLE_ENDPOINT_ENV} for DefaultAzureCredential"
            )
        return cls(
            connection_string=connection or None,
            table_name=table_name,
            account_name=account_name,
            table_endpoint=table_endpoint,
        )


@dataclass
class _CachedTableClient:
    client: Any
    service: Any
    credential: Any | None = None
    lease_count: int = 0
    in_cache: bool = True
    closed: bool = False

    def close(self) -> None:
        if self.closed:
            return
        self.closed = True

        seen: set[int] = set()
        for resource in (self.client, self.service, self.credential):
            if resource is None:
                continue
            resource_id = id(resource)
            if resource_id in seen:
                continue
            seen.add(resource_id)
            close = getattr(resource, "close", None)
            if callable(close):
                close()


class _TableClientCacheInfo(NamedTuple):
    hits: int
    misses: int
    maxsize: int
    currsize: int


_TABLE_CLIENT_CACHE: OrderedDict[TableConfig, _CachedTableClient] = OrderedDict()
_table_client_cache_hits = 0
_table_client_cache_misses = 0


def _table_service_client_class() -> Any:
    try:
        from azure.data.tables import TableServiceClient
    except ImportError as exc:  # pragma: no cover - optional dependency missing
        raise AzureTableError(
            "azure-data-tables package is required for Azure job tracking. "
            "Install it with `pip install azure-data-tables`."
        ) from exc
    return TableServiceClient


def _create_default_credential() -> Any:
    try:
        from azure.identity import DefaultAzureCredential
    except ImportError as exc:  # pragma: no cover - optional dependency missing
        raise AzureTableError(
            "azure-identity package is required for Azure job tracking with DefaultAzureCredential. "
            "Install it with `pip install azure-identity`."
        ) from exc
    return DefaultAzureCredential()


def _create_table_service_client(config: TableConfig) -> tuple[Any, Any | None]:
    table_service_client = _table_service_client_class()
    if config.connection_string:
        logger.debug(
            "Creating TableServiceClient from connection string",
            extra={"table_name": config.table_name},
        )
        return table_service_client.from_connection_string(config.connection_string), None

    endpoint = config.resolved_table_endpoint
    if not endpoint:
        raise AzureTableError(
            "Azure job tracking requires a connection string or storage account/endpoint configuration"
        )

    credential = _create_default_credential()
    logger.debug(
        "Creating TableServiceClient with DefaultAzureCredential",
        extra={"table_name": config.table_name, "table_endpoint": endpoint},
    )
    return table_service_client(endpoint=endpoint, credential=credential), credential


def _build_cached_table_client(config: TableConfig) -> _CachedTableClient:
    service, credential = _create_table_service_client(config)
    table_client = service.create_table_if_not_exists(config.table_name)
    if table_client is None:
        table_client = service.get_table_client(config.table_name)
    return _CachedTableClient(client=table_client, service=service, credential=credential)


def _get_cached_table_client(config: TableConfig, *, acquire: bool) -> _CachedTableClient:
    global _table_client_cache_hits, _table_client_cache_misses

    to_close: list[_CachedTableClient] = []

    with _TABLE_CLIENT_CACHE_LOCK:
        cached = _TABLE_CLIENT_CACHE.get(config)
        if cached is not None:
            _table_client_cache_hits += 1
            _TABLE_CLIENT_CACHE.move_to_end(config)
            if acquire:
                cached.lease_count += 1
            return cached

    created = _build_cached_table_client(config)
    if acquire:
        created.lease_count = 1

    with _TABLE_CLIENT_CACHE_LOCK:
        cached = _TABLE_CLIENT_CACHE.get(config)
        if cached is not None:
            _table_client_cache_hits += 1
            _TABLE_CLIENT_CACHE.move_to_end(config)
            if acquire:
                cached.lease_count += 1
            to_close.append(created)
            resource = cached
        else:
            _table_client_cache_misses += 1
            _TABLE_CLIENT_CACHE[config] = created
            resource = created
            if len(_TABLE_CLIENT_CACHE) > _TABLE_CLIENT_CACHE_MAXSIZE:
                _, evicted = _TABLE_CLIENT_CACHE.popitem(last=False)
                evicted.in_cache = False
                if evicted.lease_count == 0:
                    to_close.append(evicted)

    for cached_resource in to_close:
        cached_resource.close()

    return resource


def _release_cached_table_client(config: TableConfig, resource: _CachedTableClient) -> None:
    should_close = False
    with _TABLE_CLIENT_CACHE_LOCK:
        if resource.lease_count > 0:
            resource.lease_count -= 1
        if resource.lease_count == 0:
            cached = _TABLE_CLIENT_CACHE.get(config)
            if cached is resource:
                # Lease-based entries are intentionally short-lived once the last
                # store releases them. The cache deduplicates concurrent live
                # stores, while zero-lease resources are evicted so credentials
                # and transports close deterministically instead of lingering.
                _TABLE_CLIENT_CACHE.pop(config, None)
                resource.in_cache = False
            should_close = not resource.in_cache

    if should_close:
        resource.close()


def _table_client_cache_info() -> _TableClientCacheInfo:
    with _TABLE_CLIENT_CACHE_LOCK:
        return _TableClientCacheInfo(
            hits=_table_client_cache_hits,
            misses=_table_client_cache_misses,
            maxsize=_TABLE_CLIENT_CACHE_MAXSIZE,
            currsize=len(_TABLE_CLIENT_CACHE),
        )


def _clear_table_client_cache() -> None:
    global _table_client_cache_hits, _table_client_cache_misses

    with _TABLE_CLIENT_CACHE_LOCK:
        resources = list(_TABLE_CLIENT_CACHE.values())
        _TABLE_CLIENT_CACHE.clear()
        _table_client_cache_hits = 0
        _table_client_cache_misses = 0
        for resource in resources:
            resource.in_cache = False

    for resource in resources:
        if resource.lease_count == 0:
            resource.close()


def _acquire_table_client(config: TableConfig) -> tuple[Any, Callable[[], None]]:
    resource = _get_cached_table_client(config, acquire=True)
    return resource.client, lambda: _release_cached_table_client(config, resource)


class _TableClientFactory:
    def __call__(self, config: TableConfig) -> Any:
        return _get_cached_table_client(config, acquire=False).client

    def cache_info(self) -> _TableClientCacheInfo:
        return _table_client_cache_info()

    def cache_clear(self) -> None:
        _clear_table_client_cache()


_table_client = _TableClientFactory()


def reset_table_client_cache() -> None:
    """Clear the cached table clients.

    Use this to force reconnection after configuration changes or in tests
    to ensure isolation between test cases. Existing AzureTableJobStore
    instances keep their leased client until they are closed or garbage
    collected.
    """
    _clear_table_client_cache()


class AzureTableJobStore(Generic[RecordT]):
    """Use this job store for per-run metadata persisted in Azure Table Storage."""

    def __init__(
        self,
        *,
        record_type: type[RecordT],
        connection_string_var: str,
        table_name_var: str,
        fallback_connection_var: str = "AZURE_STORAGE_CONNECTION_STRING",
        config: TableConfig | None = None,
        env: Mapping[str, str] | None = None,
    ) -> None:
        self._record_type = record_type
        self._config = config or TableConfig.from_env(
            connection_var=connection_string_var,
            table_var=table_name_var,
            fallback_connection_var=fallback_connection_var,
            env=env,
        )
        self._client, release_client = _acquire_table_client(self._config)
        self._client_finalizer = weakref.finalize(self, release_client)

    # ------------------------------------------------------------------
    # CRUD helpers
    # ------------------------------------------------------------------

    def create(self, record: RecordT) -> RecordT:
        self._client.upsert_entity(record.to_entity())
        return record

    def save(self, record: RecordT) -> RecordT:
        self._client.upsert_entity(record.to_entity())
        return record

    def close(self) -> None:
        self._client_finalizer()

    def fetch(self, batch_id: str, job_id: str) -> RecordT | None:
        try:
            entity = self._client.get_entity(batch_id, job_id)
        except (ResourceNotFoundError, KeyError):  # pragma: no cover - SDK raises on missing entities
            return None
        except Exception:
            logger.exception("Failed to fetch table entity", extra={"batch_id": batch_id, "job_id": job_id})
            raise
        return cast(RecordT, self._record_type.from_entity(entity))

    # ------------------------------------------------------------------
    # Listing helpers
    # ------------------------------------------------------------------

    def list_runs(
        self,
        batch_id: str,
        *,
        limit: int = 50,
        continuation_token: str | None = None,
    ) -> tuple[list[RecordT], str | None]:
        """Return one filtered page of records for ``batch_id``.

        Records are sorted by ``updated_at`` within the returned page only.
        Across continuation tokens, the overall ordering remains whatever the
        underlying Azure Table query pagination produced.
        """
        filter_expr = f"PartitionKey eq '{_escape_odata_string(batch_id)}'"
        query_fn = getattr(self._client, "query_entities", None)

        if callable(query_fn):
            pager = query_fn(query_filter=filter_expr, results_per_page=limit).by_page(
                continuation_token=continuation_token
            )

            try:
                page = next(pager)
            except StopIteration:  # pragma: no cover - empty dataset
                return [], None

            collected = [cast(RecordT, self._record_type.from_entity(entity)) for entity in page]
            next_token = getattr(pager, "continuation_token", None)

            collected.sort(key=lambda rec: rec.updated_at or "", reverse=True)
            return collected, next_token

        raise AzureTableError(
            "query_entities is unavailable on the configured Table client; upgrade azure-data-tables "
            "or provide a client with query_entities to avoid unfiltered scans."
        )

    # ------------------------------------------------------------------
    # Accessors
    # ------------------------------------------------------------------

    @property
    def client(self) -> Any:
        """Expose the underlying TableClient for advanced scenarios (e.g., aggregation)."""

        return self._client

    def __enter__(self) -> AzureTableJobStore[RecordT]:
        return self

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: TracebackType | None,
    ) -> None:
        self.close()


__all__ = ["AzureTableError", "AzureTableJobStore", "TableConfig", "reset_table_client_cache"]
