"""`ContextHubBackend`: Store files in a LangSmith Hub agent repo (persistent)."""

from __future__ import annotations

import fnmatch
import logging
import os
import re
from typing import TYPE_CHECKING

from langsmith import Client
from langsmith.schemas import AgentEntry, FileEntry, SkillEntry
from langsmith.utils import LangSmithError, LangSmithNotFoundError

from deepagents.backends.protocol import (
    FILE_NOT_FOUND,
    INVALID_PATH,
    BackendProtocol,
    DeleteResult,
    EditResult,
    FileData,
    FileDownloadResponse,
    FileInfo,
    FileUploadResponse,
    GlobResult,
    GrepMatch,
    GrepResult,
    LsResult,
    ReadResult,
    WriteResult,
)
from deepagents.backends.utils import (
    create_file_data,
    perform_string_replacement,
    slice_read_response,
)

if TYPE_CHECKING:
    from langsmith.schemas import AgentContext

logger = logging.getLogger(__name__)

# Matches the ":<hash>" suffix appended by langsmith's _build_context_url.
_URL_COMMIT_SUFFIX_RE = re.compile(r":([0-9a-f]{8,64})$")


class ContextHubBackend(BackendProtocol):
    """Backend that stores files in a LangSmith Hub agent repo (persistent)."""

    def __init__(
        self,
        identifier: str,
        *,
        client: Client | None = None,
    ) -> None:
        """Initialize ContextHubBackend.

        Args:
            identifier: Hub agent repo, as `"owner/name"` or `"-/name"`.
            client: LangSmith client. Defaults to `Client()`.
        """
        self._identifier = identifier
        self._client = client if client is not None else Client()
        self._cache: dict[str, str] | None = None
        self._linked_entries: dict[str, str] = {}
        self._commit_hash: str | None = None

    def _load_tree(self) -> None:
        """Fetch the file tree; missing repos are treated as empty."""
        try:
            context: AgentContext = self._client.pull_agent(self._identifier)
        except LangSmithNotFoundError:
            self._cache = {}
            self._linked_entries = {}
            self._commit_hash = None
            return

        self._commit_hash = context.commit_hash
        self._cache = {}
        self._linked_entries = {}

        for path, entry in context.files.items():
            if isinstance(entry, FileEntry):
                self._cache[path] = entry.content
            else:
                self._linked_entries[path] = entry.repo_handle

    def _ensure_cache(self) -> dict[str, str]:
        """Load the file tree if not yet loaded."""
        if self._cache is None:
            self._load_tree()
        if self._cache is None:
            msg = "Context Hub cache failed to initialize"
            raise RuntimeError(msg)
        return self._cache

    def get_linked_entries(self) -> dict[str, str]:
        """Return linked-entry paths mapped to their repo handles."""
        self._ensure_cache()
        return dict(self._linked_entries)

    def has_prior_commits(self) -> bool:
        """Return `True` if the hub repo already exists with at least one commit."""
        self._ensure_cache()
        return self._commit_hash is not None

    def _commit(self, changes: dict[str, str | None]) -> None:
        """Push `changes` as one commit; update the cache on success.

        Each value is either new file content or `None`. A `None` value is the
        deletion marker: relative to `parent_commit`, the server drops that path
        from the tree.
        """
        if not changes:
            return

        payload: dict[str, FileEntry | AgentEntry | SkillEntry | None] = {
            path: FileEntry(type="file", content=content) if content is not None else None for path, content in changes.items()
        }
        url = self._client.push_agent(
            self._identifier,
            files=payload,
            parent_commit=self._commit_hash,
        )
        match = _URL_COMMIT_SUFFIX_RE.search(url)
        if match:
            self._commit_hash = match.group(1)

        if self._cache is not None:
            # Rebuild the cache rather than mutating in place: drop paths whose
            # new content is None (deletions) and overlay the rest as updates.
            deletions = {path for path, content in changes.items() if content is None}
            updates = {path: content for path, content in changes.items() if content is not None}
            self._cache = {path: content for path, content in self._cache.items() if path not in deletions} | updates

    @staticmethod
    def _strip_prefix(path: str) -> str:
        return path.lstrip("/")

    def read(self, file_path: str, offset: int = 0, limit: int = 2000) -> ReadResult:
        """Read file content for the requested line range.

        Args:
            file_path: Absolute file path.
            offset: 0-indexed starting line.
            limit: Maximum number of lines.

        Returns:
            `ReadResult` with raw (unformatted) content.
        """
        hub_path = self._strip_prefix(file_path)
        try:
            cache = self._ensure_cache()
        except LangSmithError as exc:
            logger.exception("Hub pull failed for %r", self._identifier)
            return ReadResult(error=f"Hub unavailable: {exc}")
        content = cache.get(hub_path)
        if content is None:
            return ReadResult(error=f"File '{file_path}' not found")

        file_data = create_file_data(content)
        sliced = slice_read_response(file_data, offset, limit)
        if isinstance(sliced, ReadResult):
            return sliced
        return ReadResult(
            file_data=FileData(
                content=sliced,
                encoding=file_data.get("encoding", "utf-8"),
                created_at=file_data.get("created_at", ""),
                modified_at=file_data.get("modified_at", ""),
            )
        )

    def write(self, file_path: str, content: str) -> WriteResult:
        """Commit `content` to `file_path`."""
        hub_path = self._strip_prefix(file_path)
        try:
            self._ensure_cache()  # populates _commit_hash for parent_commit on push
            self._commit({hub_path: content})
        except LangSmithError as exc:
            logger.exception("Hub write failed for %r", self._identifier)
            self._cache = None
            return WriteResult(error=f"Hub unavailable: {exc}")
        return WriteResult(path=file_path)

    def edit(
        self,
        file_path: str,
        old_string: str,
        new_string: str,
        replace_all: bool = False,  # noqa: FBT001, FBT002
    ) -> EditResult:
        """Replace `old_string` with `new_string` in a file."""
        hub_path = self._strip_prefix(file_path)
        try:
            cache = self._ensure_cache()
            current = cache.get(hub_path)
            if current is None:
                return EditResult(error=f"Error: File '{file_path}' not found")

            result = perform_string_replacement(current, old_string, new_string, replace_all)
            if isinstance(result, str):
                return EditResult(error=result)

            new_content, occurrences = result
            self._commit({hub_path: new_content})
        except LangSmithError as exc:
            logger.exception("Hub edit failed for %r", self._identifier)
            self._cache = None
            return EditResult(error=f"Hub unavailable: {exc}")
        return EditResult(path=file_path, occurrences=occurrences)

    def delete(self, file_path: str) -> DeleteResult:
        """Delete a file or directory by committing its removal from the hub repo.

        Deleting a path removes the exact file plus every nested entry under it
        (the prefix `file_path` + "/"), so a directory is removed recursively.

        Args:
            file_path: Path of the file or directory to delete.

        Returns:
            `DeleteResult` with `file_path` on success, or an error if nothing is
                stored at or under it (or the hub is unavailable).
        """
        hub_path = self._strip_prefix(file_path)
        try:
            cache = self._ensure_cache()
            base = hub_path.rstrip("/")
            prefix = base + "/"
            to_delete = [key for key in cache if key == base or key.startswith(prefix)]
            if not to_delete:
                return DeleteResult(error=f"Error: File '{file_path}' not found")
            self._commit(dict.fromkeys(to_delete, None))
        except LangSmithError as exc:
            logger.exception("Hub delete failed for %r", self._identifier)
            self._cache = None
            return DeleteResult(error=f"Hub unavailable: {exc}")
        return DeleteResult(path=file_path)

    def ls(self, path: str = "/") -> LsResult:
        """List immediate files and subdirectories under `path` (non-recursive)."""
        hub_prefix = self._strip_prefix(path).rstrip("/")
        try:
            cache = self._ensure_cache()
        except LangSmithError as exc:
            logger.exception("Hub pull failed for %r", self._identifier)
            return LsResult(error=f"Hub unavailable: {exc}")

        dirs: set[str] = set()
        entries: list[FileInfo] = []

        for file_path in cache:
            if hub_prefix and not file_path.startswith(hub_prefix + "/"):
                continue

            relative = file_path[len(hub_prefix) + 1 :] if hub_prefix else file_path
            if not relative:
                continue

            parts = relative.split("/", 1)
            if len(parts) == 1:
                entries.append(FileInfo(path=f"/{file_path}", is_dir=False))
            else:
                dir_name = parts[0]
                dir_path = f"{hub_prefix}/{dir_name}" if hub_prefix else dir_name
                if dir_path not in dirs:
                    dirs.add(dir_path)
                    entries.append(FileInfo(path=f"/{dir_path}", is_dir=True))

        return LsResult(entries=entries)

    def grep(
        self,
        pattern: str,
        path: str | None = None,
        glob: str | None = None,
    ) -> GrepResult:
        """Search contents for `pattern` (optional `path` / `glob` filters)."""
        try:
            cache = self._ensure_cache()
        except LangSmithError as exc:
            logger.exception("Hub pull failed for %r", self._identifier)
            return GrepResult(error=f"Hub unavailable: {exc}")
        matches: list[GrepMatch] = []

        try:
            regex = re.compile(pattern)
        except re.error as e:
            return GrepResult(error=f"Invalid regex pattern: {e}")

        prefix = self._strip_prefix(path).rstrip("/") if path else ""

        glob_re = re.compile(fnmatch.translate(os.path.normcase(glob))) if glob else None

        for file_path, content in cache.items():
            if prefix and not file_path.startswith(prefix):
                continue
            if glob_re is not None and not glob_re.match(os.path.normcase(file_path)):
                continue
            for i, line in enumerate(content.splitlines(), start=1):
                if regex.search(line):
                    matches.append(GrepMatch(path=f"/{file_path}", line=i, text=line))

        return GrepResult(matches=matches)

    def glob(self, pattern: str, path: str | None = None) -> GlobResult:  # noqa: ARG002
        """Return files matching `pattern` (`path` unused — flat namespace)."""
        try:
            cache = self._ensure_cache()
        except LangSmithError as exc:
            logger.exception("Hub pull failed for %r", self._identifier)
            return GlobResult(error=f"Hub unavailable: {exc}")
        results: list[FileInfo] = [
            FileInfo(path=f"/{file_path}", is_dir=False)
            for file_path in cache
            if fnmatch.fnmatch(f"/{file_path}", pattern) or fnmatch.fnmatch(file_path, pattern)
        ]
        return GlobResult(matches=results)

    def upload_files(self, files: list[tuple[str, bytes]]) -> list[FileUploadResponse]:
        """Upload text files in one commit; non-UTF-8 inputs rejected per file."""
        # Decode each input; `None` text means we'll reject this entry as invalid.
        decoded: list[tuple[str, str | None]] = []
        valid_files: dict[str, str | None] = {}
        for path, content in files:
            try:
                text = content.decode("utf-8")
            except UnicodeDecodeError:
                decoded.append((path, None))
                continue
            decoded.append((path, text))
            valid_files[self._strip_prefix(path)] = text  # last write wins

        commit_error: str | None = None
        if valid_files:
            try:
                self._ensure_cache()
                self._commit(valid_files)
            except LangSmithError as exc:
                logger.exception("Hub batch upload failed for %r", self._identifier)
                self._cache = None
                commit_error = f"Hub unavailable: {exc}"

        results: list[FileUploadResponse] = []
        for path, text in decoded:
            if text is None:
                results.append(FileUploadResponse(path=path, error=INVALID_PATH))
            elif commit_error is not None:
                # Backend-specific error string per protocol docs
                # (FileOperationError literal union doesn't cover hub failures).
                results.append(FileUploadResponse(path=path, error=commit_error))
            else:
                results.append(FileUploadResponse(path=path))
        return results

    def download_files(self, paths: list[str]) -> list[FileDownloadResponse]:
        """Download files as raw bytes. Missing paths return `file_not_found`."""
        try:
            cache = self._ensure_cache()
        except LangSmithError as exc:
            logger.exception("Hub pull failed for %r", self._identifier)
            # Backend-specific error string per protocol docs (FileOperationError
            # literal union doesn't cover hub failures).
            return [FileDownloadResponse(path=p, error=f"Hub unavailable: {exc}") for p in paths]
        results: list[FileDownloadResponse] = []
        for path in paths:
            hub_path = self._strip_prefix(path)
            content = cache.get(hub_path)
            if content is not None:
                results.append(FileDownloadResponse(path=path, content=content.encode("utf-8")))
            else:
                results.append(FileDownloadResponse(path=path, error=FILE_NOT_FOUND))
        return results
