"""File-like IO abstractions for Databricks paths."""

import base64
import io
import logging
import os
import time
from abc import ABC, abstractmethod
from threading import Thread
from typing import TYPE_CHECKING, Optional, IO, AnyStr, Union, BinaryIO

from yggdrasil.io.buffer import BytesIO as _Buffer  # spill-to-disk buffer

from databricks.sdk.errors import InternalError
from databricks.sdk.errors.platform import (
    NotFound,
    ResourceDoesNotExist,
    BadRequest,
)
from databricks.sdk.service.workspace import ImportFormat, ExportFormat

from .path_kind import DatabricksPathKind
from ...pyutils.retry import retry

if TYPE_CHECKING:
    from .path import DatabricksPath


__all__ = [
    "DatabricksIO"
]


LOGGER = logging.getLogger(__name__)
_COPY_CHUNK = 8 * 1024 * 1024   # 8 MB upload chunks

def _prepare_binaryio_and_size(
    data: Union[bytes, bytearray, memoryview, BinaryIO]
) -> tuple[int, _Buffer, bool]:
    """Wrap *data* in a :class:`_Buffer` (``yggdrasil.io.BytesIO``) and return
    ``(byte_count, buffer, owns_buffer)``.

    ``yggdrasil.io.BytesIO`` handles spill-to-disk transparently, so we never
    need a ``SpooledTemporaryFile`` here.
    """
    buf = _Buffer(data)
    return buf.size, buf, True

