# Licensed under the Apache License, Version 2.0 (the "License"); you may
# not use this file except in compliance with the License. You may obtain
# a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations
# under the License.
"""End-to-end test with mock http response.

This test relies on the snapshots of the HTTP responses from the real IBM runtime server.
Responses are stored in JSON files after sensitive information is removed.
Because the HTTP response is hard-coded, any change in the test payload doesn't update the result.
This test aims to prevent breaking logic changes against the server response.
Note that this test doesn't make any real HTTP request. Everything works in offline.

When the data schema changes in the future, JSON files must be regenerated.
Current HTTP response is generated by the IBM-API-Version "2025-01-01".
"""

import json
from pathlib import Path
from unittest.mock import ANY, AsyncMock, MagicMock

import aiohttp
import pytest
from prefect import get_client
from prefect.client.schemas.filters import ArtifactFilter
from qiskit.circuit import QuantumCircuit
from qiskit.transpiler import generate_preset_pass_manager
from utils import assert_sampler_fidelity

from prefect_qiskit import QuantumRuntime
from prefect_qiskit.vendors.ibm_quantum import IBMQuantumCredentials

RESPONSE = Path(__file__).parent / "http_response"


@pytest.fixture(autouse=True, scope="function")
def mock_http_response(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    """A factory function of dynamic monkeypatch for request of aiohttp response."""

    def get_mock_resp(*args, **kwargs) -> AsyncMock:  # type: ignore[no-untyped-def]
        mock_ctx = AsyncMock()

        match str(args[0]):  # URL
            case "backends":
                # get_resources
                with (RESPONSE / "backend_list.json").open() as fp:
                    resp_json = json.load(fp)
                mock_response = MagicMock()
                mock_response.json = AsyncMock(return_value=resp_json)
            case "backends/test_eagle_us-east/status":
                # check_resource_available
                with (RESPONSE / "backend_status.json").open() as fp:
                    resp_json = json.load(fp)
                mock_response = MagicMock()
                mock_response.json = AsyncMock(return_value=resp_json)
            case "jobs/test-job-id-123":
                # get_job_status
                with (RESPONSE / "job_status.json").open() as fp:
                    resp_json = json.load(fp)
                mock_response = MagicMock()
                mock_response.json = AsyncMock(return_value=resp_json)
            case "jobs/test-job-id-123/results":
                # get_primitive_result
                with (RESPONSE / "job_results.json").open() as fp:
                    resp_text = json.dumps(json.load(fp))
                mock_response = MagicMock()
                mock_response.text = AsyncMock(return_value=resp_text)
            case "jobs/test-job-id-123/metrics":
                # get_job_metrics
                with (RESPONSE / "job_metrics.json").open() as fp:
                    resp_json = json.load(fp)
                mock_response = MagicMock()
                mock_response.json = AsyncMock(return_value=resp_json)
            case "backends/test_eagle_us-east/configuration":
                # get_target / configuration
                with (RESPONSE / "backend_configuration.json").open() as fp:
                    resp_json = json.load(fp)
                mock_response = MagicMock()
                mock_response.json = AsyncMock(return_value=resp_json)
            case "backends/test_eagle_us-east/properties":
                # get_target / properties
                with (RESPONSE / "backend_properties.json").open() as fp:
                    resp_json = json.load(fp)
                mock_response = MagicMock()
                mock_response.json = AsyncMock(return_value=resp_json)
            case _:
                mock_response = MagicMock()
                mock_response.json = AsyncMock(return_value={})
        mock_ctx.__aenter__.return_value = mock_response

        return mock_ctx

    def post_mock_resp(*args, **kwargs) -> AsyncMock:  # type: ignore[no-untyped-def]
        mock_ctx = AsyncMock()

        match str(args[0]):  # URL
            case "jobs":
                # run_primitive
                resp_json = {
                    "backend": "test_eagle_us-east",
                    "id": "test-job-id-123",
                }
                mock_response = MagicMock()
                mock_response.json = AsyncMock(return_value=resp_json)
            case _:
                mock_response = MagicMock()
                mock_response.json = AsyncMock(return_value={})
        mock_ctx.__aenter__.return_value = mock_response

        return mock_ctx

    # Setup mock session
    mock_session = MagicMock()
    setattr(
        mock_session,
        "get",
        MagicMock(side_effect=get_mock_resp),
    )
    setattr(
        mock_session,
        "post",
        MagicMock(side_effect=post_mock_resp),
    )

    # Monkeypatch ClientSession with mock
    mock_session_ctx = AsyncMock()
    mock_session_ctx.__aenter__.return_value = mock_session
    mock_session_ctx.__aexit__.return_value = None

    monkeypatch.setattr(aiohttp, "ClientSession", MagicMock(return_value=mock_session_ctx))


def test_sampler_e2e(
    bell_circ: QuantumCircuit,
) -> None:
    """End to End test of sampler execution on IBM Quantum channel with mock server response."""
    credentials = IBMQuantumCredentials(
        api_key="test_api_key",
        crn="crn:test-name",
    )

    runtime = QuantumRuntime(
        resource_name="test_eagle_us-east",
        credentials=credentials,
        enable_job_analytics=True,
    )

    pm = generate_preset_pass_manager(
        optimization_level=2,
        target=runtime.get_target(),
    )
    isa_qc = pm.run(bell_circ)

    result = runtime.sampler(sampler_pubs=[isa_qc], options={"params": {"shots": 100}})

    assert result[0].data.meas.num_shots == 100
    assert_sampler_fidelity(result[0], {"00": 1, "11": 1}, atol=0.3)

    # Check job report artifact
    with get_client(sync_client=True) as client:
        artifacts = client.read_artifacts(
            artifact_filter=ArtifactFilter(key={"any_": ["job-metrics"]}),
        )
    assert len(artifacts) == 1

    job_metrics = dict(zip(*json.loads(artifacts[0].data)))

    ref_dict = {
        "resource": "test_eagle_us-east",
        "program_type": "sampler",
        "num_pubs": 1,
        "job_id": "test-job-id-123",
        "tags": [],
        "timestamp.created": ANY,
        "timestamp.started": ANY,
        "timestamp.completed": ANY,
        "span.queue": ANY,
        "span.work": ANY,
        "span.qpu": ANY,
        "work_efficiency": ANY,
        "pub[0].circuit.depth": isa_qc.depth(),
        "pub[0].circuit.size": isa_qc.size(),
        "pub[0].shape": [],
        "pub[0].timestamp.started": ANY,
        "pub[0].timestamp.completed": ANY,
        "pub[0].duration": ANY,
        "options.params.shots": "100",
    }

    assert job_metrics == ref_dict
