import gzip
import re
from dataclasses import dataclass, field
from datetime import datetime, timezone
from enum import Enum
from io import IOBase, TextIOBase
from typing import IO, AsyncIterator, Dict, Iterator, Optional, Union, TypedDict

import httpx

from e2b.envd.filesystem import filesystem_pb2
from e2b.exceptions import InvalidArgumentException
from e2b.io_utils import agzip_iter, aiter_io_chunks, gzip_iter, iter_io_chunks


class FileType(Enum):
    """
    Enum representing the type of filesystem object.
    """

    FILE = "file"
    """
    Filesystem object is a file.
    """
    DIR = "dir"
    """
    Filesystem object is a directory.
    """


def map_file_type(ft: filesystem_pb2.FileType):
    if ft == filesystem_pb2.FileType.FILE_TYPE_FILE:
        return FileType.FILE
    elif ft == filesystem_pb2.FileType.FILE_TYPE_DIRECTORY:
        return FileType.DIR


def map_file_type_str(value: Optional[str]) -> Optional[FileType]:
    """Map a `/files` API type string to `FileType`, `None` when unknown."""
    if value == FileType.FILE.value:
        return FileType.FILE
    elif value == FileType.DIR.value:
        return FileType.DIR
    return None


@dataclass
class WriteInfo:
    """
    Sandbox filesystem object information.
    """

    name: str
    """
    Name of the filesystem object.
    """
    type: Optional[FileType]
    """
    Type of the filesystem object.
    """
    path: str
    """
    Path to the filesystem object.
    """
    metadata: Optional[Dict[str, str]] = field(default=None, kw_only=True)
    """
    User-defined metadata stored on the file as `user.e2b.*` extended
    attributes. On writes this reflects the metadata supplied on upload; on
    reads (`get_info`, `list`, `rename`) it reflects any `user.e2b.*` xattr on
    the file, including ones set out-of-band. `None` when none is set.
    """

    @classmethod
    def from_dict(cls, payload: Dict) -> "WriteInfo":
        """Build a `WriteInfo` from a `/files` upload response entry."""
        return cls(
            name=payload["name"],
            type=map_file_type_str(payload.get("type")),
            path=payload["path"],
            metadata=map_metadata(payload.get("metadata")),
        )


@dataclass
class EntryInfo(WriteInfo):
    """
    Extended sandbox filesystem object information.
    """

    size: int
    """
    Size of the filesystem object in bytes.
    """
    mode: int
    """
    File mode and permission bits.
    """
    permissions: str
    """
    String representation of file permissions (e.g. 'rwxr-xr-x').
    """
    owner: str
    """
    Owner of the filesystem object.
    """
    group: str
    """
    Group owner of the filesystem object.
    """
    modified_time: datetime
    """
    Last modification time of the filesystem object.
    """
    symlink_target: Optional[str] = None
    """
    Target of the symlink if the filesystem object is a symlink.
    If the filesystem object is not a symlink, this field is None.
    """


def map_entry_info(entry: filesystem_pb2.EntryInfo) -> EntryInfo:
    return EntryInfo(
        name=entry.name,
        type=map_file_type(entry.type),
        path=entry.path,
        size=entry.size,
        mode=entry.mode,
        permissions=entry.permissions,
        owner=entry.owner,
        group=entry.group,
        modified_time=entry.modified_time.ToDatetime(tzinfo=timezone.utc),
        # Optional, we can't directly access symlink_target otherwise it will be "" instead of None
        symlink_target=(
            entry.symlink_target if entry.HasField("symlink_target") else None
        ),
        metadata=map_metadata(entry.metadata),
    )


class WriteEntry(TypedDict):
    """
    Contains path and data of the file to be written to the filesystem.
    """

    path: str
    data: Union[str, bytes, IO]


class FileStreamReader(Iterator[bytes]):
    """Iterator over a streamed file download.

    Returned by ``Sandbox.files.read(format="stream")``. It owns the underlying
    HTTP response and releases its pooled connection as soon as the stream is
    fully consumed, an error is raised while reading (including the idle-read
    timeout, which raises ``httpx.ReadTimeout``), or the reader is closed.

    There is no garbage-collection safety net, so always consume it fully, use
    it as a context manager, or call :meth:`close`::

        with sandbox.files.read(path, format="stream") as stream:
            for chunk in stream:
                ...
    """

    def __init__(self, response: httpx.Response):
        self._response = response
        self._iterator = response.iter_bytes()
        self._closed = False

    def __iter__(self) -> Iterator[bytes]:
        return self

    def __next__(self) -> bytes:
        try:
            return next(self._iterator)
        except BaseException:
            # Covers normal end (StopIteration) and read errors alike.
            self.close()
            raise

    def close(self) -> None:
        """Release the underlying HTTP connection. Safe to call multiple times."""
        if self._closed:
            return
        self._closed = True
        self._response.close()

    def __enter__(self) -> "FileStreamReader":
        return self

    def __exit__(self, *exc_info) -> None:
        self.close()


