File

def as_audio_file(self) -> "AudioFile":
    """Convert the file to a `AudioFile` object."""
    if isinstance(self, AudioFile):
        return self
    file = AudioFile(**self.model_dump())
    file._set_stream(self._catalog, caching_enabled=self._caching_enabled)
    return file

def as_image_file(self) -> "ImageFile":
    """Convert the file to a `ImageFile` object."""
    if isinstance(self, ImageFile):
        return self
    file = ImageFile(**self.model_dump())
    file._set_stream(self._catalog, caching_enabled=self._caching_enabled)
    return file

def as_text_file(self) -> "TextFile":
    """Convert the file to a `TextFile` object."""
    if isinstance(self, TextFile):
        return self
    file = TextFile(**self.model_dump())
    file._set_stream(self._catalog, caching_enabled=self._caching_enabled)
    return file

def as_video_file(self) -> "VideoFile":
    """Convert the file to a `VideoFile` object."""
    if isinstance(self, VideoFile):
        return self
    file = VideoFile(**self.model_dump())
    file._set_stream(self._catalog, caching_enabled=self._caching_enabled)
    return file

@classmethod
def at(
    cls, uri: str | os.PathLike[str], session: "Session | None" = None
) -> "Self":
    """Construct a File from a full URI in one call.

    Args:
        uri: Full URI or path to the file
            (e.g. ``s3://bucket/path/to/file.png`` or ``/local/path``).
        session: Optional session instance. If None, the current session is used.

    Returns:
        File: A new File object pointing to the given URI.

    Example:
        ```py
        file = File.at("s3://bucket/path/to/output.png")
        with file.open("wb") as f: ...
        ```
    """
    from datachain.client.fsspec import Client
    from datachain.query.session import Session

    if session is None:
        session = Session.get()
    catalog = session.catalog
    uri_str = stringify_path(uri)
    if uri_str.endswith(("/", os.sep)):
        raise ValueError(
            f"File.at directory URL/path given (trailing slash), got: {uri_str}"
        )
    client_cls = Client.get_implementation(uri_str)
    source, rel_path = client_cls.split_url(uri_str)
    source_uri = client_cls.storage_uri(source)
    file = cls(source=source_uri, path=rel_path)
    file._set_stream(catalog)
    return file

def ensure_cached(self) -> None:
    """Download the file to the local cache.

    `get_local_path` can be used to return the path to the cached copy on disk.
    This is useful when you need to pass the file to code that expects a local
    filesystem path (e.g. ``ffmpeg``, ``opencv``, ``pandas``, etc).

    Example:
        ```py
        file.ensure_cached()
        local_path = file.get_local_path()
        df = pandas.read_csv(local_path)
        ```
    """
    if self._catalog is None:
        raise RuntimeError(
            "cannot download file to cache because catalog is not setup"
        )
    client = self._catalog.get_client(self.source)
    client.download(self, callback=self._download_cb)

