"""Filesystem tools for the deep agent.

Each tool resolves the backend from ``tool_context.state["_backend"]`` and
delegates to the corresponding backend method. State updates (for
``StateBackend``) are applied via ``tool_context.state``.

Ported from deepagents.middleware.filesystem tool definitions.
"""

from __future__ import annotations

import base64
import os
from typing import cast

from google.adk.tools.tool_context import ToolContext

from adk_deepagents.backends.protocol import Backend, FileData
from adk_deepagents.backends.runtime import get_or_create_backend_for_session
from adk_deepagents.backends.state import StateBackend
from adk_deepagents.backends.utils import (
    format_grep_matches,
    truncate_if_too_long,
    validate_path,
)

# ---------------------------------------------------------------------------
# Image detection
# ---------------------------------------------------------------------------

_IMAGE_EXTENSIONS = frozenset({".png", ".jpg", ".jpeg", ".gif", ".webp"})

_IMAGE_MEDIA_TYPES: dict[str, str] = {
    ".png": "image/png",
    ".jpg": "image/jpeg",
    ".jpeg": "image/jpeg",
    ".gif": "image/gif",
    ".webp": "image/webp",
}


def _is_image_file(file_path: str) -> bool:
    """Return True if *file_path* has a recognised image extension."""
    ext = os.path.splitext(file_path)[1].lower()
    return ext in _IMAGE_EXTENSIONS


# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------


def _get_backend(tool_context: ToolContext) -> Backend:
    """Resolve backend without storing non-serializable objects in state."""
    state_dict = cast(dict, tool_context.state)

    backend = tool_context.state.get("_backend")
    if backend is not None:
        return cast(Backend, backend)

    factory = tool_context.state.get("_backend_factory")
    if callable(factory):
        return cast(Backend, factory(state_dict))

    session = getattr(tool_context, "session", None)
    session_id = getattr(session, "id", None)
    if isinstance(session_id, str) and session_id:
        runtime_backend = get_or_create_backend_for_session(session_id, state_dict)
        if runtime_backend is not None:
            return runtime_backend

    # Safe default for sessions with no custom backend registration.
    return StateBackend(state_dict)


def _apply_files_update(
    tool_context: ToolContext,
    files_update: dict[str, FileData] | None,
) -> None:
    """Merge files_update into session state (for StateBackend)."""
    if files_update is None:
        return
    files = tool_context.state.get("files", {})
    files.update(files_update)
    tool_context.state["files"] = files


# ---------------------------------------------------------------------------
# Tool functions
# ---------------------------------------------------------------------------


def ls(path: str, tool_context: ToolContext) -> dict:
    """List files and directories at the given path.

    Returns name, type (file/dir), size, and modification time for each entry.

    Args:
        path: Absolute path to list (must start with /).
    """
    try:
        validated = validate_path(path)
    except ValueError as e:
        return {"status": "error", "message": str(e)}

    backend = _get_backend(tool_context)
    entries = backend.ls_info(validated)
    return {"status": "success", "entries": entries}


def read_file(
    file_path: str,
    tool_context: ToolContext,
    offset: int = 0,
    limit: int = 2000,
) -> dict:
    """Read a file from the filesystem with optional pagination.

    Returns content with line numbers. For large files, use offset and limit
    to paginate through the content. Image files (.png, .jpg, .jpeg, .gif,
    .webp) are returned as base64-encoded multimodal content.

    Args:
        file_path: Absolute path to the file (must start with /).
        offset: Line number to start reading from (0-based). Defaults to 0.
        limit: Maximum number of lines to return. Defaults to 2000.
    """
    try:
        validated = validate_path(file_path)
    except ValueError as e:
        return {"status": "error", "message": str(e)}

    backend = _get_backend(tool_context)

    if _is_image_file(validated):
        responses = backend.download_files([validated])
        resp = responses[0]
        if resp.error:
            return {"status": "error", "message": f"Error: {resp.error}"}
        ext = os.path.splitext(validated)[1].lower()
        media_type = _IMAGE_MEDIA_TYPES[ext]
        encoded = base64.b64encode(resp.content).decode("ascii")  # type: ignore[arg-type]
        return {
            "status": "success",
            "content": {"type": "image", "media_type": media_type, "data": encoded},
        }

    content = backend.read(validated, offset=offset, limit=limit)

    if content.startswith("Error:"):
        return {"status": "error", "message": content}

    return {"status": "success", "content": content}


