import json
import uuid
from collections.abc import Iterator
from pathlib import Path
from unittest.mock import MagicMock, patch

import httpx
import pytest
import responses
import respx
from cognite.client.data_classes import Annotation, AnnotationList, Asset, AssetList
from cognite.client.data_classes.data_modeling import (
    DataModel,
    DataModelList,
    EdgeApply,
    NodeApply,
    NodeOrEdgeData,
    View,
    ViewId,
)
from cognite.client.data_classes.data_modeling.statistics import InstanceStatistics, ProjectStatistics

from cognite_toolkit._cdf_tk.client import ToolkitClient, ToolkitClientConfig
from cognite_toolkit._cdf_tk.client.data_classes.charts import Chart
from cognite_toolkit._cdf_tk.client.data_classes.charts_data import ChartData, ChartSource, ChartTimeseries
from cognite_toolkit._cdf_tk.client.data_classes.migration import InstanceSource
from cognite_toolkit._cdf_tk.client.testing import monkeypatch_toolkit_client
from cognite_toolkit._cdf_tk.commands._migrate.command import MigrationCommand
from cognite_toolkit._cdf_tk.commands._migrate.data_mapper import AssetCentricMapper, ChartMapper
from cognite_toolkit._cdf_tk.commands._migrate.data_model import (
    COGNITE_MIGRATION_MODEL,
    INSTANCE_SOURCE_VIEW_ID,
    MODEL_ID,
    RESOURCE_VIEW_MAPPING_VIEW_ID,
)
from cognite_toolkit._cdf_tk.commands._migrate.default_mappings import (
    ASSET_ANNOTATIONS_ID,
    ASSET_ID,
    FILE_ANNOTATIONS_ID,
    create_default_mappings,
)
from cognite_toolkit._cdf_tk.commands._migrate.migration_io import AnnotationMigrationIO, AssetCentricMigrationIO
from cognite_toolkit._cdf_tk.commands._migrate.selectors import MigrationCSVFileSelector
from cognite_toolkit._cdf_tk.exceptions import ToolkitMigrationError, ToolkitValueError
from cognite_toolkit._cdf_tk.storageio import ChartIO
from cognite_toolkit._cdf_tk.storageio.selectors import ChartExternalIdSelector
from cognite_toolkit._cdf_tk.utils.fileio import CSVReader


@pytest.fixture
def cognite_migration_model(
    toolkit_config: ToolkitClientConfig,
    rsps: responses.RequestsMock,
    cognite_core_no_3D: DataModel[View],
    cognite_extractor_views: list[View],
) -> Iterator[responses.RequestsMock]:
    """Mock the Cognite Migration Model in the CDF project."""
    config = toolkit_config
    # Migration model
    migration_model = COGNITE_MIGRATION_MODEL.dump()
    migration_model["createdTime"] = 1
    migration_model["lastUpdatedTime"] = 1
    migration_model["isGlobal"] = False
    rsps.post(config.create_api_url("models/datamodels/byids"), json={"items": migration_model})
    yield rsps


@pytest.fixture
def resource_view_mappings(
    toolkit_config: ToolkitClientConfig,
    cognite_migration_model: responses.RequestsMock,
    cognite_core_no_3D: DataModel[View],
    cognite_extractor_views: list[View],
) -> Iterator[responses.RequestsMock]:
    """Mock all the default Resource View Mappings in the Cognite Migration Model."""
    rsps = cognite_migration_model
    config = toolkit_config
    mapping_by_id = {mapping.external_id: mapping for mapping in create_default_mappings()}
    node_items: list[dict] = []
    for mapping in mapping_by_id.values():
        # Lookup of the mapping in the Migration Model
        mapping_node_response = mapping.dump(context="api")
        mapping_node_response.update({"createdTime": 0, "lastUpdatedTime": 0, "version": 1})
        sources = mapping_node_response.pop("sources", [])
        if sources:
            mapping_view_id = mapping.sources[0].source
            mapping_node_response["properties"] = {
                mapping_view_id.space: {
                    f"{mapping_view_id.external_id}/{mapping_view_id.version}": sources[0]["properties"]
                }
            }
        node_items.append(mapping_node_response)
    rsps.post(
        config.create_api_url("models/instances/byids"),
        json={"items": node_items},
        status=200,
    )
    rsps.post(
        config.create_api_url("models/views/byids"),
        json={
            "items": [view.dump() for view in cognite_core_no_3D.views]
            + [view.dump() for view in cognite_extractor_views]
        },
    )
    yield rsps


