ReMe/reme/steps/file_io/_file_io.py

"""Encoding-aware file IO, output truncation, and per-path write locks."""

import asyncio
from pathlib import Path
from typing import Iterable

import aiofiles
import aiofiles.os

from ...constants import DEFAULT_MAX_BYTES, MAX_FILE_READ_BYTES, TRUNCATION_NOTICE_MARKER
from ...schema import FileChunk
from ...utils import get_logger

logger = get_logger(log_to_file=False)

# ---------------------------------------------------------------------------
# In-process per-path write lock.
# ---------------------------------------------------------------------------
_PATH_LOCKS_MAX = 1024
_PATH_LOCKS: dict[str, asyncio.Lock] = {}
_PATH_LOCKS_REGISTRY = asyncio.Lock()


async def get_path_lock(target: Path) -> asyncio.Lock:
    """Return the asyncio.Lock for ``target``; created lazily on first request."""
    key = str(target)
    async with _PATH_LOCKS_REGISTRY:
        lock = _PATH_LOCKS.get(key)
        if lock is None:
            if len(_PATH_LOCKS) >= _PATH_LOCKS_MAX:
                to_remove = [k for k, v in _PATH_LOCKS.items() if not v.locked()]
                for k in to_remove[: len(_PATH_LOCKS) // 2]:
                    del _PATH_LOCKS[k]
            lock = asyncio.Lock()
            _PATH_LOCKS[key] = lock
    return lock


# ---------------------------------------------------------------------------
# Encoding detection
# ---------------------------------------------------------------------------

_STANDARD_TEXT_EXTS = {
    ".md",
    ".py",
    ".js",
    ".ts",
    ".json",
    ".yaml",
    ".yml",
    ".html",
    ".css",
    ".xml",
    ".log",
    ".conf",
    ".ini",
    ".txt",
    ".sh",
}
_NON_STANDARD_EXTS = {".csv", ".bat", ".cmd", ".reg"}


def _try_decode(data: bytes, encodings: Iterable[str]) -> tuple[str, str] | None:
    """Return ``(text, encoding)`` for the first encoding that decodes ``data`` cleanly."""
    for enc in encodings:
        try:
            return data.decode(enc), enc
        except (UnicodeDecodeError, LookupError):
            continue
    return None


def _decode_known_file(data: bytes, file_extension: str) -> tuple[str, str]:
    """Decode file bytes using the extension as a hint. Returns ``(text, encoding)``."""
    if data.startswith(b"\xef\xbb\xbf"):
        return data.decode("utf-8-sig"), "utf-8-sig"
    if data.startswith((b"\xff\xfe", b"\xfe\xff")):
        try:
            return data.decode("utf-16"), "utf-16"
        except UnicodeDecodeError:
            pass

    ext = (file_extension or "").lower()

    if ext in _STANDARD_TEXT_EXTS:
        try:
            return data.decode("utf-8-sig"), "utf-8"
        except UnicodeDecodeError:
            pass

    if ext in _NON_STANDARD_EXTS:
        result = _try_decode(data, ("utf-8-sig", "gbk"))
        if result is not None:
            text, enc = result
            return text, "utf-8" if enc == "utf-8-sig" else enc

    return data.decode("utf-8", errors="replace"), "utf-8"


# ---------------------------------------------------------------------------
# File read / write
# ---------------------------------------------------------------------------


async def read_file_safe(file_path, max_bytes: int = MAX_FILE_READ_BYTES) -> tuple[str, str]:
    """Read file in byte mode and decode using extension-aware strategy.

    Returns ``(text, encoding)``.
    """
    stat = await aiofiles.os.stat(str(file_path))
    read_size = min(stat.st_size, max_bytes)
    async with aiofiles.open(str(file_path), "rb") as f:
        data = await f.read(read_size)
    return _decode_known_file(data, Path(file_path).suffix)


async def read_file_lines_safe(
    file_path,
    start_line: int,
    end_line: int | None,
    *,
    max_collect_bytes: int = DEFAULT_MAX_BYTES * 2,
) -> tuple[str, int, str]:
    """Read a 1-based inclusive line range without loading the full file.

    Returns ``(text, total_lines, encoding)``.
    """
    encoding = await detect_file_encoding(file_path)
    lines: list[str] = []
    collected_bytes = 0
    total = 0
    async with aiofiles.open(str(file_path), "r", encoding=encoding, errors="replace") as f:
        async for line in f:
            total += 1
            if total >= start_line and (end_line is None or total <= end_line):
                if collected_bytes < max_collect_bytes:
                    cleaned = line.rstrip("\n")
                    lines.append(cleaned)
                    collected_bytes += len(cleaned.encode(encoding, errors="replace")) + 1
    return "\n".join(lines), total, encoding


async def detect_file_encoding(file_path, sniff_bytes: int = 8192) -> str:
    """Detect the encoding of an existing file so writes can preserve it."""
    try:
        async with aiofiles.open(str(file_path), "rb") as f:
            data = await f.read(sniff_bytes)
    except Exception:
        return "utf-8"
    _, enc = _decode_known_file(data, Path(file_path).suffix)
    return enc


async def write_file_safe(file_path: Path, content: str | bytes, encoding: str = "utf-8") -> None:
    """Write ``content`` to ``file_path`` in binary mode; creates parent dirs."""
    file_path.parent.mkdir(parents=True, exist_ok=True)
    if isinstance(content, str):
        try:
            payload = content.encode(encoding)
        except (UnicodeEncodeError, LookupError):
            logger.warning(
                "write_file_safe: %r cannot encode all chars, falling back to utf-8",
                encoding,
            )
            payload = content.encode("utf-8")
    else:
        payload = content
    async with aiofiles.open(str(file_path), "wb") as f:
        await f.write(payload)


# ---------------------------------------------------------------------------
# Output truncation
# ---------------------------------------------------------------------------


def truncate_text_output(
    text: str,
    *,
    start_line: int = 1,
    total_lines: int = 0,
    max_bytes: int = DEFAULT_MAX_BYTES,
    file_path: str | None = None,
    encoding: str = "utf-8",
) -> str:
    """Truncate text by bytes preserving line integrity; append a continuation notice."""
    if not text or max_bytes <= 0:
        return text

    try:
        text_bytes = text.encode(encoding)
        if len(text_bytes) <= max_bytes:
            return text

        truncated = text_bytes[:max_bytes]
        result = truncated.decode(encoding, errors="ignore")
        newline_count = result.count("\n")
        next_line = start_line + max(1, newline_count)

        if next_line <= total_lines:
            read_from = next_line
        elif start_line < total_lines:
            read_from = total_lines
        else:
            return result

        notice = (
            TRUNCATION_NOTICE_MARKER + f"\nThe output above was truncated."
            f"\nThe full content is saved to the file and contains {total_lines} lines in total."
            f"\nThis excerpt starts at line {start_line} and covers the next {max_bytes} bytes."
            f"\nIf the current content is not enough, call `read` with file={file_path or ''} "
            f"start_line={read_from} to read more."
        )
        return result + notice
    except Exception:
        logger.warning("truncate_text_output failed, returning original text", exc_info=True)
        return text


def truncate_session_output(
    text: str,
    *,
    start_line: int = 1,
    total_lines: int = 0,
    max_bytes: int = DEFAULT_MAX_BYTES,
    file_path: str | None = None,
    encoding: str = "utf-8",
) -> str:
    """Render session jsonl lines then truncate by bytes preserving line integrity.

    Each line of *text* is expected to be a serialized ``Msg`` JSON from a raw
    session transcript (``*.jsonl``). Lines are first rendered via
    :func:`~reme.steps.index._source_format.render_session_chunk_lines` — one
    message per line as ``[speaker @ created_at] content`` with internal
    newlines flattened — and then truncated with the same byte-budget logic
    as :func:`truncate_text_output`.
    """
    if not text or max_bytes <= 0:
        return text

    # Render each jsonl line into a human-readable [speaker @ time] content line.
    # Import locally to avoid a hard dependency from file_io on the index package
    # at module load time; the import is safe (no circular dependency).
    from ..index._source_format import render_session_chunk_lines

    chunk = FileChunk(text=text)
    rendered_lines = render_session_chunk_lines(chunk)
    rendered_text = "\n".join(rendered_lines)

    # Reuse the same byte-budget truncation as truncate_text_output.
    return truncate_text_output(
        rendered_text,
        start_line=start_line,
        total_lines=total_lines,
        max_bytes=max_bytes,
        file_path=file_path,
        encoding=encoding,
    )