"""Common code for the GeoJson exporter."""

import json
from pathlib import Path

from kili_formats.format.geojson import convert_from_kili_to_geojson_format
from kili_formats.types import Job, JobTool

from kili.domain.ontology import JobMLTask
from kili.services.export.exceptions import NotCompatibleInputType, NotCompatibleOptions
from kili.services.export.format.base import AbstractExporter
from kili.services.export.tools import is_geotiff_asset_with_lat_lon_coords
from kili.utils.tqdm import tqdm


class GeoJsonExporter(AbstractExporter):
    """Common code for GeoJson exporter."""

    def _check_arguments_compatibility(self) -> None:
        """Checks if the export label format is compatible with the export options."""
        if self.normalized_coordinates is not None:
            raise NotCompatibleOptions(
                "The GeoJson annotation format can only be exported with lat/long coordinates."
                " Please set `normalized_coordinates` to None."
            )

        if self.split_option != "merged":
            raise NotCompatibleOptions(
                "The current implementation only supports merged annotations."
            )

        if self.single_file:
            raise NotCompatibleOptions(
                "The GeoJson annotation format cannot be exported into a single file."
            )

    def _check_project_compatibility(self) -> None:
        """Checks if the export label format is compatible with the project."""
        if self.project["inputType"] not in ["IMAGE", "GEOSPATIAL"]:
            raise NotCompatibleInputType(
                f"Project with input type '{self.project['inputType']}' not compatible with"
                " GeoJson export format."
            )

    def _is_job_compatible(self, job: Job) -> bool:
        """Check if the export label format is compatible with the job."""
        if "tools" not in job:
            return False

        if job["mlTask"] != JobMLTask.OBJECT_DETECTION:
            return False

        compatible_tools = {
            JobTool.RECTANGLE,
            JobTool.POLYGON,
            JobTool.SEMANTIC,
            JobTool.MARKER,
            JobTool.POLYLINE,
        }

        return all(
            tool in compatible_tools
            for tool in job["tools"]  # pyright: ignore[reportGeneralTypeIssues]
        )

    def process_and_save(self, assets: list[dict], output_filename: Path) -> None:
        self.logger.info("Exporting to GeoJson format")

        labels_folder = self.base_folder / "labels"
        labels_folder.mkdir(parents=True, exist_ok=True)
        project_type = self.project.get("inputType")
        geotiff_assets = [
            asset
            for asset in assets
            if project_type == "GEOSPATIAL"
            or is_geotiff_asset_with_lat_lon_coords(asset, self.kili.http_client)
        ]
        if len(geotiff_assets) < len(assets):
            self.logger.warning(
                f"Among {len(assets)} assets, only {len(geotiff_assets)} are geotiff assets and"
                " will be exported."
            )

        # Get json_interface for GIS-friendly property names
        json_interface = self.project.get("jsonInterface")

        for asset in tqdm(geotiff_assets, disable=self.disable_tqdm):
            _process_asset(asset, labels_folder, json_interface, flatten_properties=True)

        self.create_readme_kili_file(self.export_root_folder)
        self.make_archive(self.export_root_folder, output_filename)

        self.logger.warning(output_filename)


def _process_asset(
    asset: dict,
    labels_folder: Path,
    json_interface: dict | None = None,
    flatten_properties: bool = False,
) -> None:
    # Collect all labels to process (handle both latestLabel and latestLabels)
    labels_to_process = []
    if "latestLabel" in asset and asset["latestLabel"]:
        labels_to_process.append(asset["latestLabel"])
    if "latestLabels" in asset and asset["latestLabels"]:
        for label in asset["latestLabels"]:
            if label is not None:
                labels_to_process.append(label)

    if not labels_to_process:
        return

    # Process each label
    for label_idx, latest_label in enumerate(labels_to_process, start=1):
        # Add label suffix if we have multiple labels
        label_suffix = f"_label{label_idx}" if len(labels_to_process) > 1 else ""

        geojson_feature_collection = convert_from_kili_to_geojson_format(
            latest_label["jsonResponse"], json_interface, flatten_properties
        )
        filepath = labels_folder / f"{asset['externalId']}{label_suffix}.geojson"
        filepath.parent.mkdir(parents=True, exist_ok=True)
        with open(filepath, "w", encoding="utf-8") as file:
            json.dump(geojson_feature_collection, file)