class DatabricksIO(ABC, IO):
    """File-like interface for Databricks workspace, volume, or DBFS paths."""

    def __init__(
        self,
        path: "DatabricksPath",
        mode: str,
        encoding: Optional[str] = None,
        position: int = 0,
        buffer: Optional[io.BytesIO] = None,
    ):
        super().__init__()

        self.encoding = encoding
        self.mode = mode

        self.path = path

        self.buffer = buffer
        self.position = position

        self._write_flag = False

    def __enter__(self) -> "DatabricksIO":
        """Enter a context manager and connect the underlying path."""
        return self.connect(clone=False)

    def __exit__(self, exc_type, exc_value, traceback):
        """Exit the context manager and close the buffer."""
        self.close()

    def __del__(self):
        if self._need_flush():
            try:
                Thread(target=self.close).start()
            except BaseException:
                pass

    def __next__(self):
        """Iterate over lines in the file."""
        line = self.readline()
        if not line:
            raise StopIteration
        return line

    def __len__(self):
        return self.content_length or 0

    def __iter__(self):
        return self.read_all_bytes().__iter__()

    def __hash__(self):
        return self.path.__hash__()

    def __str__(self):
        return "%s(path=%s)" % (
            self.__class__.__name__,
            self.path.__repr__()
        )

    def __repr__(self):
        return "%s(path=%s)" % (
            self.__class__.__name__,
            self.path.__repr__()
        )

    @classmethod
    def create_instance(
        cls,
        path: "DatabricksPath",
        mode: str,
        encoding: Optional[str] = None,
        position: int = 0,
        buffer: Optional[io.BytesIO] = None,
    ) -> "DatabricksIO":
        """Create the appropriate IO subclass for the given path kind.

        Args:
            path: DatabricksPath to open.
            mode: File mode string.
            encoding: Optional text encoding for text mode.
            position: Initial file cursor position.
            buffer: Optional pre-seeded buffer.

        Returns:
            A DatabricksIO subclass instance.
        """
        if path.kind == DatabricksPathKind.VOLUME:
            return DatabricksVolumeIO(
                path=path,
                mode=mode,
                encoding=encoding,
                position=position,
                buffer=buffer,
            )
        elif path.kind == DatabricksPathKind.DBFS:
            return DatabricksDBFSIO(
                path=path,
                mode=mode,
                encoding=encoding,
                position=position,
                buffer=buffer,
            )
        elif path.kind == DatabricksPathKind.WORKSPACE:
            return DatabricksWorkspaceIO(
                path=path,
                mode=mode,
                encoding=encoding,
                position=position,
                buffer=buffer,
            )
        else:
            raise ValueError(f"Unsupported DatabricksPath kind: {path.kind}")

    @property
    def workspace(self):
        """Return the associated Workspace instance.

        Returns:
            The Workspace bound to the path.
        """
        return self.path.workspace

    @property
    def name(self):
        """Return the name of the underlying path.

        Returns:
            The path name component.
        """
        return self.path.name

    @property
    def mode(self):
        return self._mode

    @mode.setter
    def mode(self, value: str):
        self._mode = value

        # Basic text/binary behavior:
        # - binary -> encoding None
        # - text   -> default utf-8
        if "b" in self._mode:
            self.encoding = None
        else:
            if self.encoding is None:
                self.encoding = "utf-8"

    @property
    def content_length(self) -> int:
        return self.path.content_length

    @content_length.setter
    def content_length(self, value: int):
        self.path.content_length = value

    def size(self):
        """Return the size of the file in bytes.

        Returns:
            The file size in bytes.
        """
        return self.content_length

    @property
    def buffer(self) -> _Buffer:
        """Return the IO buffer (a :class:`_Buffer`), creating it if necessary.

        :class:`_Buffer` (``yggdrasil.io.BytesIO``) spills transparently to
        disk for large payloads, avoiding OOM on big files.
        """
        if self._buffer is None:
            self._buffer = _Buffer()
            self._buffer.seek(self.position, io.SEEK_SET)
        return self._buffer

    @buffer.setter
    def buffer(self, value: Optional[_Buffer]) -> None:
        self._buffer = value

    def clear_buffer(self):
        """Clear any cached in-memory buffer.

        Returns:
            None.
        """
        self._buffer = None

    def clone_instance(self, **kwargs):
        """Clone this IO instance with optional overrides.

        Args:
            **kwargs: Field overrides for the new instance.

        Returns:
            A cloned DatabricksIO instance.
        """
        return self.__class__(
            path=kwargs.get("path", self.path),
            mode=kwargs.get("mode", self.mode),
            encoding=kwargs.get("encoding", self.encoding),
            position=kwargs.get("position", self.position),
            buffer=kwargs.get("buffer", self._buffer),
        )

    @property
    def connected(self):
        """Return True if the underlying path is connected.

        Returns:
            True if connected, otherwise False.
        """
        return self.path.connected

    def connect(self, clone: bool = False) -> "DatabricksIO":
        """Connect the underlying path and optionally return a clone.

        Args:
            clone: Whether to return a cloned instance.

        Returns:
            The connected DatabricksIO instance.
        """
        path = self.path.connect()

        if clone:
            return self.clone_instance(path=path)

        self.path = path
        return self

    def close(self, flush: bool = True):
        """Flush pending writes and close the buffer.

        Args:
            flush: Checks flush data to commit to remote location

        Returns:
            None.
        """
        if flush:
            self.flush()
        self.clear_buffer()

    @property
    def closed(self):
        return False

    def fileno(self):
        """Return a pseudo file descriptor based on object hash.

        Returns:
            An integer file descriptor-like value.
        """
        return hash(self)

    def isatty(self):
        return False

    def tell(self):
        """Return the current cursor position.

        Returns:
            The current position in bytes.
        """
        return self.position

    def seekable(self):
        """Return True to indicate seek support.

        Returns:
            True.
        """
        return True

    def seek(self, offset, whence=0, /):
        """Move the cursor to a new position.

        Args:
            offset: Offset in bytes.
            whence: Reference point (start, current, end).

        Returns:
            The new position in bytes.
        """
        if whence == io.SEEK_SET:
            new_position = offset
        elif whence == io.SEEK_CUR:
            new_position = self.position + offset
        elif whence == io.SEEK_END:
            end_position = self.content_length
            new_position = end_position + offset
        else:
            raise ValueError("Invalid value for whence")

        if new_position < 0:
            raise ValueError("New position is before the start of the file")

        if self._buffer is not None:
            self._buffer.seek(new_position, io.SEEK_SET)

        self.position = new_position
        return self.position

    def readable(self):
        """Return True to indicate read support.

        Returns:
            True.
        """
        return True

    def getvalue(self):
        """Return the buffer contents, reading from remote if needed."""
        if self._buffer is not None:
            return self._buffer.to_bytes()
        return self.read_all_bytes()

    def getbuffer(self):
        """Return the underlying buffer instance."""
        return self.buffer

    @abstractmethod
    def read_byte_range(self, start: int, length: int, allow_not_found: bool = False) -> bytes:
        """Read a byte range from the remote path.

        Args:
            start: Starting byte offset.
            length: Number of bytes to read.
            allow_not_found: Whether to suppress missing-path errors.

        Returns:
            The bytes read from the remote path.
        """
        pass

    def read_all_bytes(self, use_cache: bool = True, allow_not_found: bool = False) -> bytes:
        """Read the full contents into memory, optionally caching in the buffer.

        Args:
            use_cache: Whether to keep a cached copy in the buffer.
            allow_not_found: Whether to suppress missing-path errors.

        Returns:
            File contents as bytes.
        """
        if use_cache and self._buffer is not None:
            cached = self._buffer.to_bytes()
            if len(cached) == self.content_length:
                return cached
            self._buffer.close()
            self._buffer = None

        data = self.read_byte_range(0, self.content_length, allow_not_found=allow_not_found)

        # Keep size accurate even if the backend didn't know it upfront
        self.content_length = len(data)

        if use_cache and self._buffer is None:
            self._buffer = _Buffer(data)
            self._buffer.seek(self.position, io.SEEK_SET)

        return data

    def read(self, n=-1, use_cache: bool = True):
        """Read up to ``n`` bytes/characters from the file.

        Args:
            n: Number of bytes/characters to read; -1 for all.
            use_cache: Whether to use cached contents.

        Returns:
            The read bytes or string depending on mode.
        """
        current_position = self.position
        all_data = self.read_all_bytes(use_cache=use_cache)

        if n == -1:
            n = self.content_length - current_position

        data = all_data[current_position:current_position + n]
        read_length = len(data)

        self.position += read_length

        if self.encoding:
            return data.decode(self.encoding)
        return data

    def readline(self, limit=-1, use_cache: bool = True):
        """Read a single line from the file.

        Args:
            limit: Max characters/bytes to read; -1 for no limit.
            use_cache: Whether to use cached contents.

        Returns:
            The next line as bytes or string.
        """
        if self.encoding:
            # Text-mode: accumulate characters
            out_chars = []
            read_chars = 0

            while limit == -1 or read_chars < limit:
                ch = self.read(1, use_cache=use_cache)
                if not ch:
                    break
                out_chars.append(ch)
                read_chars += 1
                if ch == "\n":
                    break

            return "".join(out_chars)

        # Binary-mode: accumulate bytes
        line_bytes = bytearray()
        bytes_read = 0

        while limit == -1 or bytes_read < limit:
            b = self.read(1, use_cache=use_cache)
            if not b:
                break
            line_bytes.extend(b)
            bytes_read += 1
            if b == b"\n":
                break

        return bytes(line_bytes)

    def readlines(self, hint=-1, use_cache: bool = True):
        """Read all lines from the file.

        Args:
            hint: Optional byte/char count hint; -1 for no hint.
            use_cache: Whether to use cached contents.

        Returns:
            A list of lines.
        """
        lines = []
        total = 0

        while True:
            line = self.readline(use_cache=use_cache)
            if not line:
                break
            lines.append(line)
            total += len(line)
            if hint != -1 and total >= hint:
                break

        return lines

    def writable(self):
        """Return True to indicate write support.

        Returns:
            True.
        """
        return True

    @abstractmethod
    def write_all_bytes(self, data: Union[bytes, IO[bytes]]):
        """Write raw bytes to the remote path.

        Args:
            data: Bytes to write.

        Returns:
            None.
        """
        pass

    def truncate(self, size=None, /):
        """Resize the file to ``size`` bytes.

        Args:
            size: Target size in bytes (defaults to current position).

        Returns:
            The new size in bytes.
        """
        if size is None:
            size = self.position

        if self._buffer is None:
            return self.write_all_bytes(data=b"\x00" * size)

        self._buffer.truncate(size)

        self.content_length = size
        self._write_flag = True

        return size

    def _need_flush(self):
        return self._write_flag and self._buffer is not None

    def flush(self):
        """Flush buffered writes to the remote path."""
        if self._need_flush():
            self.write_all_bytes(data=self._buffer.to_bytes())
            self._write_flag = False

    def write(self, data: AnyStr) -> int:
        """Write data to the buffer and mark for flush.

        Args:
            data: String or bytes to write.

        Returns:
            The number of bytes written.
        """
        if isinstance(data, str):
            data = data.encode(self.encoding or "utf-8")

        written = self.buffer.write(data)

        self.position += written
        self.content_length = self.position
        self._write_flag = True

        return written

    def writelines(self, lines) -> None:
        """Write multiple lines to the buffer.

        Args:
            lines: Iterable of lines to write.

        Returns:
            None.
        """
        for line in lines:
            if isinstance(line, str):
                line = line.encode(self.encoding or "utf-8")
            elif not isinstance(line, (bytes, bytearray)):
                raise TypeError(
                    "a bytes-like or str object is required, not '{}'".format(type(line).__name__)
                )

            data = line + b"\n" if not line.endswith(b"\n") else line
            self.write(data)

    def get_output_stream(self, *args, **kwargs):
        """Return this instance for compatibility with Arrow APIs.

        Returns:
            The current DatabricksIO instance.
        """
        return self

    def copy_to(
        self,
        dest: Union["DatabricksIO", "DatabricksPath", str]
    ) -> None:
        """Copy the file contents to another Databricks IO/path.

        Args:
            dest: Destination IO, DatabricksPath, or path string.

        Returns:
            None.
        """
        data = self.read_all_bytes(use_cache=False)

        if isinstance(dest, DatabricksIO):
            dest.write_all_bytes(data=data)
        elif hasattr(dest, "write"):
            dest.write(data)
        else:
            from .path import DatabricksPath

            dest_path = DatabricksPath.parse(dest, client=self.workspace)

            with dest_path.open(mode="wb") as d:
                return self.copy_to(dest=d)

    # ---- format helpers ----

    def _reset_for_write(self):
        """Reset the buffer to an empty writable state."""
        if self._buffer is not None:
            self._buffer.seek(0, io.SEEK_SET)
            self._buffer.truncate(0)
        else:
            self._buffer = _Buffer()

        self.position = 0
        self.content_length = 0
        self._write_flag = True