def write_file(file_path: str, content: str, tool_context: ToolContext) -> dict:
    """Create a new file with the given content.

    Cannot overwrite existing files. Use edit_file for modifications.

    Args:
        file_path: Absolute path for the new file (must start with /).
        content: The content to write to the file.
    """
    try:
        validated = validate_path(file_path)
    except ValueError as e:
        return {"status": "error", "message": str(e)}

    backend = _get_backend(tool_context)
    result = backend.write(validated, content)

    if result.error:
        if result.error == "already_exists":
            return {
                "status": "error",
                "message": f"File already exists: {validated}. Use edit_file to modify.",
            }
        return {"status": "error", "message": result.error}

    _apply_files_update(tool_context, result.files_update)
    return {"status": "success", "path": result.path}


def edit_file(
    file_path: str,
    old_string: str,
    new_string: str,
    tool_context: ToolContext,
    replace_all: bool = False,
) -> dict:
    """Edit a file by replacing a specific string with a new string.

    The old_string must uniquely identify the text to replace, unless
    replace_all is set to True.

    Args:
        file_path: Absolute path to the file (must start with /).
        old_string: The exact text to find and replace.
        new_string: The replacement text.
        replace_all: If True, replace all occurrences. Defaults to False.
    """
    try:
        validated = validate_path(file_path)
    except ValueError as e:
        return {"status": "error", "message": str(e)}

    backend = _get_backend(tool_context)
    result = backend.edit(validated, old_string, new_string, replace_all)

    if result.error:
        return {"status": "error", "message": str(result.error)}

    _apply_files_update(tool_context, result.files_update)
    return {
        "status": "success",
        "path": result.path,
        "occurrences": result.occurrences,
    }


def glob(pattern: str, tool_context: ToolContext, path: str = "/") -> dict:
    """Find files matching a glob pattern.

    Supports ** for recursive matching and {} for alternatives.

    Args:
        pattern: Glob pattern to match (e.g., "**/*.py", "src/{a,b}/*.ts").
        path: Base directory to search from. Defaults to "/".
    """
    try:
        validated_path = validate_path(path)
    except ValueError as e:
        return {"status": "error", "message": str(e)}

    backend = _get_backend(tool_context)
    entries = backend.glob_info(pattern, validated_path)
    return {"status": "success", "entries": entries}


def grep(
    pattern: str,
    tool_context: ToolContext,
    path: str | None = None,
    glob: str | None = None,
    output_mode: str = "files_with_matches",
) -> dict:
    """Search for a text pattern within files.

    Searches for literal text matches across files. Can filter by path
    and/or glob pattern.

    Args:
        pattern: The text pattern to search for (literal match).
        path: Optional directory to limit the search to.
        glob: Optional glob pattern to filter files (e.g., "*.py").
        output_mode: One of "files_with_matches" (default), "content", or "count".
    """
    validated_path = None
    if path:
        try:
            validated_path = validate_path(path)
        except ValueError as e:
            return {"status": "error", "message": str(e)}

    backend = _get_backend(tool_context)
    raw_result = backend.grep_raw(pattern, validated_path, glob)

    if isinstance(raw_result, str):
        # Already formatted (e.g., from ripgrep backend)
        return {"status": "success", "result": truncate_if_too_long(raw_result)}

    formatted = format_grep_matches(raw_result, output_mode)
    return {"status": "success", "result": truncate_if_too_long(formatted)}
