import math
import os
from collections.abc import Generator, Iterable, Iterator
from dataclasses import asdict, dataclass
from functools import cache
from itertools import islice
from typing import Any, BinaryIO, Protocol, TypeVar
from urllib.parse import urlparse

import httpx
from fsspec import AbstractFileSystem, url_to_fs
from fsspec.implementations.http import HTTPFileSystem
from owsclient import ImpersonationOwsClient, M2MTokenManager, OwsClient

_DEFAULT_HTTPX_TIMEOUT_SECONDS = 60.0
_ASSET_UPLOAD_PRESIGNED_URL_EXPIRES_IN_SECONDS = 60 * 15  # OWS Assets limit
_KIB = 1024
_MIB = 1024 * _KIB
_PART_SIZE_BYTES = 64 * _MIB  # must be >=5 MiB except final part
_MAX_NUM_PARTS = 10_000  # S3 limit
_MAX_PRESIGNED_URL_BATCH_SIZE = 20  # OWS Assets limit
_MAX_SOURCE_FILENAME_LENGTH = 64  # OWS Assets limit
_NUM_HTTP_RETRIES = 5
_VALID_ENVIRONMENTS = {"qa", "prod"}
_ORCHARD_USER_ID = "oa:179"
_env: str | None = None
_m2m_token_manager: M2MTokenManager | None = None


def set_env(env: str) -> None:
    """Set the environment for ows-assets uploads.

    Must be called before using upload().

    Args:
        env: Environment name, must be 'qa' or 'prod'

    Raises:
        ValueError: If env is not a valid environment

    """
    global _env

    if _env is not None:
        raise RuntimeError(
            f"Environment already set to {_env}. It can only be set once."
        )

    if env not in _VALID_ENVIRONMENTS:
        raise ValueError(
            f"Invalid environment: {env}. Must be one of {_VALID_ENVIRONMENTS}"
        )
    _env = env


def get_env() -> str:
    """Get the currently configured environment for ows-assets uploads.

    Returns
        The current environment name ('qa' or 'prod')

    Raises
        RuntimeError: If no environment has been set via set_env()

    """
    if not _env:
        raise RuntimeError("Environment not set. Call set_env() to set it.")

    return _env


def set_m2m_token_manager(m2m_token_manager: M2MTokenManager) -> None:
    """Set the M2M token manager used to authenticate ows-assets uploads.

    When set, the default client authenticates with an M2M JWT.
    Must be called before using upload(). Ignored when a custom ows_client is passed.

    Raises:
        RuntimeError: If a token manager has already been set

    """
    global _m2m_token_manager

    if _m2m_token_manager is not None:
        raise RuntimeError("M2M token manager already set. It can only be set once.")
    _m2m_token_manager = m2m_token_manager


T = TypeVar("T")


def _batched(iterable: Iterable[T], batch_size: int) -> Iterator[tuple[T, ...]]:
    """Batch data into tuples of length n. The last batch may be shorter.
    Replace with itertools.batched when we require Python 3.12+.
    """
    iterator = iter(iterable)
    while True:
        batch = tuple(islice(iterator, batch_size))
        if not batch:
            break
        yield batch


def _truncate_middle(string: str, max_length: int) -> str:
    if len(string) <= max_length:
        return string

    truncation_marker = "..."

    keep = max_length - len(truncation_marker)
    left = keep // 2
    right = keep - left

    return string[:left] + truncation_marker + string[-right:]


@dataclass(frozen=True)
class _CompletedPart:
    part_number: int
    etag: str


@cache
def _get_ows_client() -> OwsClient:
    return OwsClient(
        environment=get_env(),
        service_name="python-ows-assets-upload",
        retries=_NUM_HTTP_RETRIES,
        m2m_token_manager=_m2m_token_manager,
    )


@cache
def _get_httpx_client() -> httpx.Client:
    return httpx.Client(
        timeout=httpx.Timeout(
            _DEFAULT_HTTPX_TIMEOUT_SECONDS,
            write=_ASSET_UPLOAD_PRESIGNED_URL_EXPIRES_IN_SECONDS,
        ),
        transport=httpx.HTTPTransport(retries=_NUM_HTTP_RETRIES),
    )


def _calculate_num_parts_and_part_size_bytes(file_size_bytes: int) -> tuple[int, int]:
    num_parts = math.ceil(file_size_bytes / _PART_SIZE_BYTES)
    if num_parts <= _MAX_NUM_PARTS:
        return num_parts, _PART_SIZE_BYTES

    max_num_parts_part_size_bytes = math.ceil(file_size_bytes / _MAX_NUM_PARTS)
    return math.ceil(
        file_size_bytes / max_num_parts_part_size_bytes
    ), max_num_parts_part_size_bytes


