# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.

from abc import ABC, abstractmethod
from concurrent.futures import Executor, Future
import logging
from time import monotonic
from typing import Callable, Dict, Optional, Any

from qiskit.providers import BackendV2
from qiskit.circuit import (
    QuantumCircuit,
)

from qiskit.result import Result
from qiskit.providers import JobV1, JobStatus, JobError

from ..execution import DetaultExecutor
from ... import telemetry_events
from ...estimator import EstimatorResult

logger = logging.getLogger(__name__)

RunInputCallable = Callable[[QuantumCircuit, str, Dict[str, Any]], Result]


class QsJob(JobV1, ABC):
    """
    Abstract base class for Q# Qiskit jobs.

    Manages asynchronous execution of a quantum circuit via a callable submitted
    to a thread pool. Concrete subclasses must implement :meth:`result` and
    :meth:`_submit_duration`.
    """

    def __init__(
        self,
        backend: Optional[BackendV2],
        job_id: str,
        job_callable: RunInputCallable,
        run_input: QuantumCircuit,
        input_params: Dict[str, Any],
        executor=None,
        **kwargs,
    ) -> None:
        """
        :param backend: The backend on which the job is run.
        :param job_id: A unique identifier for the job.
        :type job_id: str
        :param job_callable: The callable that executes the circuit and returns a result.
        :param run_input: The quantum circuit to execute.
        :param input_params: Parameters forwarded to ``job_callable`` at execution time.
        :type input_params: Dict
        :param executor: Thread pool executor. Uses a default single-threaded executor if
            not provided.
        :param **kwargs: Additional keyword arguments passed to ``qiskit.providers.JobV1``.
        """

        self._run_input = run_input
        self._input_params = input_params
        self._future: Optional[Future] = None
        self._executor: Executor = executor or DetaultExecutor()
        self._job_callable = job_callable
        self._status = JobStatus.INITIALIZING
        self._submit_start_time: Optional[float] = None
        super().__init__(backend, job_id, **kwargs)

    def submit(self):
        """Submit the job to the backend for execution.

        :raises JobError: If trying to re-submit the job.
        """
        if self._future is not None:
            raise JobError("Job has already been submitted.")

        self._submit_start_time = monotonic()
        self._future = self._executor.submit(
            self._job_callable, self._run_input, self.job_id(), **self._input_params
        )
        self.add_done_callback(self._submit_duration)

    @abstractmethod
    def result(self, timeout: Optional[float] = None) -> Any:
        pass

    @abstractmethod
    def _submit_duration(self, _future: Future):
        pass

    def _result(self, timeout: Optional[float] = None) -> Any:
        """Return the results of the job."""
        if self._future is None:
            raise JobError("Job has not been submitted.")

        return self._future.result(timeout=timeout)

    def status(self) -> JobStatus:
        """Return the status of the job, among the values of ``JobStatus``."""
        if self._future is None:
            return JobStatus.INITIALIZING
        if self._future.cancelled():
            return JobStatus.CANCELLED
        if self._future.done():
            if self._future.exception() is None:
                return JobStatus.DONE
            else:
                return JobStatus.ERROR
        if self._future.running():
            return JobStatus.RUNNING

        return JobStatus.INITIALIZING

    def backend(self) -> BackendV2:
        """Return the backend where this job was executed."""

        return super().backend()

    def cancel(self):
        """Attempt to cancel the job."""
        if self._future is not None:
            self._future.cancel()

    def error(self) -> Optional[JobError]:
        """Return the error that occurred during the execution of the job."""
        if self._future is not None:
            return self._future.exception()

    def add_done_callback(self, fn: Callable[[Future[Result]], object]) -> None:
        """Attaches a callable that will be called when the job finishes."""
        if self._future is not None:
            self._future.add_done_callback(fn)


class QsSimJob(QsJob):
    """
    A Qiskit job that runs a quantum circuit on the Q# simulator.

    Submits the circuit for simulation and returns a ``qiskit.result.Result``
    containing shot-level measurement outcomes.
    """

    def result(self, timeout: Optional[float] = None) -> Result:
        return self._result(timeout=timeout)

    def submit(self):
        """Submit the job to the backend for execution.

        :raises JobError: If trying to re-submit the job.
        """
        shots = self._input_params.get("shots", -1)
        telemetry_events.on_qiskit_run(shots, 1)

        super().submit()

    def _submit_duration(self, _future: Future):
        end_time = monotonic()
        # _submit_start_time is set in submit() before adding this callback
        assert self._submit_start_time is not None
        duration_in_sec = end_time - self._submit_start_time
        duration_in_ms = duration_in_sec * 1000

        shots = self._input_params.get("shots", -1)
        telemetry_events.on_qiskit_run_end(shots, 1, duration_in_ms)


class ReJob(QsJob):
    """
    A Qiskit job that runs the Q# Resource Estimator.

    Submits the circuit to the resource estimator and returns an
    :class:`~qsharp.estimator.EstimatorResult` with the computed resource estimates.
    """

    def result(self, timeout: Optional[float] = None) -> EstimatorResult:
        return self._result(timeout=timeout)

    def submit(self):
        """Submit the job to the backend for execution.

        :raises JobError: If trying to re-submit the job.
        """

        telemetry_events.on_qiskit_run_re()

        super().submit()

    def _submit_duration(self, _future: Future):
        end_time = monotonic()
        # _submit_start_time is set in submit() before adding this callback
        assert self._submit_start_time is not None
        duration_in_sec = end_time - self._submit_start_time
        duration_in_ms = duration_in_sec * 1000

        telemetry_events.on_qiskit_run_re_end(duration_in_ms)