def export(
    self,
    output: str | os.PathLike[str],
    placement: ExportPlacement = "fullpath",
    use_cache: bool = True,
    link_type: Literal["copy", "symlink"] = "copy",
    client_config: dict | None = None,
    *,
    content_type: str | None = None,
    content_disposition: str | None = None,
    cache_control: str | None = None,
    content_encoding: str | None = None,
    metadata: dict[str, str] | None = None,
    write_options: dict[str, Any] | None = None,
) -> None:
    """Copy or link this file into an output directory.

    Args:
        output: Destination directory.  Accepts a local OS path, a cloud
            prefix fsspec URI (``s3://…``, ``gs://…``, ``az://…``).
        placement: How to build the path under *output*:

            - ``"fullpath"`` (default) — ``output/bucket/dir/file.txt``
            - ``"filepath"`` — ``output/dir/file.txt``
            - ``"filename"`` — ``output/file.txt``
            - ``"etag"`` — ``output/<etag>.txt``
        use_cache: If True, download to local cache first.  Also
            required for symlinking remote files.
        link_type: ``"copy"`` (default) or ``"symlink"``.
            Symlink falls back to copy for virtual files and for
            remote files when *use_cache* is False.
        client_config: Extra kwargs forwarded to the storage client.
        content_type: ``Content-Type`` to set on the exported object.
        content_disposition: ``Content-Disposition`` to set on the object.
        cache_control: ``Cache-Control`` to set on the object.
        content_encoding: ``Content-Encoding`` to set on the object.
        metadata: Custom key/value metadata to attach to the object.
        write_options: Raw, backend-native write kwargs merged after the
            normalized fields above. Object metadata is applied only to
            copied (not symlinked) files and is ignored on the local
            filesystem.

    Example:
        ```py
        # flat export by filename
        f.export("./export", placement="filename")

        # export to a cloud prefix
        f.export("s3://output-bucket/results", placement="filepath")

        # pass storage credentials via client_config
        f.export("s3://bucket/out", client_config={"aws_access_key_id": "…"})

        # symlink from local cache (avoids re-downloading)
        f.export("./local_out", use_cache=True, link_type="symlink")
        ```
    """
    if self._catalog is None:
        raise RuntimeError("Cannot export file: catalog is not set")

    self._caching_enabled = use_cache

    suffix = self._get_destination_suffix(placement)
    output_str = stringify_path(output)
    client = self._catalog.get_client(output_str, **(client_config or {}))

    # Normalization and traversal safety: for local exports, resolve to absolute
    # and validate the suffix. Cloud exports skip this — the cloud client already
    # rejects path-traversal characters in object keys.
    if client.PREFIX == "file://":
        from datachain.fs.utils import is_subpath

        output_os = client.fs._strip_protocol(output_str)
        # On Windows, normalize backslash separators to forward slashes
        # so posixpath.join produces consistent paths.  On Linux/macOS
        # backslash is a legal filename character and must not be replaced.
        output_abs = os.path.abspath(output_os)
        if os.name == "nt":
            output_abs = output_abs.replace("\\", "/")
        dst = posixpath.join(output_abs, suffix)

        try:
            client.validate_file_path(suffix)
        except ValueError as exc:
            raise FileError(str(exc), stringify_path(output), suffix) from None

        if not is_subpath(output_abs, dst):
            raise FileError(
                "destination is not within output directory",
                stringify_path(output),
                dst,
            )
    else:
        # For cloud exports, simply join the output prefix and suffix.
        # Relying on cloud clients to do their own validation.
        dst = posixpath.join(output_str, suffix)

    client.fs.makedirs(posixpath.dirname(dst), exist_ok=True)

    if link_type == "symlink":
        try:
            return self._symlink_to(dst)
        except OSError as exc:
            if exc.errno not in (errno.ENOTSUP, errno.EXDEV, errno.ENOSYS):
                raise

    self.save(
        dst,
        client_config=client_config,
        content_type=content_type,
        content_disposition=content_disposition,
        cache_control=cache_control,
        content_encoding=content_encoding,
        metadata=metadata,
        write_options=write_options,
    )

def get_fs_path(self) -> str:
    """Combine ``source`` and ``path`` into the full location string.

    For cloud backends the result is a full URI-like (``s3://…``, ``gs://…``).
    For local files the result is a **bare OS path** (no ``file://``
    prefix).

    Examples:
        ```py
        # Cloud (S3, GCS, Azure, …)
        f = File(source="s3://my-bucket", path="data/image.jpg")
        f.get_fs_path()  # 's3://my-bucket/data/image.jpg'

        # Local files
        f = File(source="file:///home/user/project", path="out/result.csv")
        f.get_fs_path()  # '/home/user/project/out/result.csv'

        # No source — returns the relative path as-is
        f = File(source="", path="dir/file.txt")
        f.get_fs_path()  # 'dir/file.txt'
        ```

    Raises:
        FileError: If ``path`` is empty, ends with ``/``, or contains
            ``.`` / ``..`` segments.  For local files, also rejects
            absolute paths, drive letters, and empty segments.
    """
    # Fast path: return cached result when source and path are unchanged.
    cache = self._fs_path_cache
    if cache is not None and cache[0] == self.source and cache[1] == self.path:
        return cache[2]

    from datachain.client.fsspec import Client

    client_cls = Client.get_implementation(self.source or "")
    try:
        client_cls.validate_file_path(self.path)
        client_cls.validate_source(self.source)
    except ValueError as exc:
        raise FileError(str(exc), self.source, self.path) from exc
    path = self.path

    if not self.source:
        result = path
    elif self.source.startswith("file://"):
        base_path = client_cls.FS_CLASS._strip_protocol(self.source)
        result = Path(base_path, *PurePosixPath(path).parts).as_posix()
    else:
        # Cloud: build a full URI.
        name = client_cls.FS_CLASS._strip_protocol(self.source)
        base = str(client_cls.storage_uri(name))
        if base.endswith("/"):
            result = f"{base}{path}"
        else:
            result = f"{base}/{path}"

    self._fs_path_cache = (self.source, self.path, result)
    return result