@pytest.fixture
def mock_statistics(
    toolkit_config: ToolkitClientConfig,
    rsps: responses.RequestsMock,
) -> Iterator[responses.RequestsMock]:
    config = toolkit_config
    stats_response = {
        "spaces": {
            "count": 0,
            "limit": 1_000,
        },
        "containers": {
            "count": 0,
            "limit": 10_000,
        },
        "views": {
            "count": 0,
            "limit": 100_000,
        },
        "dataModels": {
            "count": 1,
            "limit": 10_000,
        },
        "containerProperties": {
            "count": 0,
            "limit": 1_000_000,
        },
        "instances": {
            "nodes": 1000,
            "edges": 0,
            "softDeletedNodes": 0,
            "softDeletedEdges": 0,
            "instancesLimit": 5_000_000,
            "softDeletedInstancesLimit": 100_000_000,
            "instances": 1000,
            "softDeletedInstances": 0,
        },
        "concurrentReadLimit": 50,
        "concurrentWriteLimit": 20,
        "concurrentDeleteLimit": 10,
    }
    rsps.get(
        config.create_api_url("/models/statistics"),
        json=stats_response,
        status=200,
    )
    yield rsps


@pytest.mark.usefixtures("disable_gzip", "disable_pypi_check")
class TestMigrationCommand:
    @pytest.mark.usefixtures("mock_statistics")
    def test_migrate_assets(
        self,
        toolkit_config: ToolkitClientConfig,
        resource_view_mappings: responses.RequestsMock,
        tmp_path: Path,
        respx_mock: respx.MockRouter,
    ) -> None:
        rsps = resource_view_mappings
        config = toolkit_config
        assets = AssetList(
            [
                Asset(
                    id=1000 + i,
                    external_id=f"asset_{i}",
                    name=f"Asset {i}",
                    description=f"This is Asset {i}",
                    last_updated_time=1,
                    created_time=0,
                    parent_external_id="asset_0" if i > 0 else None,
                )
                for i in range(2)
            ]
        )
        space = "my_space"
        csv_content = "id,space,externalId,ingestionView\n" + "\n".join(
            f"{1000 + i},{space},asset_{i},{ASSET_ID}" for i in range(len(assets))
        )

        # Asset retrieve ids
        rsps.post(
            config.create_api_url("/assets/byids"),
            json={"items": [asset.dump() for asset in assets]},
            status=200,
        )
        # Instance creation
        respx.post(
            config.create_api_url("/models/instances"),
        ).mock(
            return_value=httpx.Response(
                status_code=200,
                json={
                    "items": [
                        {
                            "instanceType": "node",
                            "space": space,
                            "externalId": f"asset_{i}",
                            "version": 1,
                            "wasModified": True,
                            "createdTime": 1,
                            "lastUpdatedTime": 1,
                        }
                        for i in range(len(assets))
                    ]
                },
            )
        )
        csv_file = tmp_path / "migration.csv"
        csv_file.write_text(csv_content, encoding="utf-8")

        client = ToolkitClient(config)
        command = MigrationCommand(silent=True)

        result = command.migrate(
            selected=MigrationCSVFileSelector(datafile=csv_file, kind="Assets"),
            data=AssetCentricMigrationIO(client),
            mapper=AssetCentricMapper(client),
            log_dir=tmp_path / "logs",
            dry_run=False,
            verbose=False,
        )

        # Check that the assets were uploaded
        last_call = respx_mock.calls[-1]
        assert last_call.request.url == config.create_api_url("/models/instances")
        assert last_call.request.method == "POST"
        actual_instances = json.loads(last_call.request.content)["items"]
        expected_instance = [
            NodeApply(
                space=space,
                external_id=asset.external_id,
                sources=[
                    NodeOrEdgeData(
                        source=ViewId("cdf_cdm", "CogniteAsset", "v1"),
                        properties={
                            "name": asset.name,
                            "description": asset.description,
                        },
                    ),
                    NodeOrEdgeData(
                        source=INSTANCE_SOURCE_VIEW_ID,
                        properties={
                            "id": asset.id,
                            "resourceType": "asset",
                            "dataSetId": None,
                            "classicExternalId": asset.external_id,
                        },
                    ),
                ],
            ).dump()
            for asset in assets
        ]
        assert actual_instances == expected_instance
        actual_results = [result.get_progress(f"asset_{asset.id}") for asset in assets]
        expected_results = [{"download": "success", "convert": "success", "upload": "success"} for _ in assets]
        assert actual_results == expected_results
        csv_file = next((tmp_path / "logs").glob("*.csv"), None)
        assert csv_file is not None, "Expected a CSV log file to be created"
        csv_results = list(CSVReader(csv_file).read_chunks_unprocessed())
        assert csv_results == [
            {"ID": f"asset_{asset.id}", "download": "success", "convert": "success", "upload": "success"}
            for asset in assets
        ]

    @pytest.mark.usefixtures("mock_statistics")
    def test_migrate_annotations(
        self,
        toolkit_config: ToolkitClientConfig,
        resource_view_mappings: responses.RequestsMock,
        tmp_path: Path,
        respx_mock: respx.MockRouter,
    ) -> None:
        rsps = resource_view_mappings
        config = toolkit_config
        asset_annotation = Annotation(
            id=2000,
            annotated_resource_type="file",
            annotated_resource_id=3000,
            data={
                "assetRef": {"id": 4000},
                "textRegion": {"xMin": 10.0, "xMax": 100.0, "yMin": 20.0, "yMax": 200.0},
            },
            status="Approved",
            creating_user="doctrino",
            creating_app="my_app",
            creating_app_version="v1",
            annotation_type="diagrams.AssetLink",
        )
        file_annotation = Annotation(
            id=2001,
            annotated_resource_type="file",
            annotated_resource_id=3001,
            data={
                "fileRef": {"id": 5000},
                "textRegion": {"xMin": 15.0, "xMax": 150.0, "yMin": 25.0, "yMax": 250.0},
            },
            status="Approved",
            creating_user="doctrino",
            creating_app="my_app",
            creating_app_version="v1",
            annotation_type="diagrams.FileLink",
        )
        annotations = AnnotationList([asset_annotation, file_annotation])
        space = "my_space"
        csv_content = "id,space,externalId,ingestionView\n" + "\n".join(
            (
                f"{2000},{space},annotation_{2000},{ASSET_ANNOTATIONS_ID}",
                f"{2001},{space},annotation_{2001},{FILE_ANNOTATIONS_ID}",
            )
        )
        # Annotation retrieve ids
        rsps.post(
            config.create_api_url("/annotations/byids"),
            json={"items": [annotation.dump() for annotation in annotations]},
            status=200,
        )
        # Lookup asset and file instance ID
        for items in [
            [("asset", 4000)],
            [("file", 5000), ("file", 3000), ("file", 3001)],
        ]:
            rsps.post(
                config.create_api_url("/models/instances/query"),
                json={
                    "items": {
                        "instanceSource": [
                            {
                                "instanceType": "node",
                                "space": space,
                                "externalId": f"{resource_type}_{resource_id}",
                                "version": 0,
                                "createdTime": 0,
                                "lastUpdatedTime": 0,
                                "properties": {
                                    "cognite_migration": {
                                        "InstanceSource/v1": {
                                            "id": resource_id,
                                            "resourceType": resource_type,
                                        }
                                    },
                                },
                            }
                            for (resource_type, resource_id) in items
                        ],
                    },
                    "nextCursor": {"instanceSource": None},
                },
                status=200,
            )

        # Instance creation
        respx.post(
            config.create_api_url("/models/instances"),
        ).mock(
            return_value=httpx.Response(
                status_code=200,
                json={
                    "items": [
                        {
                            "instanceType": "edge",
                            "space": space,
                            "externalId": f"annotation_{2000 + i}",
                            "version": 1,
                            "wasModified": True,
                            "createdTime": 1,
                            "lastUpdatedTime": 1,
                        }
                        for i in range(len(annotations))
                    ]
                },
            )
        )
        csv_file = tmp_path / "migration.csv"
        csv_file.write_text(csv_content, encoding="utf-8")

        client = ToolkitClient(config)
        command = MigrationCommand(silent=True)

        result = command.migrate(
            selected=MigrationCSVFileSelector(datafile=csv_file, kind="Annotations"),
            data=AnnotationMigrationIO(client),
            mapper=AssetCentricMapper(client),
            log_dir=tmp_path / "logs",
            dry_run=False,
            verbose=True,
        )
        actual_results = [result.get_progress(f"Annotation_{annotation.id}") for annotation in annotations]
        expected_results = [{"download": "success", "convert": "success", "upload": "success"} for _ in annotations]
        assert actual_results == expected_results

        # Check that the annotations were uploaded
        last_call = respx_mock.calls[-1]
        assert last_call.request.url == config.create_api_url("/models/instances")
        assert last_call.request.method == "POST"
        actual_instances = json.loads(last_call.request.content)["items"]
        expected_instance = [
            EdgeApply(
                space=space,
                external_id=f"annotation_{asset_annotation.id}",
                start_node=(space, f"file_{asset_annotation.annotated_resource_id}"),
                end_node=(space, f"asset_{asset_annotation.data['assetRef']['id']}"),
                type=(space, asset_annotation.annotation_type),
                sources=[
                    NodeOrEdgeData(
                        source=ViewId("cdf_cdm", "CogniteDiagramAnnotation", "v1"),
                        properties={
                            "sourceContext": asset_annotation.creating_app_version,
                            "sourceCreatedUser": asset_annotation.creating_user,
                            "sourceId": asset_annotation.creating_app,
                            "status": asset_annotation.status,
                            "startNodeXMax": asset_annotation.data["textRegion"]["xMax"],
                            "startNodeXMin": asset_annotation.data["textRegion"]["xMin"],
                            "startNodeYMax": asset_annotation.data["textRegion"]["yMax"],
                            "startNodeYMin": asset_annotation.data["textRegion"]["yMin"],
                        },
                    ),
                ],
            ).dump(),
            EdgeApply(
                space=space,
                external_id=f"annotation_{file_annotation.id}",
                start_node=(space, f"file_{file_annotation.annotated_resource_id}"),
                end_node=(space, f"file_{file_annotation.data['fileRef']['id']}"),
                type=(space, file_annotation.annotation_type),
                sources=[
                    NodeOrEdgeData(
                        source=ViewId("cdf_cdm", "CogniteDiagramAnnotation", "v1"),
                        properties={
                            "sourceContext": file_annotation.creating_app_version,
                            "sourceCreatedUser": file_annotation.creating_user,
                            "sourceId": file_annotation.creating_app,
                            "status": file_annotation.status,
                            "startNodeXMax": file_annotation.data["textRegion"]["xMax"],
                            "startNodeXMin": file_annotation.data["textRegion"]["xMin"],
                            "startNodeYMax": file_annotation.data["textRegion"]["yMax"],
                            "startNodeYMin": file_annotation.data["textRegion"]["yMin"],
                        },
                    ),
                ],
            ).dump(),
        ]
        assert actual_instances == expected_instance

    def test_migrate_charts(
        self,
        toolkit_config: ToolkitClientConfig,
        cognite_migration_model: responses.RequestsMock,
        tmp_path: Path,
        respx_mock: respx.MockRouter,
    ) -> None:
        rsps = cognite_migration_model
        config = toolkit_config
        charts = [
            Chart(
                external_id="my_chart",
                created_time=1,
                last_updated_time=1,
                visibility="PUBLIC",
                data=ChartData(
                    name="My Chart",
                    time_series_collection=[
                        ChartTimeseries(
                            ts_external_id="ts_1", type="timeseries", id="87654321-4321-8765-4321-876543218765"
                        ),
                        ChartTimeseries(ts_id=1, type="timeseries", id="12345678-1234-5678-1234-567812345678"),
                    ],
                    source_collection=[
                        ChartSource(type="timeseries", id="87654321-4321-8765-4321-876543218765"),
                        ChartSource(type="timeseries", id="12345678-1234-5678-1234-567812345678"),
                    ],
                ),
                owner_id="1234",
            )
        ]
        # Chart retrieve ids
        rsps.add(
            responses.POST,
            config.create_app_url("/storage/charts/charts/list"),
            json={
                "items": [chart.dump() for chart in charts],
            },
            status=200,
        )
        # TimeSeries Instance ID lookup
        rsps.add(
            method=responses.POST,
            url=config.create_api_url("models/instances/query"),
            json={
                "items": {
                    "instanceSource": [
                        InstanceSource(
                            space="my_space",
                            external_id="node_123",
                            version=1,
                            last_updated_time=1,
                            created_time=1,
                            resource_type="timeseries",
                            id_=1,
                            classic_external_id=None,
                            preferred_consumer_view_id=ViewId("cdf_cdm", "CogniteTimeSeries", "v1"),
                        ).dump(),
                        InstanceSource(
                            space="my_space",
                            external_id="node_ts_1",
                            version=1,
                            last_updated_time=1,
                            created_time=1,
                            resource_type="timeseries",
                            id_=2,
                            classic_external_id="ts_1",
                            preferred_consumer_view_id=ViewId("my_schema_space", "MyTimeSeries", "v1"),
                        ).dump(),
                    ]
                },
                "nextCursor": {"instanceSource": None},
            },
            status=200,
        )

        # Chart upsert
        respx.put(
            config.create_app_url("/storage/charts/charts"),
        ).mock(return_value=httpx.Response(status_code=200, json={"items": [chart.dump() for chart in charts]}))

        client = ToolkitClient(config)
        command = MigrationCommand(silent=True)
        new_uuids = [
            uuid.UUID("aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa"),
            uuid.UUID("bbbbbbbb-bbbb-bbbb-bbbb-bbbbbbbbbbbb"),
        ]
        with patch(f"{ChartMapper.__module__}.uuid4", side_effect=new_uuids):
            result = command.migrate(
                selected=ChartExternalIdSelector(external_ids=("my_chart",)),
                data=ChartIO(client),
                mapper=ChartMapper(client),
                log_dir=tmp_path / "logs",
                dry_run=False,
                verbose=True,
            )
        actual = result.get_progress("my_chart")
        expected = {"download": "success", "convert": "success", "upload": "success"}
        assert actual == expected

        calls = respx_mock.calls
        assert len(calls) == 1
        last_call = calls[-1]
        assert last_call.request.url == config.create_app_url("/storage/charts/charts")
        assert last_call.request.method == "PUT"
        actual_charts = json.loads(last_call.request.content)["items"]
        expected_charts = [
            {
                "externalId": "my_chart",
                "visibility": "PUBLIC",
                "data": {
                    "name": "My Chart",
                    "coreTimeseriesCollection": [
                        {
                            "type": "coreTimeseries",
                            "id": "aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa",
                            "nodeReference": {"space": "my_space", "externalId": "node_ts_1"},
                            "viewReference": {
                                "space": "my_schema_space",
                                "externalId": "MyTimeSeries",
                                "version": "v1",
                            },
                        },
                        {
                            "type": "coreTimeseries",
                            "id": "bbbbbbbb-bbbb-bbbb-bbbb-bbbbbbbbbbbb",
                            "nodeReference": {"space": "my_space", "externalId": "node_123"},
                            "viewReference": {"space": "cdf_cdm", "externalId": "CogniteTimeSeries", "version": "v1"},
                        },
                    ],
                    "sourceCollection": [
                        {
                            "type": "coreTimeseries",
                            "id": "aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa",
                        },
                        {
                            "type": "coreTimeseries",
                            "id": "bbbbbbbb-bbbb-bbbb-bbbb-bbbbbbbbbbbb",
                        },
                    ],
                },
            }
        ]
        assert actual_charts == expected_charts

    def test_validate_migration_model_available(self) -> None:
        with monkeypatch_toolkit_client() as client:
            client.data_modeling.data_models.retrieve.return_value = DataModelList([])
            with pytest.raises(ToolkitMigrationError):
                MigrationCommand.validate_migration_model_available(client)

    def test_validate_migration_model_available_multiple_models(self) -> None:
        """Test that multiple models raises an error."""
        with monkeypatch_toolkit_client() as client:
            # Create mock models with the expected MODEL_ID
            model1 = MagicMock(spec=DataModel)
            model1.as_id.return_value = MODEL_ID
            model2 = MagicMock(spec=DataModel)
            model2.as_id.return_value = MODEL_ID

            client.data_modeling.data_models.retrieve.return_value = DataModelList([model1, model2])

            with pytest.raises(ToolkitMigrationError) as exc_info:
                MigrationCommand.validate_migration_model_available(client)

            assert "Multiple migration models" in str(exc_info.value)

    def test_validate_migration_model_available_missing_views(self) -> None:
        """Test that a model with missing views raises an error."""
        with monkeypatch_toolkit_client() as client:
            model = MagicMock(spec=DataModel)
            model.as_id.return_value = MODEL_ID
            # Model has views but missing the required ones
            model.views = [INSTANCE_SOURCE_VIEW_ID]  # Missing VIEW_SOURCE_VIEW_ID

            client.data_modeling.data_models.retrieve.return_value = DataModelList([model])

            with pytest.raises(ToolkitMigrationError, match=r"Invalid migration model. Missing views"):
                MigrationCommand.validate_migration_model_available(client)

    def test_validate_migration_model_available_success(self) -> None:
        """Test that a valid model with all required views succeeds."""
        with monkeypatch_toolkit_client() as client:
            # Mocking the migration Model to get a response format of the model.
            # An alternative would be to write a conversion of write -> read format of the model
            # which is a significant amount of logic.
            model = MagicMock(spec=DataModel)
            model.as_id.return_value = MODEL_ID
            # Model has all required views
            model.views = [INSTANCE_SOURCE_VIEW_ID, RESOURCE_VIEW_MAPPING_VIEW_ID]

            client.data_modeling.data_models.retrieve.return_value = DataModelList([model])

            # Should not raise any exception
            MigrationCommand.validate_migration_model_available(client)

            client.data_modeling.data_models.retrieve.assert_called_once_with([MODEL_ID], inline_views=False)

    def test_validate_available_capacity_missing_capacity(self) -> None:
        cmd = MigrationCommand(silent=True)

        with monkeypatch_toolkit_client() as client:
            stats = MagicMock(spec=ProjectStatistics)
            stats.instances = InstanceStatistics(
                nodes=1000,
                edges=0,
                soft_deleted_edges=0,
                soft_deleted_nodes=0,
                instances_limit=1500,
                soft_deleted_instances_limit=10_000,
                instances=1000,
                soft_deleted_instances=0,
            )
            client.data_modeling.statistics.project.return_value = stats
            with pytest.raises(ToolkitValueError) as exc_info:
                cmd.validate_available_capacity(client, 10_000)

        assert "Cannot proceed with migration" in str(exc_info.value)

    def test_validate_available_capacity_sufficient_capacity(self) -> None:
        cmd = MigrationCommand(silent=True)

        with monkeypatch_toolkit_client() as client:
            stats = MagicMock(spec=ProjectStatistics)
            stats.instances = InstanceStatistics(
                nodes=1000,
                edges=0,
                soft_deleted_edges=0,
                soft_deleted_nodes=0,
                instances_limit=5_000_000,
                soft_deleted_instances_limit=100_000_000,
                instances=1000,
                soft_deleted_instances=0,
            )
            client.data_modeling.statistics.project.return_value = stats
            cmd.validate_available_capacity(client, 10_000)
