from collections import Counter
from collections.abc import Sequence
from functools import partial
from pathlib import Path

from rich.console import Console
from rich.markup import escape

from cognite_toolkit._cdf_tk.client import ToolkitClient
from cognite_toolkit._cdf_tk.client.http_client import (
    HTTPClient,
    ItemsFailedRequest,
    ItemsFailedResponse,
    ItemsResultMessage,
    ItemsSuccessResponse,
)
from cognite_toolkit._cdf_tk.client.resource_classes.data_modeling import ViewReference
from cognite_toolkit._cdf_tk.constants import DATA_MANIFEST_SUFFIX, DATA_RESOURCE_DIR
from cognite_toolkit._cdf_tk.cruds import ViewCRUD
from cognite_toolkit._cdf_tk.exceptions import ToolkitRuntimeError, ToolkitValueError
from cognite_toolkit._cdf_tk.protocols import T_ResourceRequest, T_ResourceResponse
from cognite_toolkit._cdf_tk.storageio import (
    FileContentIO,
    T_Selector,
    UploadableStorageIO,
    get_upload_io,
)
from cognite_toolkit._cdf_tk.storageio._base import TableUploadableStorageIO, UploadItem
from cognite_toolkit._cdf_tk.storageio.selectors import Selector, load_selector
from cognite_toolkit._cdf_tk.storageio.selectors._instances import InstanceSpaceSelector, InstanceViewSelector
from cognite_toolkit._cdf_tk.tk_warnings import HighSeverityWarning, MediumSeverityWarning, ToolkitWarning
from cognite_toolkit._cdf_tk.utils.auth import EnvironmentVariables
from cognite_toolkit._cdf_tk.utils.fileio import MultiFileReader
from cognite_toolkit._cdf_tk.utils.producer_worker import ProducerWorkerExecutor
from cognite_toolkit._cdf_tk.utils.progress_tracker import ProgressTracker
from cognite_toolkit._cdf_tk.utils.useful_types import JsonVal

from ._base import ToolkitCommand
from .deploy import DeployCommand


