# SPDX-License-Identifier: PROPRIETARY
# SPDX-FileCopyrightText: Copyright The Geneva Authors

import contextlib
import logging
import math
import os
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from typing import TYPE_CHECKING, Any, Optional

import attrs
import emoji
import lance
import pyarrow.fs as fs
from lance_namespace import DescribeTableRequest
from lance_namespace import connect as namespace_connect

from geneva import DEFAULT_UPLOAD_DIR
from geneva.config import ConfigBase
from geneva.tqdm import tqdm

if TYPE_CHECKING:
    from azure.storage.blob import BlobServiceClient

_LOG = logging.getLogger(__name__)
_DEFAULT_GCS_UPLOAD_TIMEOUT_S = 600.0
_DEFAULT_GCS_CHUNK_SIZE_MB = 32


def make_upload_path(filename: str) -> str:
    """Make the remote path for uploading a file via namespace file session.

    When using ds.new_file_session(), the session is rooted at the dataset
    directory. This function constructs the path including the _geneva_uploads
    subdirectory.
    """
    return f"{DEFAULT_UPLOAD_DIR}/{filename}"


# Suppress verbose Azure SDK HTTP logging
logging.getLogger("azure.core.pipeline.policies.http_logging_policy").setLevel(
    logging.WARNING
)


def _get_gcs_upload_timeout_s() -> float:
    raw = os.getenv("GENEVA_GCS_UPLOAD_TIMEOUT_S")
    if raw is None:
        return _DEFAULT_GCS_UPLOAD_TIMEOUT_S
    try:
        return float(raw)
    except ValueError:
        _LOG.warning(
            "Invalid GENEVA_GCS_UPLOAD_TIMEOUT_S=%r; using %.1fs",
            raw,
            _DEFAULT_GCS_UPLOAD_TIMEOUT_S,
        )
        return _DEFAULT_GCS_UPLOAD_TIMEOUT_S


def _get_gcs_chunk_size_bytes() -> int:
    raw = os.getenv("GENEVA_GCS_UPLOAD_CHUNK_SIZE_MB")
    if raw is None:
        return _DEFAULT_GCS_CHUNK_SIZE_MB * 1024 * 1024
    try:
        value = int(raw)
    except ValueError:
        _LOG.warning(
            "Invalid GENEVA_GCS_UPLOAD_CHUNK_SIZE_MB=%r; using %s MB",
            raw,
            _DEFAULT_GCS_CHUNK_SIZE_MB,
        )
        value = _DEFAULT_GCS_CHUNK_SIZE_MB
    return max(value, 1) * 1024 * 1024