def get_local_path(self) -> str | None:
    """Return path to a file in a local cache.

    Returns None if file is not cached.
    Raises an exception if cache is not setup.
    """
    if self._catalog is None:
        raise RuntimeError(
            "cannot resolve local file path because catalog is not setup"
        )
    return self._catalog.cache.get_path(self)

@contextmanager
def open(
    self,
    mode: str = "rb",
    *,
    client_config: dict[str, Any] | None = None,
    content_type: str | None = None,
    content_disposition: str | None = None,
    cache_control: str | None = None,
    content_encoding: str | None = None,
    metadata: dict[str, str] | None = None,
    write_options: dict[str, Any] | None = None,
    block_size: int | None = None,
    **open_kwargs,
) -> Iterator[Any]:
    """Open the file and return a file-like object.

    Supports both read ("rb", "r") and write modes (e.g. "wb", "w", "ab").
    When opened in a write mode, metadata is refreshed after closing.

    Reads always consult the cache (hit → local, miss → remote stream);
    ``_caching_enabled`` only governs whether a miss is also persisted.

    In write modes, ``content_type``, ``content_disposition``,
    ``cache_control``, ``content_encoding``, ``metadata`` and the raw
    ``write_options`` escape hatch set object metadata on the written
    object (mapped per backend, ignored on the local filesystem).

    ``block_size`` sets the read-ahead block for uncached remote reads, and
    the upload part size for writes. HTTP(S) reads ignore it.
    """
    writing = any(ch in mode for ch in "wax+")

    write_cfg = WriteConfig(
        content_type=content_type,
        content_disposition=content_disposition,
        cache_control=cache_control,
        content_encoding=content_encoding,
        metadata=metadata,
        write_options=write_options,
    )
    if not writing and not write_cfg.is_empty():
        raise ValueError("write metadata options are only valid in write modes")

    if self.location and writing:
        raise VFileError(
            "Writing to virtual file is not supported",
            self.source,
            self.path,
        )

    if self._catalog is None:
        raise RuntimeError("Cannot open file: catalog is not set")

    base_cfg = getattr(self._catalog, "client_config", {}) or {}
    merged_cfg = {**base_cfg, **(client_config or {})}
    client: Client = self._catalog.get_client(self.source, **merged_cfg)

    if not writing:
        if self.location:
            with VFileRegistry.open(
                self,
                self.location,  # type: ignore[arg-type]
                block_size=block_size,
            ) as f:
                with self._wrap_text(f, mode, open_kwargs=open_kwargs) as wrapped:
                    yield wrapped
            return
        if self._caching_enabled:
            self.ensure_cached()
        with client.open_object(
            self, use_cache=True, cb=self._download_cb, block_size=block_size
        ) as f:
            with self._wrap_text(f, mode, open_kwargs=open_kwargs) as wrapped:
                yield wrapped
        return

    # write path
    full_path = self.get_fs_path()
    fs_mode = mode if "b" in mode else f"{mode}b"
    binary_kwargs = {
        k: v for k, v in open_kwargs.items() if k not in self._TEXT_WRAPPER_ALLOWED
    }
    if block_size is not None:
        binary_kwargs["block_size"] = block_size
    with client.open_for_write(
        full_path, fs_mode, write_cfg, binary_kwargs
    ) as raw_handle:
        with self._wrap_text(
            raw_handle,
            mode,
            open_kwargs=open_kwargs,
        ) as wrapped:
            yield wrapped

    version_hint = self._extract_write_version(raw_handle)

    # refresh metadata pinned to the version that was just written
    refreshed = client.get_file_info(self.path, version_id=version_hint)
    for k, v in refreshed.model_dump().items():
        setattr(self, k, v)

def read(self, length: int = -1):
    """Returns file contents."""
    return self.read_bytes(length)