class AsyncFileStreamReader(AsyncIterator[bytes]):
    """Async iterator over a streamed file download.

    Returned by ``AsyncSandbox.files.read(format="stream")``. It owns the
    underlying HTTP response and releases its pooled connection as soon as the
    stream is fully consumed, an error is raised while reading (including the
    idle-read timeout, which raises ``httpx.ReadTimeout``), or the reader is
    closed.

    There is no garbage-collection safety net (releasing an async connection
    requires awaiting ``aclose()``, which a finalizer cannot do reliably), so
    always consume it fully, use it as an async context manager, or call
    :meth:`aclose`::

        async with await sandbox.files.read(path, format="stream") as stream:
            async for chunk in stream:
                ...
    """

    def __init__(self, response: httpx.Response):
        self._response = response
        self._iterator = response.aiter_bytes()
        self._closed = False

    def __aiter__(self) -> AsyncIterator[bytes]:
        return self

    async def __anext__(self) -> bytes:
        try:
            return await self._iterator.__anext__()
        except BaseException:
            # Covers normal end (StopAsyncIteration) and read errors alike.
            await self.aclose()
            raise

    async def aclose(self) -> None:
        """Release the underlying HTTP connection. Safe to call multiple times."""
        if self._closed:
            return
        self._closed = True
        await self._response.aclose()

    async def __aenter__(self) -> "AsyncFileStreamReader":
        return self

    async def __aexit__(self, *exc_info) -> None:
        await self.aclose()


def _to_httpx_file(file_path: str, file_data: Union[str, bytes, IO]):
    """Build an httpx multipart `("file", (name, data))` tuple for the upload."""
    if isinstance(file_data, (str, bytes)):
        return ("file", (file_path, file_data))
    elif isinstance(file_data, TextIOBase):
        return ("file", (file_path, file_data.read()))
    elif isinstance(file_data, IOBase):
        return ("file", (file_path, file_data))
    else:
        raise InvalidArgumentException(f"Unsupported data type for file {file_path}")


def to_upload_body(
    data: Union[str, bytes, IO],
    use_gzip: bool = False,
) -> Union[bytes, IO, Iterator[bytes]]:
    """Prepare file data for upload, optionally gzip-compressed.

    File-like objects are streamed in chunks instead of being buffered in
    memory.
    """
    if isinstance(data, (str, bytes)):
        raw = data.encode("utf-8") if isinstance(data, str) else data
        return gzip.compress(raw) if use_gzip else raw
    elif isinstance(data, (TextIOBase, IOBase)):
        if use_gzip:
            return gzip_iter(iter_io_chunks(data))
        if isinstance(data, TextIOBase):
            # Text-mode IO yields str chunks—encode them while streaming.
            return iter_io_chunks(data)
        # httpx streams binary file-like objects in chunks without buffering.
        return data
    else:
        raise InvalidArgumentException(f"Unsupported data type: {type(data)}")


def to_upload_body_async(
    data: Union[str, bytes, IO],
    use_gzip: bool = False,
) -> Union[bytes, AsyncIterator[bytes]]:
    """Prepare file data for upload with async httpx, optionally gzip-compressed.

    File-like objects are streamed in chunks instead of being buffered in
    memory. Async httpx requires an async iterable for streamed request bodies.
    """
    if isinstance(data, (str, bytes)):
        raw = data.encode("utf-8") if isinstance(data, str) else data
        return gzip.compress(raw) if use_gzip else raw
    elif isinstance(data, (TextIOBase, IOBase)):
        chunks = aiter_io_chunks(data)
        return agzip_iter(chunks) if use_gzip else chunks
    else:
        raise InvalidArgumentException(f"Unsupported data type: {type(data)}")


METADATA_HEADER_PREFIX = "X-Metadata-"

# Metadata keys travel as `X-Metadata-<key>` HTTP header names, so they must be
# valid header tokens (RFC 7230); values travel as header values, restricted to
# printable US-ASCII.
_METADATA_KEY_REGEX = re.compile(r"\A[A-Za-z0-9!#$%&'*+\-.^_`|~]+\Z")
_METADATA_VALUE_REGEX = re.compile(r"\A[\x20-\x7e]*\Z")


def validate_metadata(metadata: Optional[Dict[str, str]]) -> None:
    """Validate metadata keys/values before they are sent as upload headers."""
    if not metadata:
        return
    for key, value in metadata.items():
        if not _METADATA_KEY_REGEX.match(key):
            raise InvalidArgumentException(
                f"Invalid metadata key {key!r}: keys must be non-empty and use only "
                "HTTP token characters (letters, digits and !#$%&'*+-.^_`|~)."
            )
        if not _METADATA_VALUE_REGEX.match(value):
            raise InvalidArgumentException(
                f"Invalid metadata value for key {key!r}: values must be printable US-ASCII."
            )


def metadata_to_headers(
    metadata: Optional[Dict[str, str]],
) -> Dict[str, str]:
    """Translate user metadata into the `X-Metadata-*` upload headers envd reads."""
    if not metadata:
        return {}
    return {f"{METADATA_HEADER_PREFIX}{key}": value for key, value in metadata.items()}


def map_metadata(metadata) -> Optional[Dict[str, str]]:
    """Normalize a proto/HTTP metadata map: drop empties and return a plain dict or None."""
    if not metadata:
        return None
    return dict(metadata)