import shutil
from collections import defaultdict
from collections.abc import Callable
from pathlib import Path
from typing import Any

from rich import print
from rich.table import Table

from cognite_toolkit._cdf_tk.client import ToolkitClient
from cognite_toolkit._cdf_tk.commands._base import ToolkitCommand
from cognite_toolkit._cdf_tk.exceptions import AuthenticationError, ToolkitValueError
from cognite_toolkit._cdf_tk.tk_warnings import LowSeverityWarning
from cognite_toolkit._cdf_tk.utils import read_yaml_file, safe_write
from cognite_toolkit._cdf_tk.utils.file import safe_rmtree, yaml_safe_dump


class ImportTransformationCLI(ToolkitCommand):
    def __init__(
        self,
        get_client: Callable[[], ToolkitClient] | None = None,
        print_warning: bool = True,
        skip_tracking: bool = False,
    ):
        super().__init__(print_warning, skip_tracking)
        self._dataset_external_id_by_id: dict[int, str] = {}
        # We only initialize the client if we need to look up dataset ids.
        self._client: ToolkitClient | None = None
        self._get_client = get_client

    def execute(
        self,
        source: Path,
        destination: Path,
        overwrite: bool,
        flatten: bool,
        clean: bool,
        verbose: bool = False,
    ) -> None:
        # Manifest files are documented at
        # https://cognite-transformations-cli.readthedocs-hosted.com/en/latest/quickstart.html#transformation-manifest
        if source.is_file() and source.suffix in {".yaml", ".yml"}:
            yaml_files = [source]
        elif source.is_file():
            raise ToolkitValueError(f"File {source} is not a YAML file.")
        elif source.is_dir():
            yaml_files = list(source.rglob("*.yaml")) + list(source.rglob("*.yml"))
        else:
            raise ToolkitValueError(f"Source {source} is not a file or directory.")

        if not yaml_files:
            self.warn(LowSeverityWarning("No YAML files found in the source directory."))
            return None

        count_by_resource_type: dict[str, int] = defaultdict(int)
        for yaml_file in yaml_files:
            data = self._load_file(yaml_file)
            if data is None:
                continue

            # The convert schedule and notifications pop off the schedule and notifications
            # keys from the transformation
            schedule = self._convert_schedule(data, data["externalId"], yaml_file)
            notifications = self._convert_notifications(data, data["externalId"], yaml_file)
            transformation, source_query_relative_path = self._convert_transformation(data, yaml_file)
            source_query_path = yaml_file.parent / source_query_relative_path if source_query_relative_path else None

            if source_query_path and not source_query_path.exists():
                raise ToolkitValueError(
                    f"Query file {source_query_path.as_posix()!r} does not exist. "
                    f"This is referenced in {yaml_file.as_posix()!r}."
                )

            if flatten:
                destination_folder = destination
            else:
                destination_folder = destination / yaml_file.relative_to(source).parent
            destination_folder.mkdir(parents=True, exist_ok=True)

            destination_transformation = destination_folder / f"{yaml_file.stem}.Transformation.yaml"
            if not overwrite and destination_transformation.exists():
                self.warn(LowSeverityWarning(f"File already exists at {destination_transformation}. Skipping."))
                continue
            safe_write(destination_transformation, yaml_safe_dump(transformation))
            if source_query_path is not None:
                destination_query_path = destination_folder / f"{destination_transformation.stem}.sql"
                shutil.copyfile(source_query_path, destination_query_path)

            if schedule is not None:
                destination_schedule = destination_folder / f"{yaml_file.stem}.Schedule.yaml"
                safe_write(destination_schedule, yaml_safe_dump(schedule))
            if notifications:
                destination_notification = destination_folder / f"{yaml_file.stem}.Notification.yaml"
                safe_write(destination_notification, yaml_safe_dump(notifications))
            if verbose:
                print(f"Imported {yaml_file} to {destination_folder}.")
            count_by_resource_type["transformation"] += 1
            count_by_resource_type["schedule"] += 1 if schedule is not None else 0
            count_by_resource_type["notification"] += len(notifications)

        print(f"Finished importing from {source} to {destination}.")
        table = Table(title="Import transformation-cli Summary")
        table.add_column("Resource Type", justify="right", style="cyan")
        table.add_column("Count", justify="right", style="magenta")
        for resource_type, count in count_by_resource_type.items():
            table.add_row(resource_type, str(count))
        print(table)

        if clean:
            if source.is_dir():
                print(f"Cleaning up source directory {source}.")
                safe_rmtree(source)
            else:
                print(f"Cleaning up source file {source}.")
                source.unlink()

    def _load_file(self, yaml_file: Path) -> dict[str, Any] | None:
        content = read_yaml_file(yaml_file, expected_output="dict")
        required_keys = {"externalId", "name", "destination", "query"}
        if missing_keys := required_keys - content.keys():
            self.warn(
                LowSeverityWarning(
                    f"Missing required keys {missing_keys} in {yaml_file}. Likely not a Transformation manifest. Skipping."
                )
            )
            return None
        return content

    def _convert_transformation(
        self, transformation: dict[str, Any], source_file: Path
    ) -> tuple[dict[str, Any], Path | None]:
        if "shared" in transformation:
            transformation["isPublic"] = transformation.pop("shared")
        if "action" in transformation:
            transformation["conflictMode"] = transformation.pop("action")
        if "ignoreNullFields" not in transformation:
            # This is required by the API, but the transformation-cli sets it to true by default.
            transformation["ignoreNullFields"] = True
        if isinstance(transformation["destination"], str):
            # The API expects the destination to be a dict with a type key, while the transformation-cli
            # also allows a string.
            transformation["destination"] = {"type": transformation["destination"]}

        source_query_path: Path | None = None
        if isinstance(transformation["query"], dict):
            query = transformation.pop("query")
            if "file" in query:
                source_query_path = Path(query.pop("file"))

        if "dataSetId" in transformation:
            if "dataSetExternalId" in transformation:
                self.warn(LowSeverityWarning(f"Both dataSetId and dataSetExternalId are present in {source_file}."))
            else:
                data_set_external_id = self._lookup_dataset(transformation.pop("dataSetId"))
                if data_set_external_id is None:
                    self.warn(
                        LowSeverityWarning(f"Failed to find DataSet with id {transformation['dataSetId']} in CDF.")
                    )
                else:
                    transformation["dataSetExternalId"] = data_set_external_id

        if "authentication" in transformation:
            authentication = transformation["authentication"]
            if not isinstance(authentication, dict):
                self.warn(LowSeverityWarning(f"Invalid authentication format in {source_file}."))
            else:
                if "tokenUrl" in authentication:
                    authentication["tokenUri"] = authentication.pop("tokenUrl")

                if "read" in authentication or "write" in authentication:
                    # Read or Write in authentication.
                    transformation.pop("authentication")

                read = authentication.pop("read", None)
                if read and isinstance(read, dict):
                    if "tokenUrl" in read:
                        read["tokenUri"] = read.pop("tokenUrl")
                    transformation["sourceOidcCredentials"] = read

                write = authentication.pop("write", None)
                if write and isinstance(write, dict):
                    if "tokenUrl" in write:
                        write["tokenUri"] = write.pop("tokenUrl")
                    transformation["destinationOidcCredentials"] = write
        return transformation, source_query_path

    def _convert_notifications(
        self, transformation: dict[str, Any], external_id: str, source_file: Path
    ) -> list[dict[str, Any]]:
        notifications = []
        if "notifications" in transformation:
            notifications_raw = transformation.pop("notifications")
            if isinstance(notifications_raw, list) and all(isinstance(n, str) for n in notifications_raw):
                notifications = [
                    {"destination": email, "transformationExternalId": external_id} for email in notifications_raw
                ]
            else:
                self.warn(LowSeverityWarning(f"Invalid notifications format in {source_file}."))
        return notifications

    def _convert_schedule(
        self, transformation: dict[str, Any], external_id: str, source_file: Path
    ) -> dict[str, Any] | None:
        schedule: dict[str, Any] | None = None
        if "schedule" in transformation:
            schedule_raw = transformation.pop("schedule")
            if isinstance(schedule_raw, str):
                schedule = {"interval": schedule_raw}
            elif isinstance(schedule_raw, dict):
                schedule = schedule_raw
            else:
                self.warn(LowSeverityWarning(f"Invalid schedule format in {source_file}."))
            if isinstance(schedule, dict):
                schedule["externalId"] = external_id
        return schedule

    def _lookup_dataset(self, dataset_id: int) -> str | None:
        if dataset_id in self._dataset_external_id_by_id:
            return self._dataset_external_id_by_id[dataset_id]
        dataset = self.client.data_sets.retrieve(id=dataset_id)
        if dataset is None or dataset.external_id is None:
            return None
        self._dataset_external_id_by_id[dataset.id] = dataset.external_id
        return dataset.external_id

    @property
    def client(self) -> ToolkitClient:
        if self._client is None:
            if self._get_client is None:
                raise AuthenticationError(
                    "No Cognite Client available. Are you missing a .env file?"
                    "\nThis is required to look up dataset ids in "
                    "the transformation-cli manifest(s)."
                )
            self._client = self._get_client()
        return self._client