@attrs.define
class Uploader(ConfigBase):
    """
    This class is used to upload files to a specified directory.

    Supports uploads via:
    - GCS optimized path (google.cloud.storage SDK) for gs:// URLs
    - Azure optimized path (azure-storage-blob SDK) for az:// URLs
    - Dataset file session (with multi-part upload) when namespace is available
    - PyArrow FileSystem as fallback

    Upload directories are automatically derived to be table-specific:
    - For namespace tables: {table_location}/_geneva_uploads
    - For local tables: {db_uri}/{table_name}.lance/_geneva_uploads
    """

    upload_dir: Optional[str] = attrs.field(default=None)
    namespace_impl: Optional[str] = attrs.field(default=None)
    namespace_properties: Optional[dict[str, str]] = attrs.field(default=None)
    table_id: Optional[list[str]] = attrs.field(default=None)
    db_uri: Optional[str] = attrs.field(default=None)

    # Cached namespace resources (initialized lazily in __attrs_post_init__)
    _file_session: Any = attrs.field(default=None, init=False, repr=False)

    def __attrs_post_init__(self) -> None:
        # Derive upload_dir to be table-specific
        if self.namespace_impl and self.namespace_properties and self.table_id:
            # Namespace table: get location from namespace and initialize file session
            namespace_client = namespace_connect(
                self.namespace_impl, self.namespace_properties
            )
            response = namespace_client.describe_table(
                DescribeTableRequest(id=self.table_id)
            )
            if response.location is None:
                raise ValueError(f"Table location is None for table {self.table_id}")

            # Set upload_dir to table-specific location
            self.upload_dir = f"{response.location.rstrip('/')}/{DEFAULT_UPLOAD_DIR}"
            _LOG.info(f"Derived namespace table upload_dir: {self.upload_dir}")

            # Initialize file session for reuse
            ds = lance.dataset(namespace=namespace_client, table_id=self.table_id)
            self._file_session = ds.new_file_session()
        elif self.db_uri and self.table_id:
            # Local table: construct path from db_uri and table_name
            # Handle URIs with query strings (e.g. s3+ddb://bucket/path?ddbTableName=x)
            from urllib.parse import urlparse, urlunparse

            table_name = self.table_id[-1] if self.table_id else ""
            parsed = urlparse(self.db_uri)
            path = parsed.path.rstrip("/")
            base_path = f"{path}/{table_name}.lance/{DEFAULT_UPLOAD_DIR}"
            if parsed.scheme:
                self.upload_dir = urlunparse(parsed._replace(path=base_path))
            else:
                self.upload_dir = base_path
            _LOG.info(f"Derived local table upload_dir: {self.upload_dir}")
        elif self.upload_dir:
            # Remove trailing slash if present
            self.upload_dir = self.upload_dir.removesuffix("/")
        else:
            raise ValueError(
                "Either upload_dir must be provided, or table context "
                "(table_id + db_uri or namespace credentials) must be provided"
            )

    @classmethod
    def name(cls) -> str:
        return "uploader"

    @property
    def fs_and_path(self) -> tuple[fs.FileSystem, str]:
        assert self.upload_dir is not None, "upload_dir must be set"
        # Note: az:// URIs are handled separately via Azure SDK, not PyArrow
        return fs.FileSystem.from_uri(self.upload_dir)

    def _make_upload_path(self, filename: str) -> str:
        """Make the remote path for uploading a file via namespace file session."""
        return make_upload_path(filename)

    def _upload_gcs(self, f: Path) -> str:
        """
        Upload to GCS -- dispatch here when we detect the dir is GCS
        Because the google client is much more performant than pyarrow
        """
        from google.cloud import storage

        # we don't call this frequently so just create client on the fly
        storage_client = storage.Client()

        assert self.upload_dir is not None, "upload_dir must be set"
        path = self.upload_dir.removeprefix("gs://")
        bucket_name, destination_blob_prefix = path.split("/", 1)
        destination_blob_name = f"{destination_blob_prefix}/{f.name}"
        bucket = storage_client.bucket(bucket_name)
        blob = bucket.blob(destination_blob_name)
        upload_timeout_s = _get_gcs_upload_timeout_s()
        blob.chunk_size = _get_gcs_chunk_size_bytes()

        # determine upload strategy
        # if the file is small enough, we can just upload it directly
        # otherwise we will upload in chunks
        # and then compose the chunks together
        # this lets us maximize upload bandwidth for large files
        if f.stat().st_size < 1024 * 1024 * 1024:  # 1GiB
            blob.upload_from_filename(f.as_posix(), timeout=math.ceil(upload_timeout_s))
        else:
            # we will upload in 32 chunks, which is the maximum number of
            # chunks allowed for composing blobs
            chunk_size = -(-f.stat().st_size // 32)

            pbar = tqdm(
                total=f.stat().st_size, unit="B", unit_scale=True, unit_divisor=1024
            )
            pbar.set_description(
                emoji.emojize(f":cloud: uploading {f.name:5.5} to {self.upload_dir}")
            )

            def _upload_part(idx: int) -> storage.Blob:
                start = idx * chunk_size
                end = min(start + chunk_size, f.stat().st_size)

                length = end - start
                part = bucket.blob(f"{destination_blob_name}-{idx}")
                with (
                    part.open("wb", timeout=upload_timeout_s) as f_out,
                    open(f, "rb") as f_in,
                ):
                    f_in.seek(start)
                    while length:
                        read_size = min(length, 1024 * 64)
                        data = f_in.read(read_size)
                        f_out.write(data)  # type: ignore[arg-type]
                        length -= read_size
                        pbar.update(read_size)
                return part

            # int_divceil
            num_chunks = -(-f.stat().st_size // chunk_size)
            with ThreadPoolExecutor(max_workers=8) as executor:
                parts = executor.map(_upload_part, range(num_chunks))

            pbar.close()
            blob.compose(list(parts))

        return f"gs://{bucket_name}/{destination_blob_name}"

    def _get_azure_blob_client(self) -> "BlobServiceClient":
        """
        Get Azure BlobServiceClient with appropriate credentials.

        Supports:
        - AZURE_STORAGE_ACCOUNT_KEY: Account key auth
        - AZURE_STORAGE_TOKEN: Bearer token auth (from az account get-access-token)
        - DefaultAzureCredential: Azure CLI, managed identity, workload identity, etc.
        """
        from azure.core.credentials import AccessToken
        from azure.storage.blob import BlobServiceClient

        from geneva.utils.storage import get_azure_storage_account

        account_name = get_azure_storage_account()

        account_url = f"https://{account_name}.blob.core.windows.net"

        # Try account key first (SharedKey auth)
        account_key = os.environ.get("AZURE_STORAGE_ACCOUNT_KEY")
        if account_key:
            return BlobServiceClient(account_url=account_url, credential=account_key)

        # Try bearer token (from az account get-access-token)
        # Must be wrapped in a TokenCredential-like object
        token = os.environ.get("AZURE_STORAGE_TOKEN")
        if token:

            class _StaticTokenCredential:
                """Simple credential that returns a static token."""

                def __init__(self, token: str) -> None:
                    self._token = token

                def get_token(self, *scopes, **kwargs) -> Any:
                    # Token expires in 1 hour, but we don't track expiry here
                    return AccessToken(self._token, 0)

            return BlobServiceClient(
                account_url=account_url, credential=_StaticTokenCredential(token)
            )

        # Fall back to DefaultAzureCredential
        from azure.identity import DefaultAzureCredential

        credential = DefaultAzureCredential()
        return BlobServiceClient(account_url=account_url, credential=credential)

    def _parse_azure_uri(self) -> tuple[str, str]:
        """Parse az://container/path into (container, blob_prefix)."""
        assert self.upload_dir is not None
        path = self.upload_dir.removeprefix("az://")
        if "/" in path:
            container, blob_prefix = path.split("/", 1)
        else:
            container, blob_prefix = path, ""
        return container, blob_prefix

    def _upload_azure(self, f: Path) -> str:
        """
        Upload to Azure Blob Storage using azure-storage-blob SDK.

        Similar to GCS optimized path - uses native SDK for better performance
        and credential handling.
        """
        blob_service = self._get_azure_blob_client()
        container_name, blob_prefix = self._parse_azure_uri()
        blob_name = f"{blob_prefix}/{f.name}" if blob_prefix else f.name

        container_client = blob_service.get_container_client(container_name)
        blob_client = container_client.get_blob_client(blob_name)

        with (
            tqdm(
                total=f.stat().st_size, unit="B", unit_scale=True, unit_divisor=1024
            ) as pbar,
            f.open("rb") as f_in,
        ):
            pbar.set_description(
                emoji.emojize(f":cloud: uploading {f.name} to {self.upload_dir}")
            )
            blob_client.upload_blob(f_in, overwrite=True)
            pbar.update(f.stat().st_size)

        return f"{self.upload_dir}/{f.name}"

    def _file_exists(self, f: Path) -> bool:
        try:
            # Use GCS SDK for GCS
            if self.upload_dir and self.upload_dir.startswith("gs://"):
                from google.cloud import storage

                storage_client = storage.Client()
                path = self.upload_dir.removeprefix("gs://")
                bucket_name, prefix = path.split("/", 1)
                bucket = storage_client.bucket(bucket_name)
                blob = bucket.blob(f"{prefix}/{f.name}")
                return blob.exists()

            # Use Azure SDK for Azure
            if self.upload_dir and self.upload_dir.startswith("az://"):
                blob_service = self._get_azure_blob_client()
                container_name, blob_prefix = self._parse_azure_uri()
                blob_name = f"{blob_prefix}/{f.name}" if blob_prefix else f.name
                blob_client = blob_service.get_blob_client(container_name, blob_name)
                return blob_client.exists()

            # Use cached file session when namespace is available
            if self._file_session is not None:
                return self._file_session.contains(self._make_upload_path(f.name))

            # Fallback to PyArrow FileSystem
            handle, prefix = self.fs_and_path
            upload_path = str(Path(prefix) / f.name)
            return handle.get_file_info(upload_path).type == fs.FileType.File
        except Exception:
            _LOG.exception(f"Failed to check if file exists: {f.name}")
            return False

    def _upload_lance_session(self, f: Path) -> str:
        """
        Upload using dataset's file session (with automatic multi-part upload support).

        Uses the cached file session initialized in __attrs_post_init__ which
        inherits the dataset's storage configuration.
        """
        assert self._file_session is not None, "file_session must be initialized"
        assert self.upload_dir is not None, "upload_dir must be set"

        remote_path = self._make_upload_path(f.name)

        # Upload with progress bar
        with tqdm(
            total=f.stat().st_size, unit="B", unit_scale=True, unit_divisor=1024
        ) as pbar:
            pbar.set_description(
                emoji.emojize(f":cloud: uploading {f.name} to {self.upload_dir}")
            )

            # Upload the file - multi-part upload happens automatically for files > 5MB
            self._file_session.upload_file(str(f), remote_path)
            pbar.update(f.stat().st_size)
        return f"{self.upload_dir}/{f.name}"

    def upload(self, f: Path) -> str:
        """
        Upload a file to the specified directory.

        The name of the object will be in the form of
        <path_to_upload_dir>/<name_of_file>
        """
        if self._file_exists(f):
            _LOG.debug(
                f"File {f.name} already exists in {self.upload_dir}, skipping upload"
            )
            return f"{self.upload_dir}/{f.name}"

        # GCS optimized path using google.cloud.storage SDK
        assert self.upload_dir is not None, "upload_dir must be set"
        if self.upload_dir.startswith("gs://"):
            return self._upload_gcs(f)

        # Azure optimized path using azure-storage-blob SDK
        if self.upload_dir.startswith("az://"):
            return self._upload_azure(f)

        # Use cached file session for non-GCS/Azure uploads when namespace available
        # This provides multi-part upload and proper credential management
        if self._file_session is not None:
            return self._upload_lance_session(f)

        # Fallback to PyArrow FileSystem
        handle, prefix = self.fs_and_path
        upload_path = str(Path(prefix) / f.name)

        # Create upload directory if it doesn't exist
        with contextlib.suppress(Exception):
            handle.create_dir(prefix, recursive=True)

        with (
            handle.open_output_stream(upload_path, buffer_size=1024 * 1024 * 64) as out,
            tqdm(
                total=f.stat().st_size, unit="B", unit_scale=True, unit_divisor=1024
            ) as pbar,
            f.open("rb") as f_in,
        ):
            pbar.set_description(
                emoji.emojize(f":cloud: uploading {f.name} to {self.upload_dir}")
            )
            chunk_size = 1024 * 1024 * 8  # 8MiB chunks

            while data := f_in.read(chunk_size):
                out.write(data)
                pbar.update(len(data))

        return f"{self.upload_dir}/{f.name}"