def _get_source_filename(
    filesystem: AbstractFileSystem, path: str, source_filename: str | None
) -> str:
    if source_filename:
        return source_filename

    parsed_path = path

    if isinstance(filesystem, HTTPFileSystem):
        parsed_path = urlparse(path).path

    parsed_filename = os.path.basename(parsed_path)

    return parsed_filename


def _read_part_data(
    file: BinaryIO,
    part_number: int,
    part_size_bytes: int,
    file_size_bytes: int,
) -> bytes:
    read_offset_bytes = (part_number - 1) * part_size_bytes
    read_length_bytes = min(part_size_bytes, file_size_bytes - read_offset_bytes)
    file.seek(read_offset_bytes)

    return file.read(read_length_bytes)


class _OwsMethod(Protocol):
    def __call__(
        self, service_name: str, path: str, **kwargs: Any
    ) -> httpx.Response: ...


class _OwsClientWrapper:
    """Normalises OwsClient and ImpersonationOwsClient to a single calling convention."""

    def __init__(
        self,
        client: OwsClient | ImpersonationOwsClient,
        impersonated_identity_uuid: str | None = None,
    ) -> None:
        self.get: _OwsMethod
        self.post: _OwsMethod
        self.patch: _OwsMethod

        if isinstance(client, ImpersonationOwsClient):
            if not impersonated_identity_uuid:
                raise ValueError(
                    "impersonated_identity_uuid is required when using ImpersonationOwsClient"
                )
            # Close over impersonated_identity_uuid — ImpersonationOwsClient requires it as a positional arg.
            self.get = lambda service_name, path, **kw: client.get(
                service_name=service_name,
                path=path,
                impersonated_identity_uuid=impersonated_identity_uuid,
                **kw,
            )
            self.post = lambda service_name, path, **kw: client.post(
                service_name=service_name,
                path=path,
                impersonated_identity_uuid=impersonated_identity_uuid,
                **kw,
            )
            self.patch = lambda service_name, path, **kw: client.patch(
                service_name=service_name,
                path=path,
                impersonated_identity_uuid=impersonated_identity_uuid,
                **kw,
            )
        else:
            if impersonated_identity_uuid:
                raise ValueError(
                    "impersonated_identity_uuid should not be provided when using standard OwsClient"
                )
            # Standard OwsClient — delegate directly.
            self.get = client.get
            self.post = client.post
            self.patch = client.patch