class DatabricksWorkspaceIO(DatabricksIO):
    """IO adapter for Workspace files."""

    @retry(exceptions=(InternalError,))
    def read_byte_range(self, start: int, length: int, allow_not_found: bool = False) -> bytes:
        """Read bytes from a Workspace file.

        Args:
            start: Starting byte offset.
            length: Number of bytes to read.
            allow_not_found: Whether to suppress missing-path errors.

        Returns:
            Bytes read from the file.
        """
        if length == 0:
            return b""

        sdk = self.workspace.workspace_client()
        client = sdk.workspace
        full_path = self.path.full_path()

        result = client.download(
            path=full_path,
            format=ExportFormat.AUTO,
        )

        if result is None:
            return b""

        data = result.read()

        end = start + length
        return data[start:end]

    @retry(exceptions=(InternalError,))
    def write_all_bytes(self, data: Union[bytes, IO[bytes]]):
        """Write bytes to a Workspace file.

        Args:
            data: Union[bytes, IO[bytes]] to write.

        Returns:
            The DatabricksWorkspaceIO instance.
        """
        sdk = self.workspace.workspace_client()
        workspace_client = sdk.workspace
        full_path = self.path.full_path()

        if isinstance(data, bytes):
            bsize = len(data)
        elif isinstance(data, _Buffer):
            bsize = data.size
        elif isinstance(data, io.BytesIO):
            bsize = len(data.getvalue())
        else:
            bsize = None

        LOGGER.debug(
            "Writing %s(size=%s) in %s",
            type(data),
            bsize,
            self
        )

        try:
            workspace_client.upload(
                full_path,
                data,
                format=ImportFormat.AUTO,
                overwrite=True
            )
        except (NotFound, ResourceDoesNotExist, BadRequest):
            self.path.parent.mkdir(parents=True, exist_ok=True)

            workspace_client.upload(
                full_path,
                data,
                format=ImportFormat.AUTO,
                overwrite=True
            )

        self.path.reset_metadata(
            is_file=True,
            is_dir=False,
            size=bsize,
            mtime=time.time()
        )

        LOGGER.info(
            "Written %s bytes in %s",
            bsize,
            self
        )

        return self


