"""Slices Client."""

import importlib.metadata
from pathlib import Path

import httpx
import msgspec
from httpx import HTTPError
from slices_clientlib_core.base.config import SlicesCentralApiClientsConfig
from slices_clientlib_core.base.exceptions import from_httpx_exception
from slices_clientlib_core.client import (
    SlicesClient as SlicesCoreClient,
)
from slices_clientlib_core.client import (
    SlicesExperimentClient as SlicesCoreExperimentClient,
)
from slices_clientlib_core.client import (
    SlicesProjectClient as SlicesCoreProjectClient,
)

from .aiis import AIInfraServiceClient
from .base.config import SlicesAIApiClientsConfig, load_default_aiis_api_clients_config


def _load_file_or_url(location: str, client: httpx.Client) -> str:
    if location.startswith("http"):
        response = client.get(location, headers={"Accept": "application/json"})
        try:
            response.raise_for_status()
            return response.content.decode()
        except HTTPError as err:
            raise from_httpx_exception(err) from err
    elif Path(location).exists():
        with Path(location).open() as f:
            return f.read()
    else:
        raise ValueError(f"{location!r} is not a valid path or URL")


def _get_http_user_agent():
    """Get HTTP User Agent."""
    return f"slices-aiis-client/{importlib.metadata.version('slices_clientlib_aiis')}"


DEFAULT_HTTP_USER_AGENT = _get_http_user_agent()


def _get_site_id(site_or_infra_id: str) -> str:
    """Get site ID from site or infrastructure ID."""

    id_parts = site_or_infra_id.split("-")

    if len(id_parts) > 2 and id_parts[2] == "ai":  # noqa: PLR2004
        # It's an infrastructure/cluster ID: {site_id}-ai-{cluster_id}
        return f"{id_parts[0]}-{id_parts[1]}"
    # Assume it's a site ID
    return site_or_infra_id


class SlicesClientWithAIInfraService(SlicesCoreClient):
    """Slices Client, with support for AI Infrastructure Service."""

    _aiis_api_clients_config: SlicesAIApiClientsConfig

    def __init__(
        self,
        central_config: SlicesCentralApiClientsConfig | str | None = None,
        ai_config: SlicesAIApiClientsConfig | str | None = None,
        ignore_config_cache: bool = False,
        http_user_agent: str = DEFAULT_HTTP_USER_AGENT,
    ) -> None:
        """Create a Slices AIIS client."""

        super().__init__(
            central_config=central_config,
            ignore_config_cache=ignore_config_cache,
            http_user_agent=http_user_agent,
        )

        self._init_param_ai_config = ai_config

        # Load AIIS config
        if ai_config and isinstance(ai_config, str):
            self._aiis_api_clients_config = msgspec.json.decode(
                _load_file_or_url(ai_config, self._no_auth_client), type=SlicesAIApiClientsConfig
            )
        elif ai_config and isinstance(ai_config, SlicesAIApiClientsConfig):
            self._aiis_api_clients_config = ai_config
        elif ai_config is None:
            self._aiis_api_clients_config = load_default_aiis_api_clients_config(self._no_auth_client)
        else:
            raise TypeError(f"Unsupported ai_config {ai_config!r}")

    def reload_ai_config(self) -> bool:
        """Reload the AIIS slices config from the default online source.

        This is useful because it is locally cached.

        This does nothing if the central slices config was not loaded from
        the default online source or its cached version.

        Returns: True if reload was done, False if reload was skipped.
        """
        if self._init_param_ai_config is not None:
            return False
        self._aiis_api_clients_config = load_default_aiis_api_clients_config(self._no_auth_client, ignore_cache=True)
        return True

    @property
    def ai_api_clients_config(self) -> SlicesAIApiClientsConfig:
        """Get AIIS API Clients Config."""
        return self._aiis_api_clients_config

    def for_project(self, project_id: str) -> "SlicesProjectClientWithAIInfraService":
        """Get Slices Project Client."""
        return SlicesProjectClientWithAIInfraService(project_id, self)

    def for_experiment(self, experiment_id: str, project_id: str) -> "SlicesExperimentClientWithAIInfraService":
        """Get Slices Project Client."""
        return SlicesExperimentClientWithAIInfraService(project_id, experiment_id, self)


class SlicesProjectClientWithAIInfraService(SlicesCoreProjectClient):
    """Slices client for a specific project, with support for AI Infrastructure Service."""

    def __init__(self, project_id: str, generic_client: SlicesClientWithAIInfraService) -> None:
        """Create Slices Project Client, with support for AI Infrastructure Service."""
        super().__init__(project_id, generic_client)
        if not project_id.startswith("proj_"):
            raise ValueError(f"Invalid project ID: {project_id!r}")
        self._aiis_api_clients_config = generic_client._aiis_api_clients_config
        self._generic_aiis_client = generic_client

    def for_experiment(self, experiment_id: str) -> "SlicesExperimentClientWithAIInfraService":
        """Get Slices Project Client."""
        return SlicesExperimentClientWithAIInfraService(experiment_id, self.project_id, self._generic_aiis_client)

    def ai(self, site_or_infra_id: str | None = None) -> AIInfraServiceClient:
        """Get AI Infrastructure Service API Client for a specific site."""

        if site_or_infra_id:
            site_id = _get_site_id(site_or_infra_id)
        else:
            site_id = self._aiis_api_clients_config.ai_service_default_site

        aiis_api_entry = self._aiis_api_clients_config.ai_service_sites.get(site_id)

        if not aiis_api_entry:
            raise ValueError(f"Unknown Site ID: {site_id}")

        return AIInfraServiceClient(
            site_id,
            aiis_api_entry,
            self._generic_client._create_httpx_client_for_api_entry(
                api_entry=aiis_api_entry, project_id=self.project_id
            ),
            self,
            self._generic_aiis_client,
        )


class SlicesExperimentClientWithAIInfraService(SlicesCoreExperimentClient):
    """Slices client for a specific experiment, with support for AI Infrastructure Service."""

    def __init__(self, experiment_id: str, project_id: str, generic_client: SlicesClientWithAIInfraService) -> None:
        """Create Slices Experiment Client, with support for AI Infrastructure Service."""
        super().__init__(experiment_id, project_id, generic_client)
        if not experiment_id.startswith("exp_"):
            raise ValueError(f"Invalid experiment ID: {experiment_id!r}")
        if not project_id.startswith("proj_"):
            raise ValueError(f"Invalid project ID: {project_id!r}")

        self._aiis_api_clients_config = generic_client._aiis_api_clients_config
        self._generic_aiis_client = generic_client

    def ai(self, site_or_infra_id: str | None = None) -> AIInfraServiceClient:
        """Get AI Infrastructure Service API Client for a specific site."""

        if site_or_infra_id:
            site_id = _get_site_id(site_or_infra_id)
        else:
            site_id = self._aiis_api_clients_config.ai_service_default_site

        ai_api_entry = self._aiis_api_clients_config.ai_service_sites.get(site_id)

        if not self.project_id:
            raise ValueError("Problem with Project ID")
        if not self.experiment_id:
            raise ValueError("Problem with Experiment ID")

        if not ai_api_entry:
            raise ValueError(f"Unknown site ID: {site_id}")

        client = self._generic_client._create_httpx_client_for_api_entry(
            api_entry=ai_api_entry, project_id=self.project_id, experiment_id=self.experiment_id
        )
        return AIInfraServiceClient(site_id, ai_api_entry, client, self, self._generic_aiis_client)
