import logging
import os
import random
import string
import time
import typing

from authlib.integrations.httpx_client import OAuth2Client
from qhub.api.service import (
    HubServiceClient as BaseHubServiceClient,
    ServiceExecution,
    ServiceExecutionStatus,
    ResultResponse,
    LogEntry,
)
from qhub.api.service.service_api.client import ServiceApiClient

DEFAULT_TOKEN_ENDPOINT = "https://gateway.hub.kipu-quantum.com/token"


class HubServiceAuth:
    _logger = logging.getLogger(__name__)

    def __init__(
        self,
        access_key_id: str,
        secret_access_key: str,
        token_endpoint: str = DEFAULT_TOKEN_ENDPOINT,
    ):
        self._access_key_id = access_key_id
        self._secret_access_key = secret_access_key
        self._token_endpoint = token_endpoint
        self._token = None
        self._last_token_fetch = None

    def get_token(self) -> str:
        if self._token is None or self._last_token_fetch is None:
            self._refresh_token()

        # Refresh token 120 seconds before expiration to prevent race conditions
        if time.time() - self._last_token_fetch >= (self._token["expires_in"] - 120):
            self._refresh_token()

        return self._token["access_token"]

    def _refresh_token(self):
        self._logger.debug("Creating new token using client credentials flow")
        with OAuth2Client(self._access_key_id, self._secret_access_key) as client:
            token = client.fetch_token(
                self._token_endpoint, grant_type="client_credentials"
            )

        self._token = token
        self._last_token_fetch = time.time()


class HubServiceExecution:
    def __init__(self, client: "HubServiceClient", service_execution: ServiceExecution):
        self._client = client
        self._service_execution = service_execution

    @property
    def id(self) -> str:
        return self._service_execution.id

    @property
    def status(self) -> ServiceExecutionStatus:
        return self._service_execution.status

    @property
    def created_at(self) -> str:
        return self._service_execution.created_at

    @property
    def started_at(self) -> typing.Optional[str]:
        return self._service_execution.started_at

    @property
    def ended_at(self) -> typing.Optional[str]:
        return self._service_execution.ended_at

    @property
    def has_finished(self) -> bool:
        self.refresh()
        return self.status in ["SUCCEEDED", "FAILED", "CANCELLED"]

    def wait_for_final_state(
        self, timeout: typing.Optional[float] = None, wait: float = 5
    ) -> None:
        """
        Poll the status until it progresses to a final state.

        Parameters:
            - timeout: Seconds to wait for the job. If ``None``, wait indefinitely.
            - wait: Seconds between queries.

        Raises:
            Exception: If the service execution does not reach a final state before the specified timeout.
        """
        start_time = time.time()
        while not self.has_finished:
            elapsed_time = time.time() - start_time
            if timeout is not None and elapsed_time >= timeout:
                raise TimeoutError(
                    f"Timeout while waiting for service execution '{self.id}'."
                )
            time.sleep(wait)

    def refresh(self):
        self._service_execution = self._client.api.get_status(id=self.id)

    def result(self) -> ResultResponse:
        self.wait_for_final_state()
        delay = 1  # Start with a small delay
        max_delay = 16  # Maximum delay
        while True:
            try:
                result = self._client.api.get_result(id=self.id)
                break  # If the operation succeeds, break out of the loop
            except Exception as e:
                time.sleep(delay)  # If the operation fails, wait
                delay *= 2  # Double the delay
                if delay >= max_delay:
                    raise e  # If the delay is too long, raise the exception
        return result

    def result_files(self) -> typing.List[str]:
        file_names = []
        links = self.result().links
        for link in links:
            file_name = link[0]
            if file_name in ["status", "self"]:
                continue
            file_names.append(file_name)
        return file_names

    def result_file_stream(self, file_name: str) -> typing.Iterator[bytes]:
        return self._client.api.get_result_file(id=self.id, file=file_name)

    def download_result_file(self, file_name: str, target_path: str) -> None:
        # check if target path exists and is a directory
        if not os.path.isdir(target_path):
            raise ValueError(
                f"Target path '{target_path}' does not exist or is not a directory."
            )

        abs_target_path = os.path.abspath(target_path)
        file_path = os.path.join(abs_target_path, file_name)

        iterator = self.result_file_stream(file_name)
        with open(file_path, "wb") as f:
            for chunk in iterator:
                f.write(chunk)

    def cancel(self) -> None:
        self._client.api.cancel_execution(id=self.id)

    def logs(self) -> typing.List[LogEntry]:
        return self._client.api.get_logs(id=self.id)


class HubServiceClient:
    def __init__(
        self,
        service_endpoint: str,
        access_key_id: typing.Union[str, None],
        secret_access_key: typing.Union[str, None],
        token_endpoint: str = DEFAULT_TOKEN_ENDPOINT,
    ):
        self._service_endpoint = service_endpoint
        self._access_key_id = access_key_id
        self._secret_access_key = secret_access_key
        self._token_endpoint = token_endpoint

        if (self._access_key_id is not None) or (self._secret_access_key is not None):
            self._auth = HubServiceAuth(
                access_key_id=self._access_key_id,
                secret_access_key=self._secret_access_key,
                token_endpoint=self._token_endpoint,
            )
            self._api = BaseHubServiceClient(
                base_url=self._service_endpoint, token=self._auth.get_token
            )
        else:
            random_token = "".join(
                random.choices(string.ascii_letters + string.digits, k=21)
            )
            self._api = BaseHubServiceClient(
                base_url=self._service_endpoint, token=random_token
            )

    @property
    def api(self) -> ServiceApiClient:
        return self._api.service_api

    def run(
        self,
        request: typing.Dict[str, typing.Any],
        secrets: typing.Dict[str, typing.Any] = None,
        tags: typing.List[str] = None,
    ) -> HubServiceExecution:
        if tags is not None and "$tags" not in request:
            request["$tags"] = tags
        if secrets is not None and "$secrets" not in request:
            request["$secrets"] = secrets
        service_execution = self.api.start_execution(request=request)
        return HubServiceExecution(client=self, service_execution=service_execution)

    def get_service_execution(self, service_execution_id: str) -> HubServiceExecution:
        service_execution = self.api.get_status(id=service_execution_id)
        return HubServiceExecution(client=self, service_execution=service_execution)

    def get_service_executions(self) -> typing.List[HubServiceExecution]:
        service_executions = self.api.get_service_executions()
        return [
            HubServiceExecution(client=self, service_execution=e)
            for e in service_executions
        ]