class DatabricksVolumeIO(DatabricksIO):
    """IO adapter for Unity Catalog volume files."""

    @retry(exceptions=(InternalError,))
    def read_byte_range(self, start: int, length: int, allow_not_found: bool = False) -> bytes:
        """Read bytes from a volume file.

        Args:
            start: Starting byte offset (0-based).
            length: Number of bytes to read.
            allow_not_found: Whether to suppress missing-path errors.

        Returns:
            Bytes read from the file.
        """
        if length <= 0:
            return b""
        if start < 0:
            raise ValueError(f"start must be >= 0, got {start}")
        if length < 0:
            raise ValueError(f"length must be >= 0, got {length}")

        sdk = self.workspace.workspace_client()
        client = sdk.files
        full_path = self.path.full_path()

        try:
            resp = client.download(full_path)
        except (NotFound, ResourceDoesNotExist, BadRequest, InternalError) as e:
            # Databricks SDK exceptions vary a bit by version; keep it pragmatic.
            if allow_not_found:
                return b""
            raise

        data = resp.contents.read()

        # If start is past EOF, return empty (common file-like behavior).
        if start >= len(data):
            return b""

        end = start + length
        return data[start:end]

    def write_all_bytes(
        self,
        data: Union[bytes, bytearray, memoryview, BinaryIO],
        *,
        overwrite: bool = True,
        part_size: Optional[int] = None,
        use_parallel: bool = True,
        parallelism: Optional[int] = None,
    ):
        """Write bytes/stream to a volume file safely (BinaryIO upload)."""
        sdk = self.workspace.workspace_client()
        client = sdk.files
        full_path = self.path.full_path()

        LOGGER.debug("Writing all bytes in %s", self)

        size, bio, should_close = _prepare_binaryio_and_size(data)

        def _upload():
            return client.upload(
                full_path,
                bio,
                overwrite=overwrite,
                part_size=part_size,
                use_parallel=use_parallel,
                parallelism=parallelism,
            )

        try:
            _ = _upload()
        except (NotFound, ResourceDoesNotExist, BadRequest, InternalError):
            self.path.parent.mkdir(parents=True, exist_ok=True)
            # Important: rewind if possible before retry
            try:
                bio.seek(0)
            except Exception:
                pass
            _ = _upload()
        finally:
            if should_close:
                try:
                    bio.close()
                except Exception:
                    pass

        self.path.reset_metadata(
            is_file=True,
            is_dir=False,
            size=size,
            mtime=time.time(),
        )

        LOGGER.info("Written %s bytes in %s", size or "all", self.path)

        return self  # or return result if your API prefers that


