"""
Kipu Quantum Hub Runtime Job V2 Implementation

This class extends HubRuntimeJob to provide full compatibility with IBM's
RuntimeJobV2 interface, including all public methods from RuntimeJobV2 and its parent classes.
"""

import json
from datetime import datetime
from typing import Optional, Type, Any, Dict, Union, Sequence, List

from qiskit.providers import Backend
from qiskit_ibm_runtime import IBMError, RuntimeJobMaxTimeoutError, RuntimeJobFailureError, RuntimeInvalidStateError
from qiskit_ibm_runtime.constants import DEFAULT_DECODERS, API_TO_JOB_ERROR_MESSAGE
from qiskit_ibm_runtime.models.backend_properties import BackendProperties
from qiskit_ibm_runtime.runtime_job_v2 import JobStatus
from qiskit_ibm_runtime.utils import utc_to_local
from qiskit_ibm_runtime.utils.result_decoder import ResultDecoder

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

JOB_STATUS_MAP: Dict[str, JobStatus] = {
    "CREATED": "INITIALIZING",
    "PENDING": "QUEUED",
    "RUNNING": "RUNNING",
    "COMPLETED": "DONE",
    "FAILED": "ERROR",
    "CANCELLING": "RUNNING",
    "CANCELLED": "CANCELLED",
    "UNKNOWN": "INITIALIZING",
}