class UploadCommand(ToolkitCommand):
    """Command for uploading data to CDF from files in a specified directory.

    This command reads files matching a specific pattern in the input directory,
    processes them using a StorageIO instance, and uploads the data to CDF.

    Attributes:
        _MAX_QUEUE_SIZE (int): The maximum size of the queue for processing items.
            Set to 80 to balance memory usage and processing speed.
    """

    _MAX_QUEUE_SIZE = 80
    _MAX_VERBOSE_PRINTED_FAILED_IDS = 10
    _UPLOAD = "upload"

    def upload(
        self,
        input_dir: Path,
        client: ToolkitClient,
        deploy_resources: bool,
        dry_run: bool,
        verbose: bool,
        kind: str | None = None,
    ) -> None:
        """Uploads data from files in the specified input directory to CDF.

        Args:
            input_dir: The directory containing the files to upload. It is expected to
                have the structure defined below.
            client: An instance of ToolkitClient to interact with CDF.
            deploy_resources: If True, deploys resources from the 'resources' subdirectory.
            dry_run: If True, performs a dry run without actually uploading the data
                (or deploying resources).
            verbose: If True, prints detailed information about the upload process.
            kind: Optional; if provided, only data files of this kind will be processed.

        The expected structure of the input directory is as follows:
        ```
        input_dir/
        ├── resources/                # Optional, only if deploy_resources is True
        │   ├── raw/
        │   │   ├── table1.Table.yaml
        │   │   └── table2.Table.yaml
        │   └── ...
        ├── datafile1.kind.ndjson # Data file of a specific kind
        ├── datafile1.Manifest.yaml       # Manifest for datafile1
        ├── datafile2.kind2.ndjson # Another data file of the same or different kind
        ├── datafile2.Manifest.yaml       # Manifest file for datafile2
        └── ...
        """
        console = client.console
        data_files_by_selector = self._find_data_files(input_dir)

        self._deploy_resource_folder(input_dir / DATA_RESOURCE_DIR, deploy_resources, client, console, dry_run, verbose)

        data_files_by_selector = self._topological_sort_if_instance_selector(data_files_by_selector, client)
        self._upload_data(data_files_by_selector, client, dry_run, input_dir, console, verbose)

    def _topological_sort_if_instance_selector(
        self, data_files_by_selector: dict[Selector, list[Path]], client: ToolkitClient
    ) -> dict[Selector, list[Path]]:
        """Topologically sorts InstanceSelectors (if they are present) to determine the order of upload based on container dependencies from the views.

        Args:
            data_files_by_selector: A dictionary mapping selectors to their data files.
            client: The cognite client to use for the upload.

        Returns:
            A dictionary mapping selectors to their data files with necessary preprocessing.
        """
        counts = Counter(type(selector) for selector in data_files_by_selector.keys())
        if (counts[InstanceSpaceSelector] + counts[InstanceViewSelector]) <= 1:
            return data_files_by_selector

        selector_by_view_id: dict[ViewReference, Selector] = {}
        for selector in data_files_by_selector:
            if isinstance(selector, InstanceSpaceSelector | InstanceViewSelector) and selector.view is not None:
                view_ref = selector.view.as_id()
                if not isinstance(view_ref, ViewReference):
                    raise RuntimeError(
                        f"View {view_ref} does not have a version, which is required for topological sorting."
                    )
                selector_by_view_id[view_ref] = selector

        view_dependencies, cyclic_views = ViewCRUD.create_loader(client).topological_sort_container_constraints(
            list(selector_by_view_id.keys())
        )
        if cyclic_views:
            MediumSeverityWarning(
                f"Cyclic dependencies detected between views: {', '.join(str(view) for view in cyclic_views)}. "
                f"If their constraints are not already satisfied by instances in existing containers / views "
                f"in CDF, the upload may fail for these views."
            ).print_warning(console=client.console)
            view_dependencies = view_dependencies + cyclic_views

        prepared_selectors: dict[Selector, list[Path]] = {}

        # Reorder selectors according to the dependency-sorted view list
        for view_id in view_dependencies:
            selector = selector_by_view_id[view_id]
            prepared_selectors[selector] = data_files_by_selector[selector]

        # Preserve selectors that aren't affected by view dependencies
        # (e.g., raw tables, time series, non-view instance data)
        for selector in data_files_by_selector.keys():
            if selector not in prepared_selectors:
                prepared_selectors[selector] = data_files_by_selector[selector]

        return prepared_selectors

    def _find_data_files(
        self,
        input_dir: Path,
    ) -> dict[Selector, list[Path]]:
        """Finds data files and their corresponding metadata files in the input directory."""
        data_files_by_metadata: dict[Selector, list[Path]] = {}
        for manifest_file in input_dir.glob(f"*{DATA_MANIFEST_SUFFIX}"):
            selector_or_warning = load_selector(manifest_file)
            if isinstance(selector_or_warning, ToolkitWarning):
                self.warn(selector_or_warning)
                continue
            selector: Selector = selector_or_warning
            data_files = selector.find_data_files(input_dir, manifest_file)
            if not data_files:
                self.warn(
                    MediumSeverityWarning(
                        f"Metadata file {manifest_file.as_posix()!r} has no corresponding data files, skipping.",
                    )
                )
                continue
            data_files_by_metadata[selector] = data_files
        return data_files_by_metadata

    def _deploy_resource_folder(
        self,
        resource_dir: Path,
        deploy_resources: bool,
        client: ToolkitClient,
        console: Console,
        dry_run: bool,
        verbose: bool,
    ) -> None:
        """Deploy resources from the specified resource directory if it exists and deployment is enabled."""
        if deploy_resources and resource_dir.exists():
            deploy_command = DeployCommand()
            deploy_results = deploy_command.deploy_all_resources(
                client=client,
                build_dir=resource_dir,
                env_vars=EnvironmentVariables.create_from_environment(),
                dry_run=dry_run,
                verbose=verbose,
            )
            self.warning_list.extend(deploy_command.warning_list)
            if deploy_results.has_counts:
                console.print(deploy_results.counts_table())
        elif deploy_resources:
            self.warn(
                MediumSeverityWarning(
                    f"Resource directory {resource_dir.as_posix()!r} does not exist, skipping resource deployment."
                )
            )

    def _upload_data(
        self,
        data_files_by_selector: dict[Selector, list[Path]],
        client: ToolkitClient,
        dry_run: bool,
        input_dir: Path,
        console: Console,
        verbose: bool,
    ) -> None:
        total_file_count = sum(len(files) for files in data_files_by_selector.values())
        if verbose:
            input_dir_display = self._path_as_display_name(input_dir)
            console.print(f"Found {total_file_count} files to upload in {input_dir_display.as_posix()!r}.")
        action = "Would upload" if dry_run else "Uploading"
        with HTTPClient(config=client.config) as upload_client:
            file_count = 1
            for selector, datafiles in data_files_by_selector.items():
                io = self._create_selected_io(selector, datafiles[0], client)
                if io is None:
                    continue
                reader = MultiFileReader(datafiles)
                # FileContentIO supports uploading any file format.
                if reader.is_table and not isinstance(io, TableUploadableStorageIO | FileContentIO):
                    raise ToolkitValueError(f"{selector.display_name} does not support {reader.format!r} files.")

                item_count = io.count_items(reader, selector)
                iteration_count = item_count // io.CHUNK_SIZE + (1 if item_count % io.CHUNK_SIZE > 0 else 0)

                tracker = ProgressTracker[str]([self._UPLOAD])
                executor = ProducerWorkerExecutor[list[tuple[str, dict[str, JsonVal]]], Sequence[UploadItem]](
                    download_iterable=io.read_chunks(reader, selector),
                    process=partial(io.rows_to_data, selector=selector)
                    if reader.is_table and isinstance(io, TableUploadableStorageIO)
                    else io.json_chunk_to_data,
                    write=partial(
                        self._upload_items,
                        upload_client=upload_client,
                        io=io,
                        dry_run=dry_run,
                        selector=selector,
                        tracker=tracker,
                        console=console,
                        verbose=verbose,
                    ),
                    iteration_count=iteration_count,
                    max_queue_size=self._MAX_QUEUE_SIZE,
                    download_description=f"Reading {selector.display_name!r} files",
                    process_description="Processing",
                    write_description=f"{action} {selector.display_name!r}",
                    console=console,
                )
                executor.run()
                file_count += len(datafiles)
                executor.raise_on_error()
                final_action = "Uploaded" if not dry_run else "Would upload"
                suffix = " successfully" if not dry_run else ""
                results = tracker.aggregate()
                success = results.get((self._UPLOAD, "success"), 0)
                failed = results.get((self._UPLOAD, "failed"), 0)
                if failed > 0:
                    suffix += f", {failed:,} failed"
                console.print(
                    f"{final_action} {success:,} {selector.display_name} from {len(datafiles)} files{suffix}."
                )

    @staticmethod
    def _path_as_display_name(input_path: Path, cwd: Path = Path.cwd()) -> Path:
        display_name = input_path
        if input_path.is_relative_to(cwd):
            display_name = input_path.relative_to(cwd)
        return display_name

    def _create_selected_io(
        self, selector: Selector, data_file: Path, client: ToolkitClient
    ) -> UploadableStorageIO | None:
        try:
            io_cls = get_upload_io(selector)
        except ValueError as e:
            self.warn(HighSeverityWarning(f"Could not find StorageIO for selector {selector}: {e}"))
            return None
        return io_cls(client)

    @classmethod
    def _upload_items(
        cls,
        data_chunk: Sequence[UploadItem[T_ResourceRequest]],
        upload_client: HTTPClient,
        io: UploadableStorageIO[T_Selector, T_ResourceResponse, T_ResourceRequest],
        selector: T_Selector,
        dry_run: bool,
        tracker: ProgressTracker[str],
        console: Console,
        verbose: bool,
    ) -> None:
        if dry_run:
            for item in data_chunk:
                tracker.set_progress(item.source_id, cls._UPLOAD, "success")
            return
        results = io.upload_items(data_chunk, upload_client, selector)
        all_failed = True
        # Group failed ids by error description for cleaner output
        failures_by_error: dict[str, list[str]] = {}
        for message in results:
            if isinstance(message, ItemsSuccessResponse):
                all_failed = False
                for id_ in message.ids:
                    tracker.set_progress(id_, step=cls._UPLOAD, status="success")
            elif isinstance(message, ItemsResultMessage):
                for id_ in message.ids:
                    tracker.set_progress(id_, step=cls._UPLOAD, status="failed")
                if verbose:
                    if isinstance(message, ItemsFailedResponse):
                        error_description = f"(HTTP {message.status_code}): {escape(message.error.message)}"
                        failures_by_error.setdefault(error_description, []).extend(message.ids)
                    elif isinstance(message, ItemsFailedRequest):
                        failures_by_error.setdefault(escape(message.error_message), []).extend(message.ids)
            else:
                console.log(f"[red]Unexpected result from upload: {escape(str(message))}[/red]")
        if verbose:
            for error_description, failed_ids in failures_by_error.items():
                ids_display = escape(", ".join(failed_ids[: cls._MAX_VERBOSE_PRINTED_FAILED_IDS]))
                if len(failed_ids) > cls._MAX_VERBOSE_PRINTED_FAILED_IDS:
                    ids_display += f" ... and {len(failed_ids) - cls._MAX_VERBOSE_PRINTED_FAILED_IDS} more"
                console.print(
                    f"[red]Failed to upload[/red] {len(failed_ids)} items from [bold]{selector}[/bold] "
                    f"{error_description}\n"
                    f"  Failed items: {ids_display}"
                )
        if all_failed and results:
            raise ToolkitRuntimeError(
                "Upload process was stopped due to repeatedly failed uploads. Rerun with --verbose to see detailed failure information."
            )