class DatabricksDBFSIO(DatabricksIO):
    """IO adapter for DBFS files."""

    @retry(exceptions=(InternalError,))
    def read_byte_range(self, start: int, length: int, allow_not_found: bool = False) -> bytes:
        """Read bytes from a DBFS file.

        Args:
            start: Starting byte offset.
            length: Number of bytes to read.
            allow_not_found: Whether to suppress missing-path errors.

        Returns:
            Bytes read from the file.
        """
        if length == 0:
            return b""

        sdk = self.workspace.workspace_client()
        client = sdk.dbfs
        full_path = self.path.full_path()

        read_bytes = bytearray()
        bytes_to_read = length
        current_position = start

        while bytes_to_read > 0:
            chunk_size = min(bytes_to_read, 2 * 1024 * 1024)

            resp = client.read(
                path=full_path,
                offset=current_position,
                length=chunk_size
            )

            if not resp.data:
                break

            # resp.data is base64; decode and move offsets by *decoded* length
            resp_data_bytes = base64.b64decode(resp.data)

            read_bytes.extend(resp_data_bytes)
            bytes_read = len(resp_data_bytes)  # <-- FIX (was base64 string length)
            current_position += bytes_read
            bytes_to_read -= bytes_read

        return bytes(read_bytes)

    @retry(exceptions=(InternalError,))
    def write_all_bytes(self, data: Union[bytes, IO[bytes]]):
        """Write bytes to a DBFS file.

        Args:
            data: Union[bytes, IO[bytes]] to write.

        Returns:
            The DatabricksDBFSIO instance.
        """
        sdk = self.workspace.workspace_client()
        client = sdk.dbfs
        full_path = self.path.full_path()

        LOGGER.debug(
            "Writing all bytes in %s",
            self
        )

        try:
            with client.open(
                path=full_path,
                read=False,
                write=True,
                overwrite=True
            ) as f:
                f.write(data)
        except (NotFound, ResourceDoesNotExist, BadRequest):
            self.path.parent.mkdir(parents=True, exist_ok=True)

            with client.open(
                path=full_path,
                read=False,
                write=True,
                overwrite=True
            ) as f:
                f.write(data)

        LOGGER.info(
            "Written all bytes in %s",
            self
        )

        self.path.reset_metadata(
            is_file=True,
            is_dir=False,
            size=len(data),
            mtime=time.time()
        )