def read_bytes(self, length: int = -1):
    """Returns file contents as bytes."""
    block_size = min(length, PARTIAL_READ_BLOCK_SIZE) if length > 0 else None
    with self.open(mode="rb", block_size=block_size) as stream:
        return stream.read(length)

def read_text(self, **open_kwargs):
    """Return file contents decoded as text.

    **open_kwargs : Any
        Extra keyword arguments forwarded to ``open(mode="r", ...)``
        (e.g. ``encoding="utf-8"``, ``errors="ignore"``)
    """
    with self.open(mode="r", **open_kwargs) as stream:
        return stream.read()

def rebase(
    self,
    old_base: str,
    new_base: str,
    suffix: str = "",
    extension: str = "",
) -> str:
    """
    Rebase the file's URI from one base directory to another.

    Args:
        old_base: Base directory to remove from the file's URI
        new_base: New base directory to prepend
        suffix: Optional suffix to add before file extension
        extension: Optional new file extension (without dot)

    Returns:
        str: Rebased URI with new base directory

    Raises:
        ValueError: If old_base is not found in the file's URI

    Examples:
        ```py
        file = File(source="s3://bucket", path="data/2025-05-27/file.wav")
        file.rebase("s3://bucket/data", "s3://output-bucket/processed",
                    extension="mp3")
        # 's3://output-bucket/processed/2025-05-27/file.mp3'

        file.rebase("data/audio", "/local/output", suffix="_ch1",
                    extension="npy")
        # '/local/output/file_ch1.npy'
        ```
    """
    return rebase_path(self.get_fs_path(), old_base, new_base, suffix, extension)

def resolve(self) -> "Self":
    """
    Resolve a File object by checking its existence and updating its metadata.

    Returns:
        File: The resolved File object with updated metadata.
    """
    if self._catalog is None:
        raise RuntimeError("Cannot resolve file: catalog is not set")

    if self.location:
        raise VFileError(
            "Resolving a virtual file is not supported",
            self.source,
            self.path,
        )

    try:
        client = self._catalog.get_client(self.source)
    except NotImplementedError as e:
        raise RuntimeError(
            f"Unsupported protocol for file source: {self.source}"
        ) from e

    try:
        converted_info = client.get_file_info(self.path)
        # get_file_info always returns a base File; re-wrap as type(self)
        # so resolve() preserves the subclass (TextFile, ImageFile, …).
        res = type(self)(**converted_info.model_dump())
        res._set_stream(self._catalog)
        return res
    except (FileNotFoundError, PermissionError, OSError) as e:
        logger.warning(
            "Error when resolving %s/%s: %s",
            self.source,
            self.path,
            str(e),
        )

    res = type(self)(
        path=self.path,
        source=self.source,
        size=0,
        etag="",
        version="",
        is_latest=True,
        last_modified=TIME_ZERO,
    )
    res._set_stream(self._catalog)
    return res

def save(
    self,
    destination: str,
    client_config: dict | None = None,
    *,
    content_type: str | None = None,
    content_disposition: str | None = None,
    cache_control: str | None = None,
    content_encoding: str | None = None,
    metadata: dict[str, str] | None = None,
    write_options: dict[str, Any] | None = None,
) -> "File":
    """Write file contents to *destination*.

    Args:
        destination: Target path or URI.  Accepts a local OS path, a
            cloud URI (``s3://…``, ``gs://…``, ``az://…``), or an
            unencoded ``file://`` URI as produced by
            :func:`~datachain.fs.utils.path_to_fsspec_uri`.
            Do **not** pass ``Path.as_uri()`` output — that produces
            RFC percent-encoded URIs (e.g. ``file:///my%20dir``) which
            fsspec does not decode, causing ``FileNotFoundError`` for
            paths containing spaces, ``#``, or ``%``.
        client_config: Optional extra kwargs forwarded to the storage
            client (e.g. credentials, endpoint URL).
        content_type: ``Content-Type`` to set on the written object.
        content_disposition: ``Content-Disposition`` to set on the object.
        cache_control: ``Cache-Control`` to set on the object.
        content_encoding: ``Content-Encoding`` to set on the object.
        metadata: Custom key/value metadata to attach to the object.
        write_options: Raw, backend-native write kwargs merged after the
            normalized fields; forwarded on S3 only (e.g.
            ``{"ACL": "public-read"}``) and rejected on GCS/Azure.

    The normalized metadata fields are mapped to each backend's native
    parameters and ignored by the local filesystem.

    Returns:
        A ``File`` representing the newly written destination, with
        ``version`` populated on versioned backends.

    Example:
        ```py
        file.save("/local/output/result.bin")
        file.save("s3://my-bucket/output/result.bin")
        file.save("~/output/result.bin")
        file.save(
            "s3://my-bucket/report.pdf",
            content_type="application/pdf",
            content_disposition='attachment; filename="report.pdf"',
        )
        ```
    """
    if self._catalog is None:
        raise RuntimeError("Cannot save file: catalog is not set")

    destination = stringify_path(destination)
    client, rel_path = self._resolve_destination(destination, client_config)
    write_config = WriteConfig(
        content_type=content_type,
        content_disposition=content_disposition,
        cache_control=cache_control,
        content_encoding=content_encoding,
        metadata=metadata,
        write_options=write_options,
    )
    with self.open(mode="rb") as src:
        result = client.upload(src, rel_path, write_config=write_config)
    result._set_stream(self._catalog)
    return result

