from typing import Dict, List, Optional, Union

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

from qhub.api.quantum import HubQuantumClient, Job as JobModel
from qhub.quantum.sdk.job import JobRequest, JOB_FINAL_STATES
from qhub.quantum.sdk.qiskit import HubQiskitJob
from qhub.quantum.sdk.qiskit.providers.iqm.hub_iqm_client import _HubIqmClient


class HubIqmQiskitJob(HubQiskitJob):

    def __init__(self, backend: Optional[Backend], job_id: Optional[str] = None,
                 job_details: Optional[Union[JobModel, JobRequest]] = None,
                 hub_client: Optional[Union[HubQuantumClient, _HubIqmClient]] = None):
        super().__init__(backend, job_id, job_details, hub_client)

    def _result_states(self) -> set[str]:
        return JOB_FINAL_STATES

    def _create_experiment_result(self, provider_result: dict) -> ExperimentResult:
        """Transform IQM measurement results to Qiskit ExperimentResult format."""
        result = self.iqm_to_experiment_result(provider_result)
        if result.results:
            return result.results[0]
        else:
            return ExperimentResult(
                shots=0,
                success=True,
                data=ExperimentResultData(counts={})
            )

    def iqm_to_experiment_result(self,
            iqm_meas_results: List[Dict[str, Union[List[str], Dict[str, int]]]],
    ) -> Result:
        """Convert IQM-style measurement results into a Qiskit Result."""
        if not iqm_meas_results:
            return Result(
                backend_name=self._backend.name if self._backend else "unknown",
                backend_version="1.0",
                job_id=self._job_id,
                qobj_id=0,
                success=True,
                results=[],
                status="COMPLETED",
                date=None,
            )

        experiment_results = []
        for i, measurement_batch in enumerate(iqm_meas_results):
            counts_dict = measurement_batch.get("counts", {})
            total_shots = sum(counts_dict.values())

            memory = []
            for state, count in counts_dict.items():
                hex_state = hex(int(state, 2))
                memory.extend([hex_state] * count)

            experiment_result = {
                "shots": total_shots,
                "success": True,
                "data": {
                    "memory": memory,
                    "counts": Counts(counts_dict),
                },
                "header": {"name": f"circ{i}"},
                "status": "DONE",
            }
            experiment_results.append(experiment_result)

        result_dict = {
            "backend_name": self._backend.name if self._backend else "unknown",
            "backend_version": "1.0",
            "qobj_id": 0,
            "job_id": self._job_id,
            "success": True,
            "results": experiment_results,
            "date": None,
        }
        return Result.from_dict(result_dict)