class HubRuntimeJobV2(HubQiskitJob):
    """Kipu Quantum Hub Runtime Job V2 with full IBM RuntimeJobV2 compatibility."""

    def __init__(
            self,
            backend: Optional[Backend],
            job_id: Optional[str] = None,
            job_details: Optional[Union[JobModel, JobRequest]] = None,
            result_decoder: Optional[Union[Type[ResultDecoder], Sequence[Type[ResultDecoder]]]] = None,
            hub_client: Optional[HubQuantumClient] = None,
            session_id: Optional[str] = None,
            program_id: Optional[str] = None,
            provider: "HubQiskitRuntimeService" = None
    ):
        # Validate mandatory provider parameter
        if provider is None:
            raise ValueError("provider parameter is required for HubRuntimeJobV2")

        super().__init__(backend=backend, job_id=job_id, job_details=job_details, hub_client=hub_client)

        # Store provider for backend creation
        self._provider = provider

        # Set runtime-specific parameters from constructor parameters or job details
        self._session_id = session_id
        self._program_id = program_id

        # Try to extract from job details input_params if not provided directly
        if self._job_details and hasattr(self._job_details, 'input_params') and self._job_details.input_params:
            params = self._job_details.input_params
            if isinstance(params, dict):
                if self._session_id is None:
                    self._session_id = params.get('session_id', None)
                if self._program_id is None:
                    self._program_id = params.get('program_id', None)

        # Set up result decoders (interim and final)
        decoder = result_decoder or DEFAULT_DECODERS.get(self._program_id, None) or ResultDecoder
        if isinstance(decoder, Sequence):
            self._interim_result_decoder, self._final_result_decoder = decoder
        else:
            self._interim_result_decoder = self._final_result_decoder = decoder

        self._error_message = None
        self._reason = None
        self._reason_code = None
        self._result = None

    def result(
            self,
            timeout: Optional[float] = None,
            decoder: Optional[Type[ResultDecoder]] = None,
    ) -> Any:
        """Return the results of the job."""
        _decoder = decoder or self._final_result_decoder

        if self._result is None or (_decoder != self._final_result_decoder):
            self.wait_for_final_state(timeout=timeout)
            status = self.status()

            if status == "ERROR":
                error_message = self._reason if self._reason else self._error_message
                if self._reason == 1305:
                    raise RuntimeJobMaxTimeoutError(error_message)
                raise RuntimeJobFailureError(f"Unable to retrieve job result. {error_message}")

            if status == "CANCELLED":
                raise RuntimeInvalidStateError(
                    f"Unable to retrieve result for job {self.job_id()}. Job was cancelled."
                )

            result_raw = self._get_raw_result()
            self._result = _decoder.decode(json.dumps(result_raw)) if result_raw else None

        return self._result

    def status(self) -> JobStatus:
        """Return the status of the job."""
        self._set_status_and_error_message()
        return JOB_STATUS_MAP[self._status]

    def _get_raw_result(self) -> Dict:
        """Fetch raw job result from Kipu Quantum Hub. Returns error data if job failed."""
        try:
           return super()._result()
        except RuntimeError:
            return self._client.jobs.get_job_result(self.id)

    def cancelled(self) -> bool:
        return self._update_state() == "CANCELLED"

    def done(self) -> bool:
        return self._update_state() == "COMPLETED"

    def running(self) -> bool:
        return self._update_state() == "RUNNING"

    def in_final_state(self) -> bool:
        return self._update_state() in JOB_FINAL_STATES

    def errored(self) -> bool:
        return self._update_state() == "FAILED"

    def submit(self) -> None:
        raise NotImplementedError("Job submission is not supported for Kipu Quantum Hub runtime jobs V2. Use the primitive's run function instead.")

    # noinspection PyMethodOverriding
    def wait_for_final_state(self, timeout: Optional[float] = None) -> None:
        self._wait_for_final_state(timeout=timeout)

    def queue_info(self) -> Optional[Dict[str, Any]]:
        raise NotImplementedError("Queue information is not supported for Kipu Quantum Hub runtime jobs V2.")

    def queue_position(self) -> Optional[int]:
        raise NotImplementedError("Queue position is not supported for Kipu Quantum Hub runtime jobs V2.")

    def error_message(self) -> Optional[str]:
        self._set_status_and_error_message()
        return self._error_message

    def backend(self, timeout: Optional[float] = None) -> Optional[Backend]:
        if self._backend is None and self._provider:
                self._backend = self._provider.backend(self._job_details.backend_id)
        return self._backend

    def _set_status_and_error_message(self) -> None:
        """Fetch and set status and error message."""
        if self._update_state() in {"FAILED", "CANCELLED"}:
            error_data = self._get_raw_result()

            if error_data:
                self._set_status(error_data)
                self._set_error_message(error_data)

    def _set_status(self, job_result: Dict) -> None:
        try:
            reason = job_result.get("reason")
            reason_code = job_result.get("reasonCode")
            if reason:
                self._reason = reason
                if reason_code:
                    self._reason = f"Error code {reason_code}; {self._reason}"
                    self._reason_code = reason_code
            self._status = self._status_from_job_response()
        except KeyError:
            raise IBMError(f"Unknown status: {job_result['status']}")

    def _set_error_message(self, job_result: Dict) -> None:
        if self.errored():
            self._error_message = self._error_msg_from_job_response(str(job_result))
        else:
            self._error_message = None

    def _error_msg_from_job_response(self, job_response_str: str) -> str:
        index = job_response_str.rfind("Traceback")
        if index != -1:
            job_response_str = job_response_str[index:]

        if self.cancelled() and self._reason_code == 1305:
            error_msg = API_TO_JOB_ERROR_MESSAGE["CANCELLED - RAN TOO LONG"]
            return error_msg.format(self.job_id(), job_response_str)
        else:
            error_msg = API_TO_JOB_ERROR_MESSAGE["FAILED"]
            return error_msg.format(self.job_id(), self._reason or job_response_str)

    def _status_from_job_response(self) -> Union[JobStatus, str]:
        if self._status == "CANCELLED" and self._reason_code == 1305:
            return "FAILED"
        return self._status

    def scheduling_mode(self) -> Optional[str]:
        raise NotImplementedError("Scheduling mode is not supported for Kipu Quantum Hub runtime jobs V2.")

    def usage_estimation(self) -> Dict[str, Any]:
        raise NotImplementedError("Usage estimation is not supported for Kipu Quantum Hub runtime jobs V2.")

    def update_name(self, name: str) -> str:
        raise NotImplementedError("Name update is not supported for Kipu Quantum Hub runtime jobs V2.")

    def update_tags(self, replacement_tags: List[str]) -> List[str]:
        raise NotImplementedError("Tag update is not supported for Kipu Quantum Hub runtime jobs V2.")

    def usage(self) -> float:
        raise NotImplementedError("Usage tracking is not supported for Kipu Quantum Hub runtime jobs V2.")

    def metrics(self) -> Dict[str, Any]:
        raise NotImplementedError("Metrics collection is not supported for Kipu Quantum Hub runtime jobs V2.")

    def logs(self) -> str:
        raise NotImplementedError("Log retrieval is not supported for Kipu Quantum Hub runtime jobs V2.")

    def properties(self, refresh: bool = False) -> Optional[BackendProperties]:
        return self._backend.properties(refresh)

    @property
    def name(self) -> Optional[str]:
        return self.job_id()

    @property
    def tags(self) -> List[str]:
        raise NotImplementedError("Job tags are not supported for Kipu Quantum Hub runtime jobs V2.")

    @property
    def time_per_step(self) -> Optional[Dict[str, Any]]:
        raise NotImplementedError("Time per step is not supported for Kipu Quantum Hub runtime jobs V2.")

    @property
    def estimated_start_time(self) -> Optional[datetime]:
        raise NotImplementedError("Estimated start time is not supported for Kipu Quantum Hub runtime jobs V2.")

    @property
    def estimated_completion_time(self) -> Optional[datetime]:
        raise NotImplementedError("Estimated completion time is not supported for Kipu Quantum Hub runtime jobs V2.")

    @property
    def version(self) -> int:
        raise NotImplementedError("Job version is not supported for Kipu Quantum Hub runtime jobs V2.")

    @property
    def backend_version(self) -> str:
        raise NotImplementedError("Backend version is not supported for Kipu Quantum Hub runtime jobs V2.")

    @property
    def image(self) -> str:
        return ""

    @property
    def inputs(self) -> Dict:
        if not self._job_details:
            return {}
        # For JobModel, the input field is available
        input_val = getattr(self._job_details, 'input', None)
        if input_val is None:
            return {}
        return input_val

    @property
    def primitive_id(self) -> str:
        return self._program_id

    @property
    def creation_date(self) -> Optional[datetime]:
        return utc_to_local(self._job_details.created_at)

    @property
    def session_id(self) -> str:
        return self._session_id

    @property
    def instance(self) -> Optional[str]:
        raise NotImplementedError("instances method not implemented. Instance management is handled by Kipu Quantum Hub.")

    def __repr__(self) -> str:
        return f"<{self.__class__.__name__}('{self._job_id}', '{self._program_id}')>"