@classmethod
def upload(
    cls,
    data: bytes,
    path: str | os.PathLike[str],
    catalog: "Catalog | None" = None,
    *,
    content_type: str | None = None,
    content_disposition: str | None = None,
    cache_control: str | None = None,
    content_encoding: str | None = None,
    metadata: dict[str, str] | None = None,
    write_options: dict[str, Any] | None = None,
) -> "Self":
    """Upload bytes to a storage path and return a File pointing to it.

    Args:
        data: The raw bytes to upload.
        path: Destination path (local or remote, e.g. ``s3://bucket/file.txt``).
        catalog: Optional catalog instance. If None, the current session
            catalog is used.
        content_type: ``Content-Type`` to set on the object.
        content_disposition: ``Content-Disposition`` to set on the object.
        cache_control: ``Cache-Control`` to set on the object.
        content_encoding: ``Content-Encoding`` to set on the object.
        metadata: Custom key/value metadata to attach to the object.
        write_options: Raw, backend-native write kwargs merged after the
            normalized fields; forwarded on S3 only (e.g.
            ``{"ACL": "public-read"}``) and rejected on GCS/Azure.

    The normalized metadata fields are mapped to each backend's native
    parameters and ignored by the local filesystem.

    Returns:
        File: A new File object with metadata populated from the upload.

    Example:
        ```py
        file = File.upload(b"hello world", "s3://bucket/hello.txt")
        file = File.upload(
            pdf_bytes,
            "s3://bucket/report.pdf",
            content_type="application/pdf",
            content_disposition='attachment; filename="report.pdf"',
        )
        ```

    Note:
        To write data as a stream, use `File.at` with `open` instead:
        ```py
        file = File.at("s3://bucket/output.txt")
        with file.open("wb") as f:
            f.write(b"hello world")
        ```
    """
    if catalog is None:
        from datachain.query.session import Session

        catalog = Session.get().catalog
    from datachain.client.fsspec import Client

    path_str = stringify_path(path)

    client_cls = Client.get_implementation(path_str)
    source, rel_path = client_cls.split_url(path_str)

    client = catalog.get_client(client_cls.storage_uri(source))
    file = client.upload(
        data,
        rel_path,
        write_config=WriteConfig(
            content_type=content_type,
            content_disposition=content_disposition,
            cache_control=cache_control,
            content_encoding=content_encoding,
            metadata=metadata,
            write_options=write_options,
        ),
    )
    if not isinstance(file, cls):
        file = cls(**file.model_dump())
    file._set_stream(catalog)
    return file

@classmethod
def open(cls, file: "File", location: list[dict], block_size: int | None = None):
    """Stream file from tar archive based on location in archive.

    Reads always consult the cache (hit → local, miss → remote stream);
    ``_caching_enabled`` only governs whether a miss is also persisted.
    """
    tar_file = cls.parent(file, location)

    loc = location[0]

    if (offset := loc.get("offset", None)) is None:
        raise VFileError("'offset' is not specified", file.source, file.path)

    if (size := loc.get("size", None)) is None:
        raise VFileError("'size' is not specified", file.source, file.path)

    client = file._catalog.get_client(tar_file.source)
    fd = client.open_object(tar_file, use_cache=True, block_size=block_size)
    return FileSlice(fd, offset, size, file.name)