r"""Representation of objects stored in AWS Simple Storage Service (S3)."""

from __future__ import annotations

import re

from copy import deepcopy
from pathlib import Path
from tempfile import gettempdir
from typing import Generator, Optional

from .client import S3Client
from ..resource import PathStr, Resource


class S3Object(Resource):
    r"""Object stored in AWS Simple Storage Service (S3)."""

    def __init__(self, bucket: str, key: str, path: PathStr = None) -> None:
        r"""Initialize AWS S3 object.

        Args:
            bucket: Name of AWS S3 bucket containing this object.
            key: Key of object in AWS S3 bucket with forward slashes as path separators.
                When string ends with a forward slash, the S3 object represents a set
                of S3 objects which share this common key prefix. The local path is
                in this case a directory tree rather than a single file path.
            path: File path of local file or directory corresponding to AWS S3 object.
                When not specified, a temporary path in the tempfile.gettempdir() is
                constructed from the bucket name and S3 object key.

        """
        if bucket is None:
            raise TypeError("S3 bucket name must be str")
        if bucket == "":
            raise ValueError("S3 bucket name must not be an empty string")
        key = self.normkey(key)
        path = Path(path) if path else self.default_path(bucket=bucket, key=key)
        if not key.endswith("/") and path.is_dir():
            key += "/"
        super().__init__(path)
        self._bucket = bucket
        self._key = key

    @property
    def s3(self) -> S3Client:
        r"""
        Returns:
            client: Underlying connected S3 client.

        """
        return S3Client.default().connect()

    @staticmethod
    def normkey(key: str) -> str:
        r"""Normalize S3 object key string."""
        if not isinstance(key, str):
            raise TypeError("S3 object key must be str")
        if key == "":
            return "/"
        key = re.sub(r"[/\\]{1,}", "/", key)
        if len(key) > 1 and key.startswith("/"):
            key = key[1:]
        return key

    @staticmethod
    def default_path(bucket: str, key: str) -> Path:
        r"""Get default local path of S3 object copy.

        Args:
            bucket: Name of S3 bucket.
            key: Key of S3 object.

        Returns:
            path: Absolute path of local copy of S3 object.

        """
        key = S3Object.normkey(key)
        if key.startswith("/"):
            key = key[1:]
        return Path(gettempdir()).joinpath("deepali", "cache", "s3", bucket, key)

    def reset_path(self) -> S3Object:
        r"""Reset ``path`` to ``default_path`` for given ``bucket`` and ``key``.

        Returns:
            self: This instance.

        """
        self._path = self.default_path(bucket=self.bucket, key=self.key)
        return self

    @classmethod
    def from_uri(cls, uri: str) -> S3Object:
        r"""Create AWS S3 object from URI.

        Args:
            uri: URI of S3 object. Must start with 's3://' followed by the bucket name.
                The remainder of the URI represents the object key, excluding the forward
                slash separating the bucket name from the object key.

        Returns:
            obj: S3 object instance.

        """
        match = re.match(r"[sS]3://(?P<bucket>[^/]+)/(?P<key>.*)", uri)
        if not match:
            raise ValueError("Invalid AWS S3 object URI: %s", uri)
        return cls(bucket=match["bucket"], key=match["key"])

    @property
    def uri(self) -> str:
        r"""URI of storage object."""
        return "s3://" + self.bucket + "/" + self.key

    @property
    def bucket(self) -> str:
        r"""Name of S3 bucket containing storage object."""
        return self._bucket

    @property
    def key(self) -> str:
        r"""Name of S3 key corresponding to storage object referenced by this instance."""
        return self._key

    @property
    def name(self) -> str:
        r"""Name of storage object including file name extension, excluding directory path."""
        key = self.key
        if key.endswith("/"):
            key = key[:-1]
        return key.rsplit("/", 1)[-1]

    def with_bucket(self, bucket: str) -> S3Object:
        r"""Create copy of storage object reference with modified ``bucket``.

        Args:
            bucket: New bucket name.

        Returns:
            self: New storage object reference with modified ``bucket`` property.

        """
        if bucket is None:
            raise TypeError("Bucket name must be str")
        if bucket == "":
            raise ValueError("Invalid S3 bucket name")
        obj = deepcopy(self)
        obj._bucket = bucket
        return obj

    def with_key(self, key: str) -> S3Object:
        r"""Create copy of storage object reference with modified ``key``.

        Args:
            key: New S3 object key.

        Returns:
            self: New storage object reference with modified ``key`` property.

        """
        if key is None:
            raise TypeError("S3 object key must be str")
        if key == "":
            raise ValueError("Invalid S3 object key")
        obj = deepcopy(self)
        obj._key = key
        return obj

    def exists(self) -> bool:
        r"""Whether object exists in AWS S3."""
        return self.s3.exists(bucket=self.bucket, key=self.key)

    def is_file(self) -> bool:
        r"""Whether AWS S3 object exists and represents a file."""
        return not self.key.endswith("/") and self.exists()

    def is_dir(self) -> bool:
        r"""Whether AWS S3 object exists and represents a directory."""
        return self.key.endswith("/") and self.exists()

    def iterdir(self, prefix: str = None) -> Generator[Resource, None, None]:
        r"""List S3 objects within directory, excluding subfolder contents.

        Args:
            prefix: Name prefix.

        Returns:
            iterable: Generator of S3 objects.

        """
        if self.key.endswith("/"):
            if prefix is None:
                prefix = ""
            if self.key != "/":
                prefix = self.key + prefix
            for key in self.s3.iterdir(bucket=self.bucket, prefix=prefix):
                yield S3Object(bucket=self.bucket, key=key)

    def pull(self, force: bool = False) -> S3Object:
        r"""Download content from AWS S3 to local path.

        Args:
            force: Whether to force download even if local path already exists.

        Returns:
            self: This instance.

        """
        if self.key.endswith("/"):
            self.s3.download_files(
                bucket=self.bucket, prefix=self.key, path=self.path, overwrite=force
            )
        else:
            try:
                self.s3.download_file(
                    bucket=self.bucket, key=self.key, path=self.path, overwrite=force
                )
            except FileExistsError:
                pass
        return self

    def push(self, force: bool = False) -> S3Object:
        r"""Upload content of local path to AWS S3.

        Args:
            force: Whether to force upload even if S3 object already exists.

        Returns:
            self: This instance.

        """
        if self.key.endswith("/"):
            self.s3.upload_files(
                bucket=self.bucket, prefix=self.key, path=self.path, overwrite=force
            )
        else:
            self.s3.upload_file(bucket=self.bucket, key=self.key, path=self.path, overwrite=force)
        return self

    def read_bytes(self) -> bytes:
        r"""Read file content from local path if it exists, or corresponding S3 object otherwise.

        Returns:
            Binary file content of storage object.

        """
        assert not self.key.endswith("/"), "S3 object key must referrence a file object"
        if self.path.is_file():
            return self.path.read_bytes()
        return self.s3.read_bytes(bucket=self.bucket, key=self.key)

    def write_bytes(self, data: bytes) -> S3Object:
        r"""Write bytes directly to storage object.

        Args:
            data: Binary data to write.

        Returns:
            self: This storage object.

        """
        assert not self.key.endswith("/"), "S3 object key must referrence a file object"
        self.s3.write_bytes(bucket=self.bucket, key=self.key, data=data)
        return self

    def read_text(self, encoding: Optional[str] = None) -> str:
        r"""Read text file content from local path if it exists, or corresponding S3 object otherwise.

        Args:
            encoding: Text encoding.

        Returns:
            text: Decoded text file content of storage object.

        """
        data = self.read_bytes()
        return data.decode(encoding) if encoding else data.decode()

    def write_text(self, text: str, encoding: Optional[str] = None) -> S3Object:
        r"""Write text to storage object.

        Args:
            text: Text to write.
            encoding: Text encoding.

        Returns:
            self: This storage object.

        """
        self.write_bytes(text.encode(encoding) if encoding else text.encode())
        return self

    def rmdir(self) -> S3Object:
        r"""Remove S3 objects both locally and from cloud storage."""
        if not self.key.endswith("/"):
            raise NotADirectoryError("S3 object key prefix must end with '/' character")
        super().rmdir()
        self.s3.delete_files(bucket=self.bucket, prefix=self.key)
        return self

    def unlink(self) -> S3Object:
        r"""Remove S3 object both locally and from cloud storage."""
        assert not self.key.endswith("/"), "S3 object key must not end with '/' character"
        super().unlink()
        self.s3.delete_file(bucket=self.bucket, key=self.key)
        return self

    def __repr__(self) -> str:
        r"""Get human-readable string representation of storage object reference."""
        return type(self).__name__ + "(bucket='{}', key='{}', path='{}')".format(
            self.bucket, self.key, self.path
        )
