import time
from collections.abc import Iterator
from typing import Any, Literal

from cognite.client import ClientConfig, CogniteClient
from cognite.client.data_classes.data_modeling import Edge, Node, ViewId
from cognite.client.data_classes.filters import Filter
from cognite.client.exceptions import CogniteAPIError
from cognite.client.utils.useful_types import SequenceNotStr


class NeatInstancesAPI:
    """Custom DMS instances API that handles 408 timeouts by reducing page size."""

    def __init__(self, client: "NeatClient") -> None:
        self._client = client

    def iterate(
        self,
        instance_type: Literal["node", "edge"] = "node",
        space: str | SequenceNotStr[str] | None = None,
        filter_: Filter | None = None,
        source: ViewId | None = None,
    ) -> Iterator[Node] | Iterator[Edge]:
        """Iterate instances with 408 timeout handling and optimized sort.

        Uses sort by externalId to leverage the (project_id, space, external_id) index
        and avoid soft-deleted instances.
        """
        body: dict[str, Any] = {"limit": 1_000, "cursor": None, "instanceType": instance_type}
        body["sort"] = [
            {
                "property": [instance_type, "externalId"],
                "direction": "ascending",
            }
        ]
        url = f"/api/{self._client._API_VERSION}/projects/{self._client.config.project}/models/instances/list"
        filter_ = self._client.data_modeling.instances._merge_space_into_filter(instance_type, space, filter_)
        if filter_:
            body["filter"] = filter_.dump() if isinstance(filter_, Filter) else filter_
        if source:
            body["sources"] = [{"source": source.dump(include_type=True, camel_case=True)}]
        last_limit_change: float = 0.0
        while True:
            try:
                response = self._client.post(url=url, json=body)
            except CogniteAPIError as e:
                if e.code == 408 and body["limit"] > 1:
                    body["limit"] = body["limit"] // 2
                    last_limit_change = time.perf_counter()
                    continue
                raise e
            response_body = response.json()
            yield from (
                Node.load(item) if item.get("instanceType") == "node" else Edge.load(item)
                for item in response_body["items"]
            )
            next_cursor = response_body.get("nextCursor")
            if next_cursor is None:
                break
            body["cursor"] = next_cursor
            # Increase the limit every 30 seconds to recover from 408s
            if (time.perf_counter() - last_limit_change) > 30.0 and body["limit"] < 1_000:
                body["limit"] = min(int(body["limit"] * 1.5), 1_000)
                last_limit_change = time.perf_counter()


class NeatClient(CogniteClient):
    def __init__(self, config: ClientConfig | CogniteClient | None = None) -> None:
        if isinstance(config, CogniteClient):
            config = config.config
        super().__init__(config=config)
        self.instances = NeatInstancesAPI(self)
