from __future__ import annotations

import os
from pathlib import Path
from typing import Any
from urllib.parse import urlparse

try:
    from airflow.sdk import ObjectStoragePath
except ImportError:
    try:
        from airflow.io.path import ObjectStoragePath
    except ImportError:
        pass

from cosmos import settings
from cosmos.constants import DEFAULT_TARGET_PATH, FILE_SCHEME_AIRFLOW_DEFAULT_CONN_ID_MAP
from cosmos.exceptions import CosmosValueError


def upload_to_aws_s3(
    project_dir: str,
    bucket_name: str,
    aws_conn_id: str | None = None,
    source_subpath: str = DEFAULT_TARGET_PATH,
    **kwargs: Any,
) -> None:
    """
    Helper function demonstrating how to upload files to AWS S3 that can be used as a callback.

    :param project_dir: Path of the cloned project directory which Cosmos tasks work from.
    :param bucket_name: Name of the S3 bucket to upload to.
    :param aws_conn_id: AWS connection ID to use when uploading files.
    :param source_subpath: Path of the source directory sub-path to upload files from.
    """
    from airflow.providers.amazon.aws.hooks.s3 import S3Hook

    target_dir = f"{project_dir}/{source_subpath}"
    aws_conn_id = aws_conn_id if aws_conn_id else S3Hook.default_conn_name
    hook = S3Hook(aws_conn_id=aws_conn_id)
    context = kwargs["context"]

    # Airflow 3 and Airflow 2 compatibility, respectively:
    try_number = getattr(context["task_instance"], "try_number") or getattr(context["task_instance"], "_try_number")

    # Iterate over the files in the target dir and upload them to S3
    for dirpath, _, filenames in os.walk(target_dir):
        for filename in filenames:
            s3_key = (
                f"{context['dag'].dag_id}"
                f"/{context['run_id']}"
                f"/{context['task_instance'].task_id}"
                f"/{try_number}"
                f"{dirpath.split(project_dir)[-1]}/{filename}"
            )
            hook.load_file(
                filename=f"{dirpath}/{filename}",
                bucket_name=bucket_name,
                key=s3_key,
                replace=True,
            )


def upload_to_gcp_gs(
    project_dir: str,
    bucket_name: str,
    gcp_conn_id: str | None = None,
    source_subpath: str = DEFAULT_TARGET_PATH,
    **kwargs: Any,
) -> None:
    """
    Helper function demonstrating how to upload files to GCP GS that can be used as a callback.

    :param project_dir: Path of the cloned project directory which Cosmos tasks work from.
    :param bucket_name: Name of the GCP GS bucket to upload to.
    :param gcp_conn_id: GCP connection ID to use when uploading files.
    :param source_subpath: Path of the source directory sub-path to upload files from.
    """
    from airflow.providers.google.cloud.hooks.gcs import GCSHook

    target_dir = f"{project_dir}/{source_subpath}"
    gcp_conn_id = gcp_conn_id if gcp_conn_id else GCSHook.default_conn_name
    # bucket_name = kwargs["bucket_name"]
    hook = GCSHook(gcp_conn_id=gcp_conn_id)
    context = kwargs["context"]

    # Iterate over the files in the target dir and upload them to GCP GS
    for dirpath, _, filenames in os.walk(target_dir):
        for filename in filenames:
            object_name = (
                f"{context['dag'].dag_id}"
                f"/{context['run_id']}"
                f"/{context['task_instance'].task_id}"
                f"/{context['task_instance']._try_number}"
                f"{dirpath.split(project_dir)[-1]}/{filename}"
            )
            hook.upload(
                filename=f"{dirpath}/{filename}",
                bucket_name=bucket_name,
                object_name=object_name,
            )


