from typing import Optional, Union

from qiskit.providers import JobV1, JobStatus, Backend
from qiskit.result import Result
from qiskit.result.models import ExperimentResult, ExperimentResultData

from qhub.api.quantum import HubQuantumClient, Job as JobModel
from qhub.quantum.sdk.job import HubBaseJob, JobRequest

JobStatusMap = {
    "CREATED": JobStatus.INITIALIZING,
    "PENDING": JobStatus.QUEUED,
    "RUNNING": JobStatus.RUNNING,
    "COMPLETED": JobStatus.DONE,
    "FAILED": JobStatus.ERROR,
    "CANCELLING": JobStatus.RUNNING,
    "CANCELLED": JobStatus.CANCELLED,
    "UNKNOWN": JobStatus.INITIALIZING,
}


class HubQiskitJob(HubBaseJob, JobV1):
    def __init__(self, backend: Optional[Backend], job_id: Optional[str] = None,
                 job_details: Optional[Union[JobModel, JobRequest]] = None,
                 hub_client: Optional[HubQuantumClient] = None):
        """
        Constructor for internal use only.

        Args:
            backend: The backend where the job was executed
            job_id: The job ID - this must be only provided if an existing job is retrieved
            job_details: The job details (JobModel) or job request (JobRequest) for new jobs
            hub_client: The Kipu Quantum Hub client
        """
        HubBaseJob.__init__(self, backend=backend, job_id=job_id, job_details=job_details, hub_client=hub_client)
        JobV1.__init__(self, backend=backend, job_id=job_id, shots=self._job_details.shots)

    def submit(self):
        super()._submit()

    def _create_experiment_result(self, provider_result: dict) -> ExperimentResult:
        return ExperimentResult(
            shots=self.shots,
            success=True,
            status=JobStatus.DONE.name,
            data=ExperimentResultData(
                counts=provider_result.get("counts") or {},
                memory=provider_result.get("memory") or []
            ),
        )

    def result(self) -> Result:
        """
        Return the result of the job.
        """
        provider_result_data = super()._result()

        experiment_result = self._create_experiment_result(provider_result_data)

        result = Result(
            backend_name=self._backend.name,
            backend_version=self._backend.version,
            job_id=self._job_id,
            qobj_id=0,
            success=True,
            results=[experiment_result],
            status=JobStatus.DONE,
            date=self._job_details.ended_at,
        )

        return result

    def status(self) -> JobStatus:
        """
        Return the status of the job.
        """
        status = super()._update_state()
        return JobStatusMap[status]

    def to_dict(self) -> dict:
        """
        Return a dictionary representation of the job.
        """
        return {key: value for key, value in vars(self).items() if not key.startswith('_')}

    def queue_position(self):
        """
        Return the position of the job in the server queue.
        """
        return None


# Keep for backward compatibility reasons
HubJob = HubQiskitJob
