import logging
import time
from abc import ABC
from dataclasses import dataclass, field
from typing import Optional, Dict, Any, Union

from qhub.api.quantum import HubQuantumClient, Job as JobModel
from qhub.api.quantum.core import ApiError
from qhub.quantum.sdk.exceptions import HubError

logger = logging.getLogger(__name__)

JOB_FINAL_STATES = {"ABORTED", "COMPLETED", "CANCELLED", "FAILED"}


@dataclass
class JobRequest:
    """Data class for job submission parameters."""
    backend_id: str
    shots: int
    input: Any = None
    input_format: Optional[str] = None
    input_params: Optional[Any] = None
    sdk_provider: Optional[str] = None
    session_id: Optional[str] = None
    name: Optional[str] = None
    tags: Optional[list] = None


class HubBaseJob(ABC):

    def __init__(self, backend: Optional, job_id: Optional[str] = None,
                 job_details: Optional[Union[JobModel, JobRequest]] = None,
                 hub_client: Optional[HubQuantumClient] = None):
        self._client = hub_client
        if self._client is None:
            raise RuntimeError("hub_client must not be None")

        if job_id is None and job_details is None:
            raise ValueError("Either 'job_id', 'job_details' or both must be provided.")

        self._result = None
        self._backend = backend
        self._job_details = None
        self._job_request = None
        self._status = None

        if isinstance(job_details, JobRequest):
            self._job_request = job_details
            self._submit()
        elif isinstance(job_details, JobModel):
            self._job_details = job_details
        elif job_id is not None and job_details is None:
            self._refresh(job_id)
        else:
            self._job_details = job_details

    def _submit(self):
        """
        Submits the job for execution.
        """
        if self._job_request is None:
            raise RuntimeError("Cannot submit job as no job request is set.")

        req = self._job_request
        self._job_details = self._client.jobs.create_job(
            backend_id=req.backend_id,
            shots=req.shots,
            input=req.input,
            input_format=req.input_format,
            input_params=req.input_params,
            session_id=req.session_id,
            name=req.name,
            tags=req.tags,
        )

    def _result_states(self) -> set[str]:
        """Returns the set of job states that are considered final for result retrieval."""
        return {"COMPLETED"}

    def _result(self) -> any:
        """Polls for the job result until the job has completed successfully and returns the result.

        Returns:
            The job result if the job completed successfully. The result format is backend specific.

        Raises:
            An error if the job did reach another end state as `COMPLETED`
        """
        if self._result is not None:
            return self._result

        status = self._update_state()
        if status not in JOB_FINAL_STATES:
            status = self._wait_for_final_state()

        if status in self._result_states():
            self._result = self._client.jobs.get_job_result(self.id)
            return self._result

        if status == "FAILED":
            error_result_data = self._client.jobs.get_job_result(self.id)
            msg = (
                f"Cannot retrieve results because the job execution failed with status '{status}'."
            )
            if error_result_data:
                msg += f" Reason: {error_result_data}."
            raise RuntimeError(msg)

        raise RuntimeError(
            f'{"Cannot retrieve results as job execution has not completed."}'
            + f"(status: {self._status}.")

    def _wait_for_final_state(self, timeout: Optional[float] = None, wait: float = 5) -> str:
        """Poll the job status until it progresses to a final state such as ``DONE`` or ``ERROR``.

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

        Raises:
            HubError: If the job does not reach a final state before the specified timeout.
        """
        start_time = time.time()
        status = self._update_state()
        while status not in JOB_FINAL_STATES:
            logger.debug('Waiting for job %s to complete for retrieving its result. Current job status %s', self.id, status)
            elapsed_time = time.time() - start_time
            if timeout is not None and elapsed_time >= timeout:
                raise HubError(f"Timeout while waiting for job {self.id}.")
            time.sleep(wait)
            status = self._update_state()
        return status

    def _update_state(self, use_cached_value: bool = False) -> str:
        """
        Return the status of the job.

        Args:
            use_cached_value (bool): If `True`, uses the value most recently retrieved from Kipu Quantum Hub.
        """
        if not use_cached_value and self._status not in JOB_FINAL_STATES:
            response = self._client.jobs.get_job_status(self.id)
            self._status = response.status

        return self._status

    @property
    def id(self) -> str:
        """
        This job's id.
        """
        return self._job_details.id if self._job_details is not None else None

    def job_id(self) -> str:
        """
        This job's id.
        """
        return self.id

    @property
    def shots(self) -> int:
        return self._job_details.shots

    def cancel(self):
        """
        Attempt to cancel the job.
        """
        self._client.jobs.cancel_job(self.id)

    def _refresh(self, job_id: str = None):
        """
        Refreshes the job details from the server.
        """
        if job_id is None and self.id is None:
            raise ValueError("Job Id is not set.")

        job_id = job_id if job_id is not None else self.id
        self._job_details = self._client.jobs.get_job(job_id)
        self._update_state()

    def calibration(self) -> Optional[Dict[str, Any]]:
        """
        Retrieve the backend calibration that was effective when the job execution started.

        Note:
            The backend calibration data is unavailable if the job was executed on a simulator
            or if the job has not been executed yet.

        Returns:
            Optional[Dict[str, Any]]: Backend calibration data or None if no calibration data is available.
            The data format is backend specific.
        """
        try:
            response = self._client.jobs.get_job_calibration(self.id)
            return response.calibration if response else None
        except ApiError as e:
            if e.status_code == 204:
                return None
            raise