def upload_to_azure_wasb(
    project_dir: str,
    container_name: str,
    azure_conn_id: str | None = None,
    source_subpath: str = DEFAULT_TARGET_PATH,
    **kwargs: Any,
) -> None:
    """
    Helper function demonstrating how to upload files to Azure WASB that can be used as a callback.

    :param project_dir: Path of the cloned project directory which Cosmos tasks work from.
    :param container_name: Name of the Azure WASB container to upload files to.
    :param azure_conn_id: Azure connection ID to use when uploading files.
    :param source_subpath: Path of the source directory sub-path to upload files from.
    """
    from airflow.providers.microsoft.azure.hooks.wasb import WasbHook

    target_dir = f"{project_dir}/{source_subpath}"
    azure_conn_id = azure_conn_id if azure_conn_id else WasbHook.default_conn_name
    # container_name = kwargs["container_name"]
    hook = WasbHook(wasb_conn_id=azure_conn_id)
    context = kwargs["context"]

    # Iterate over the files in the target dir and upload them to WASB container
    for dirpath, _, filenames in os.walk(target_dir):
        for filename in filenames:
            blob_name = (
                f"{context['dag'].dag_id}"
                f"/{context['run_id']}"
                f"/{context['task_instance'].task_id}"
                f"/{context['task_instance']._try_number}"
                f"{dirpath.split(project_dir)[-1]}/{filename}"
            )
            hook.load_file(
                file_path=f"{dirpath}/{filename}",
                container_name=container_name,
                blob_name=blob_name,
                overwrite=True,
            )


def _configure_remote_target_path() -> tuple[Path | ObjectStoragePath, str] | tuple[None, None]:
    """Configure the remote target path if it is provided."""

    if not settings.remote_target_path:
        return None, None

    _configured_target_path = None

    target_path_str = str(settings.remote_target_path)

    remote_conn_id = settings.remote_target_path_conn_id
    if not remote_conn_id:
        target_path_schema = urlparse(target_path_str).scheme
        remote_conn_id = FILE_SCHEME_AIRFLOW_DEFAULT_CONN_ID_MAP.get(target_path_schema, None)  # type: ignore[assignment]
    if remote_conn_id is None:
        return None, None

    _configured_target_path = ObjectStoragePath(target_path_str, conn_id=remote_conn_id)

    if not _configured_target_path.exists():  # type: ignore[no-untyped-call]
        _configured_target_path.mkdir(parents=True, exist_ok=True)

    return _configured_target_path, remote_conn_id


def _construct_dest_file_path(
    dest_root: Any | None,
    rel_path: str,
    *prefixes: str,
) -> str:
    """
    Helper function for constructing normalized destination file paths for remote object storage uploads.

    This function provides a consistent way to join multiple path segments — such as task-specific identifiers,
    subpaths, and relative file locations — into a single, normalized path string. It is typically used alongside
    helpers like ``upload_to_cloud_storage`` to generate fully qualified remote storage paths.

    :param dest_root: Optional root prefix (e.g., S3 bucket URI or base path).
    :param rel_path: Relative file path (e.g., "target/run_results.json").
    :param prefixes: Optional prefix segments (e.g., DAG ID, run ID, task ID) that are prepended to the path.
    :return: A normalized string representing the destination path for uploading or referencing in object storage.
    """
    clean_prefixes = [p.strip("/") for p in prefixes if p]
    clean_rel = rel_path.lstrip("/")

    segments: list[str] = [*clean_prefixes, clean_rel]
    if dest_root is None:
        return "/".join(segments)

    root = str(dest_root).rstrip("/")
    if not segments:
        return root

    return f"{root}/{'/'.join(segments)}"


def upload_to_cloud_storage(project_dir: str, source_subpath: str = DEFAULT_TARGET_PATH, **kwargs: Any) -> None:
    """
    Helper function demonstrating how to upload files to remote object stores that can be used as a callback. This is
    an example of a helper function that can be used if on Airflow >= 2.8 and cosmos configurations like
    ``remote_target_path`` and ``remote_target_path_conn_id`` when set can be leveraged.

    :param project_dir: Path of the cloned project directory which Cosmos tasks work from.
    :param source_subpath: Path of the source directory sub-path to upload files from.
    """
    dest_target_dir, dest_conn_id = _configure_remote_target_path()

    if not dest_target_dir:
        raise CosmosValueError("You're trying to upload artifact files, but the remote target path is not configured.")

    source_target_dir = Path(project_dir) / f"{source_subpath}"
    files = [str(file) for file in source_target_dir.rglob("*") if file.is_file()]

    context = kwargs["context"]
    task_run_identifier = (
        f"{context['dag'].dag_id}"
        f"/{context['run_id']}"
        f"/{context['task_instance'].task_id}"
        f"/{context['task_instance'].try_number}"
    )

    for file_path in files:
        rel_path = os.path.relpath(file_path, source_target_dir)
        dest_file_path = _construct_dest_file_path(
            dest_target_dir,
            rel_path,
            task_run_identifier,
            source_subpath,
        )
        dest_object_storage_path = ObjectStoragePath(dest_file_path, conn_id=dest_conn_id)
        ObjectStoragePath(file_path).copy(dest_object_storage_path)