class AssetsUploader:
    def __init__(
        self,
        *,
        ows_client: OwsClient | ImpersonationOwsClient,
        impersonated_identity_uuid: str | None = None,
    ) -> None:
        """Initialize AssetsUploader with either a standard or OBO OWS client."""
        self._ows_client = _OwsClientWrapper(ows_client, impersonated_identity_uuid)
        self._httpx_client = _get_httpx_client()

    def _create_multipart_upload(
        self,
        source_filename: str,
        asset_upload_type: str,
        product_id: int,
        track_id: int | None = None,
    ) -> str:
        body = {
            "product_id": product_id,
            "track_unique_id": track_id or 0,
            "original_filename": source_filename,
            "asset_upload_type": asset_upload_type,
        }
        headers = {
            "Orchard-User-Id": _ORCHARD_USER_ID,
            "Content-Type": "application/json",
        }
        response = self._ows_client.post(
            service_name="ows-assets",
            path="/v2/assets/upload",
            json=body,
            headers=headers,
        )
        response.raise_for_status()
        payload = response.json()

        return str(payload["filename"])

    def _get_part_upload_url_batch(
        self,
        destination_filename: str,
        part_numbers: Iterable[int],
    ) -> dict[int, str]:
        params = {
            "part_numbers": ",".join(str(part_number) for part_number in part_numbers)
        }
        headers = {
            "Orchard-User-Id": _ORCHARD_USER_ID,
        }
        path = f"/v2/assets/upload/{destination_filename}"
        response = self._ows_client.get(
            service_name="ows-assets",
            path=path,
            params=params,
            headers=headers,
        )
        response.raise_for_status()
        part_number_to_presigned_url = response.json()["part_number_to_presigned_url"]

        return {
            int(part_number): url
            for part_number, url in part_number_to_presigned_url.items()
        }

    def _generate_part_upload_urls(
        self,
        destination_filename: str,
        num_parts: int,
    ) -> Generator[tuple[int, str], None, None]:
        for part_numbers in _batched(
            range(1, num_parts + 1), _MAX_PRESIGNED_URL_BATCH_SIZE
        ):
            part_number_to_presigned_url = self._get_part_upload_url_batch(
                destination_filename,
                part_numbers,
            )
            for part_number in part_numbers:
                yield part_number, part_number_to_presigned_url[part_number]

    def _complete_multipart_upload(
        self,
        destination_filename: str,
        completed_parts: list[_CompletedPart],
    ) -> None:
        payload = {
            "parts": [
                asdict(completed_part)
                for completed_part in sorted(
                    completed_parts,
                    key=lambda completed_part: completed_part.part_number,
                )
            ]
        }
        headers = {
            "Orchard-User-Id": _ORCHARD_USER_ID,
            "Content-Type": "application/json",
        }
        path = f"/v2/assets/upload/{destination_filename}"
        response = self._ows_client.patch(
            service_name="ows-assets",
            path=path,
            json=payload,
            headers=headers,
        )
        response.raise_for_status()

    def upload(
        self,
        *,
        source_file_location: str,
        asset_upload_type: str,
        product_id: int,
        track_id: int | None = None,
        source_filename: str | None = None,
    ) -> str:
        """Upload product or track asset via ows-assets.

        Works for both standard and OBO uploads — the client type and impersonated_identity_uuid
        are fixed at construction time via the AssetsUploader constructor.

        Args:
            source_file_location: Path to local file, S3 URL (s3://bucket/key), or HTTP(S) URL
            asset_upload_type: Intended use of the asset
            product_id: product_id to associate the asset with
            track_id: Optional track_id for track-level assets (None for product-level)
            source_filename: Optional filename override (auto-detected from path if not provided)

        Returns:
            Destination filename in ows-assets

        """
        filesystem, path = url_to_fs(source_file_location)
        file_size_bytes = int(filesystem.size(path))
        source_filename = _truncate_middle(
            _get_source_filename(filesystem, path, source_filename),
            _MAX_SOURCE_FILENAME_LENGTH,
        )

        destination_filename = self._create_multipart_upload(
            source_filename,
            asset_upload_type,
            product_id,
            track_id,
        )
        num_parts, part_size_bytes = _calculate_num_parts_and_part_size_bytes(
            file_size_bytes
        )
        completed_parts: list[_CompletedPart] = []

        with filesystem.open(path, "rb") as file:
            # We upload parts sequentially for simplicity and to avoid threading issues in UWSGI environments.
            # We don't handle presigned URL expiration because MAX_PRESIGNED_URL_BATCH_SIZE * PART_SIZE_BYTES
            # should upload well within ASSET_UPLOAD_PRESIGNED_URL_EXPIRES_IN_SECONDS.
            for part_number, part_upload_url in self._generate_part_upload_urls(
                destination_filename,
                num_parts,
            ):
                part_data = _read_part_data(
                    file, part_number, part_size_bytes, file_size_bytes
                )

                part_upload_response = self._httpx_client.put(
                    part_upload_url, content=part_data
                )
                part_upload_response.raise_for_status()
                etag = part_upload_response.headers.get("ETag")
                if not etag:
                    msg = (
                        f"Missing ETag header in multipart upload response for part "
                        f"{part_number}."
                    )
                    raise ValueError(msg)
                completed_parts.append(
                    _CompletedPart(part_number=part_number, etag=etag)
                )

        self._complete_multipart_upload(
            destination_filename,
            completed_parts,
        )

        return destination_filename


def upload(
    *,
    source_file_location: str,
    asset_upload_type: str,
    product_id: int,
    track_id: int | None = None,
    source_filename: str | None = None,
    ows_client: OwsClient | None = None,
) -> str:
    """Upload product or track asset via ows-assets.

    This is a convenience wrapper around AssetsUploader.upload().
        See AssetsUploader.upload() for full documentation.

    The default client authenticates with an M2M JWT when a token manager has been
    registered via set_m2m_token_manager().

    CAVEAT:
        This wrapper does not support impersonation (OBO) uploads
        For that, construct an AssetsUploader directly with an ImpersonationOwsClient,
            and call upload() on it.
    """
    client = ows_client or _get_ows_client()
    return AssetsUploader(ows_client=client).upload(
        source_file_location=source_file_location,
        asset_upload_type=asset_upload_type,
        product_id=product_id,
        track_id=track_id,
        source_filename=source_filename,
    )