# Mypy; for the `|` operator purpose
# Remove this __future__ import once the oldest supported Python is 3.10
from __future__ import annotations

import contextlib
import logging
import os
import time
from typing import TYPE_CHECKING, Any

from . import job_metadata_constants
from .api import jobs
from .backend_connection import BackendConnection
from .check_version import check_version
from .estimate_result import EstimateResult
from .exceptions import (
    BQJobCouldNotCancelError,
    BQJobInvalidDeviceTypeError,
    BQJobNotCompleteError,
)
from .job_result import JobResult
from .version import __version__

if TYPE_CHECKING:
    import datetime

# TODO this requires imports of actual quantum libraries for proper type
# checking.
CircuitT = Any

logger = logging.getLogger("bluequbit-python-sdk")


class BQClient:
    """Client for managing jobs on BlueQubit platform.

    :param api_token: API token of the user. If ``None``, the token will be looked
                      in default configuration file ``$HOME/.config/bluequbit/config.json``.
                      If not ``None``, the token will also be saved in the same
                      default configuration file.
    :param update_config_file: if True, update default configuration file
                               if api_token is not None.
    """

    def __init__(self, api_token: str | None = None, update_config_file: bool = True):
        if os.environ.get("BLUEQUBIT_TESTING") is None:
            with contextlib.suppress(Exception):
                check_version(__version__)

        self._backend_connection = BackendConnection(api_token, update_config_file)

    @staticmethod
    def validate_device(device):
        if not isinstance(device, str):
            raise BQJobInvalidDeviceTypeError(device)
        converted_device = device.lower()
        if converted_device not in job_metadata_constants.DEVICE_TYPES:
            raise BQJobInvalidDeviceTypeError(device)
        return converted_device

    def estimate(self, circuit: CircuitT, device: str = "cpu") -> EstimateResult:
        """Estimate job runtime

        :param circuit: quantum circuit
        :type circuit: Cirq, Qiskit, circuit
        :param device: device for which to estimate the circuit. Can be one of
                       ``"cpu"`` | ``"gpu"`` | ``"quantum"``
        :return: estimate result metadata
        """
        device = self.validate_device(device)
        return EstimateResult(
            jobs.submit_job(
                self._backend_connection, circuit, device, estimate_only=True
            )
        )

    # PLR0913 is for "Too many arguments to function call (6/5)"
    def run(  # noqa: PLR0913
        self,
        circuit: CircuitT,
        device: str = "cpu",
        asynchronous: bool = False,
        job_name: str | None = None,
        shots: int = 1000,
    ) -> JobResult:
        """Submit a job to run on BlueQubit platform

        :param circuit: quantum circuit
        :type circuit: Cirq, Qiskit, circuit
        :param device: device on which to run the circuit. Can be one of
                       ``"cpu"`` | ``"gpu"`` | ``"quantum"``
        :param asynchronous: if set to ``False``, wait for job completion before
                             returning. If set to ``True``, return immediately
        :param job_name: customizable job name
        :param shots: number of shots to run if device is quantum.
        :return: job metadata
        """
        device = self.validate_device(device)
        submitted_job = JobResult(
            jobs.submit_job(
                self._backend_connection, circuit, device, job_name, shots=shots
            )
        )
        if submitted_job.run_status == "FAILED_VALIDATION":
            raise BQJobNotCompleteError(
                submitted_job.job_id,
                submitted_job.run_status,
                submitted_job.error_message,
            )
        logger.info("Submitted: %s", submitted_job)
        if not asynchronous:
            return self.wait(submitted_job.job_id)
        return submitted_job

    def wait(self, job_id: str) -> JobResult:
        """Wait for job completion

        :param job_id: job ID that can be found as property of :class:`JobResult` metadata
                       of :func:`~run` method
        :return: job metadata
        """
        while True:
            job_result = self.get(job_id)
            if job_result.run_status in job_metadata_constants.JOB_RESULTS_READY_STATES:
                return job_result
            elif (
                job_result.run_status
                in job_metadata_constants.JOB_NO_RESULT_TERMINAL_STATES
            ):
                raise BQJobNotCompleteError(
                    job_id, job_result.run_status, job_result.error_message
                )
            else:
                time.sleep(1.0)

    def get(self, job_id: str) -> JobResult:
        """Get current metadata of job

        :param job_id: job ID that can be found as property of :class:`JobResult` metadata
                       of :func:`~run` method
        :return: job metadata
        """
        return JobResult(jobs.get(self._backend_connection, job_id))

    def cancel(self, job_id: str) -> JobResult:
        """Submit job cancel request

        :param job_id: job ID that can be found as property of :class:`JobResult` metadata
                       of :func:`run` method
        :return: job metadata
        """
        jobs.cancel_job(self._backend_connection, job_id)
        try:
            self.wait(job_id)
        except BQJobNotCompleteError as e:
            if not e.run_status == "CANCELED":
                raise BQJobCouldNotCancelError(
                    e.job_id, e.run_status, e.error_message
                ) from None
        return self.get(job_id)

    def search(
        self,
        run_status: str | None = None,
        created_later_than: str | datetime.datetime | None = None,
    ) -> list[JobResult]:
        """Search jobs

        :param run_status: if not ``None``, run status of jobs to filter.
                           Can be one of ``"FAILED_VALIDATION"`` | ``"PENDING"`` |
                           ``"QUEUED"`` | ``"RUNNING"`` | ``"TERMINATED"`` | ``"CANCELED"`` |
                           ``"NOT_ENOUGH_FUNDS"`` | ``"COMPLETED"``

        :param created_later_than: if not ``None``, filter by latest job creation datetime.
                                   Please add timezone for clarity, otherwise UTC
                                   will be assumed

        :return: metadata of jobs
        """
        job_results = jobs.search_jobs(
            self._backend_connection, run_status, created_later_than
        )
        return [JobResult(r) for r in job_results["data"]]
