"""
Deep Agents Remote Backends

S3 backend implementations for LangChain's Deep Agents.
Supports any S3-compatible storage (AWS S3, MinIO, etc.) 
with connection pooling for optimal performance.
"""

from __future__ import annotations

import asyncio
import fnmatch
import json
import re
import threading
from contextlib import asynccontextmanager
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import PurePosixPath
from typing import TYPE_CHECKING, Any, AsyncIterator, Coroutine

import aioboto3
import wcmatch.glob as wcglob
from botocore.config import Config as BotoConfig
from botocore.exceptions import ClientError
from deepagents.backends.protocol import (
    BackendProtocol,
    EditResult,
    FileDownloadResponse,
    FileInfo,
    FileUploadResponse,
    GrepMatch,
    WriteResult,
)
from deepagents.backends.utils import (
    check_empty_content,
    format_content_with_line_numbers,
    perform_string_replacement,
)

if TYPE_CHECKING:
    from types_aiobotocore_s3 import S3Client

__all__ = ["S3Backend", "S3Config"]


def run_async_safely[T](coroutine: Coroutine[Any, Any, T], timeout: float | None = None) -> T:
    try:
        asyncio.get_running_loop()
    except RuntimeError:
        return asyncio.run(coroutine)

    result: dict[str, T] = {}
    error: dict[str, Exception] = {}

    def _run() -> None:
        try:
            result["value"] = asyncio.run(coroutine)
        except Exception as exc:  # noqa: BLE001
            error["value"] = exc

    thread = threading.Thread(target=_run, daemon=True)
    thread.start()
    thread.join(timeout)

    if thread.is_alive():
        raise TimeoutError("Timed out while waiting for coroutine result")
    if "value" in error:
        raise error["value"]

    return result["value"]


# =============================================================================
# S3 Backend (S3-compatible: AWS S3, MinIO, etc.)
# =============================================================================


@dataclass
class S3Config:
    """Configuration for S3-compatible storage."""

    bucket: str
    prefix: str = ""
    region: str = "us-east-1"
    endpoint_url: str | None = None
    access_key_id: str | None = None
    secret_access_key: str | None = None
    use_ssl: bool = True
    max_pool_connections: int = 50
    connect_timeout: float = 5.0
    read_timeout: float = 30.0
    max_retries: int = 3


class S3Backend(BackendProtocol):
    """
    S3-compatible backend for Deep Agents file operations.

    Supports AWS S3, MinIO, and any S3-compatible object storage.
    All operations are async-native using aioboto3.

    Files are stored as objects with paths mapping to S3 keys.
    Content is stored as JSON with the structure:
    {"content": [...lines], "created_at": "...", "modified_at": "..."}
    """

    def __init__(self, config: S3Config) -> None:
        self._config = config
        self._prefix = config.prefix.strip("/")
        if self._prefix:
            self._prefix += "/"

        self._boto_config = BotoConfig(
            region_name=config.region,
            signature_version="s3v4",
            retries={"max_attempts": config.max_retries, "mode": "adaptive"},
            max_pool_connections=config.max_pool_connections,
            connect_timeout=config.connect_timeout,
            read_timeout=config.read_timeout,
        )

        session_kwargs: dict[str, Any] = {}
        if config.access_key_id:
            session_kwargs["aws_access_key_id"] = config.access_key_id
        if config.secret_access_key:
            session_kwargs["aws_secret_access_key"] = config.secret_access_key

        self._session = aioboto3.Session(**session_kwargs)
        self._bucket = config.bucket

    def _s3_key(self, path: str) -> str:
        """Convert virtual path to S3 key."""
        clean = path.lstrip("/")
        return f"{self._prefix}{clean}"

    def _virtual_path(self, key: str) -> str:
        """Convert S3 key to virtual path."""
        if self._prefix and key.startswith(self._prefix):
            key = key[len(self._prefix) :]
        return "/" + key.lstrip("/")

    @asynccontextmanager
    async def _client(self) -> AsyncIterator["S3Client"]:
        """Get S3 client context."""
        async with self._session.client(
            "s3",
            config=self._boto_config,
            endpoint_url=self._config.endpoint_url,
            use_ssl=self._config.use_ssl,
        ) as client:
            yield client

    async def _get_file_data(self, path: str) -> dict[str, Any] | None:
        """Get file data dict from S3."""
        key = self._s3_key(path)
        try:
            async with self._client() as client:
                response = await client.get_object(Bucket=self._bucket, Key=key)
                async with response["Body"] as stream:
                    content = await stream.read()
                return json.loads(content.decode("utf-8"))
        except ClientError as e:
            if e.response["Error"]["Code"] == "NoSuchKey":
                return None
            raise

    async def _put_file_data(
        self, path: str, data: dict[str, Any], *, update_modified: bool = True
    ) -> None:
        """Put file data dict to S3."""
        key = self._s3_key(path)
        if update_modified:
            data["modified_at"] = datetime.now(timezone.utc).isoformat()
        content = json.dumps(data).encode("utf-8")
        async with self._client() as client:
            await client.put_object(
                Bucket=self._bucket,
                Key=key,
                Body=content,
                ContentType="application/json",
            )

    async def _exists(self, path: str) -> bool:
        """Check if file exists in S3."""
        key = self._s3_key(path)
        try:
            async with self._client() as client:
                await client.head_object(Bucket=self._bucket, Key=key)
            return True
        except ClientError as e:
            if e.response["Error"]["Code"] == "404":
                return False
            raise

    async def _list_keys(self, prefix: str = "") -> list[dict[str, Any]]:
        """List all keys with a prefix."""
        full_prefix = self._s3_key(prefix)
        results: list[dict[str, Any]] = []
        async with self._client() as client:
            paginator = client.get_paginator("list_objects_v2")
            async for page in paginator.paginate(
                Bucket=self._bucket, Prefix=full_prefix
            ):
                for obj in page.get("Contents", []):
                    results.append(obj)
        return results

    # -------------------------------------------------------------------------
    # BackendProtocol Implementation
    # -------------------------------------------------------------------------

    def ls_info(self, path: str) -> list[FileInfo]:
        """Sync wrapper for als_info."""
        return run_async_safely(self.als_info(path))

    async def als_info(self, path: str) -> list[FileInfo]:
        """List files in a directory."""
        prefix = path.lstrip("/")
        if prefix and not prefix.endswith("/"):
            prefix += "/"

        objects = await self._list_keys(prefix)
        results: list[FileInfo] = []
        seen_dirs: set[str] = set()

        for obj in objects:
            key = obj["Key"]
            vpath = self._virtual_path(key)

            # Check if this is a direct child or nested
            rel = vpath[len("/" + prefix) :] if prefix else vpath[1:]
            if "/" in rel:
                # This is in a subdirectory, add the directory entry
                dir_name = rel.split("/")[0]
                dir_path = "/" + prefix + dir_name + "/"
                if dir_path not in seen_dirs:
                    seen_dirs.add(dir_path)
                    results.append({"path": dir_path, "is_dir": True})
            else:
                # Direct file
                results.append(
                    {
                        "path": vpath,
                        "is_dir": False,
                        "size": obj.get("Size", 0),
                        "modified_at": obj["LastModified"].isoformat()
                        if "LastModified" in obj
                        else None,
                    }
                )

        results.sort(key=lambda x: x.get("path", ""))
        return results

    def read(self, file_path: str, offset: int = 0, limit: int = 2000) -> str:
        """Sync wrapper for aread."""
        return run_async_safely(
            self.aread(file_path, offset, limit)
        )

    async def aread(self, file_path: str, offset: int = 0, limit: int = 2000) -> str:
        """Read file content with line numbers."""
        data = await self._get_file_data(file_path)
        if data is None:
            return f"Error: File '{file_path}' not found"

        lines = data.get("content", [])
        if not lines:
            empty_msg = check_empty_content("")
            if empty_msg:
                return empty_msg

        if offset >= len(lines):
            return f"Error: Line offset {offset} exceeds file length ({len(lines)} lines)"

        selected = lines[offset : offset + limit]
        return format_content_with_line_numbers(selected, start_line=offset + 1)

    def write(self, file_path: str, content: str) -> WriteResult:
        """Sync wrapper for awrite."""
        return run_async_safely(
            self.awrite(file_path, content)
        )

    async def awrite(self, file_path: str, content: str) -> WriteResult:
        """Create a new file."""
        if await self._exists(file_path):
            return WriteResult(
                error=f"Cannot write to {file_path} because it already exists. "
                "Read and then make an edit, or write to a new path."
            )

        now = datetime.now(timezone.utc).isoformat()
        data = {
            "content": content.splitlines(),
            "created_at": now,
            "modified_at": now,
        }
        try:
            await self._put_file_data(file_path, data, update_modified=False)
            return WriteResult(path=file_path, files_update=None)
        except Exception as e:
            return WriteResult(error=f"Error writing file '{file_path}': {e}")

    def edit(
        self,
        file_path: str,
        old_string: str,
        new_string: str,
        replace_all: bool = False,
    ) -> EditResult:
        """Sync wrapper for aedit."""
        return run_async_safely(
            self.aedit(file_path, old_string, new_string, replace_all)
        )

    async def aedit(
        self,
        file_path: str,
        old_string: str,
        new_string: str,
        replace_all: bool = False,
    ) -> EditResult:
        """Edit file by replacing strings."""
        data = await self._get_file_data(file_path)
        if data is None:
            return EditResult(error=f"Error: File '{file_path}' not found")

        content = "\n".join(data.get("content", []))
        result = perform_string_replacement(content, old_string, new_string, replace_all)

        if isinstance(result, str):
            return EditResult(error=result)

        new_content, occurrences = result
        data["content"] = new_content.splitlines()

        try:
            await self._put_file_data(file_path, data)
            return EditResult(
                path=file_path, files_update=None, occurrences=int(occurrences)
            )
        except Exception as e:
            return EditResult(error=f"Error editing file '{file_path}': {e}")

    def grep_raw(
        self, pattern: str, path: str | None = None, glob: str | None = None
    ) -> list[GrepMatch] | str:
        """Sync wrapper for agrep_raw."""
        return run_async_safely(
            self.agrep_raw(pattern, path, glob)
        )

    async def agrep_raw(
        self, pattern: str, path: str | None = None, glob: str | None = None
    ) -> list[GrepMatch] | str:
        """Search for pattern in files."""
        try:
            regex = re.compile(pattern)
        except re.error as e:
            return f"Invalid regex pattern: {e}"

        search_prefix = (path or "/").lstrip("/")
        objects = await self._list_keys(search_prefix)
        matches: list[GrepMatch] = []

        for obj in objects:
            vpath = self._virtual_path(obj["Key"])
            filename = PurePosixPath(vpath).name

            if glob and not wcglob.globmatch(filename, glob, flags=wcglob.BRACE):
                continue

            data = await self._get_file_data(vpath)
            if data is None:
                continue

            for line_num, line in enumerate(data.get("content", []), 1):
                if regex.search(line):
                    matches.append({"path": vpath, "line": line_num, "text": line})

        return matches

    def glob_info(self, pattern: str, path: str = "/") -> list[FileInfo]:
        """Sync wrapper for aglob_info."""
        return run_async_safely(
            self.aglob_info(pattern, path)
        )

    async def aglob_info(self, pattern: str, path: str = "/") -> list[FileInfo]:
        """Find files matching a glob pattern."""
        search_prefix = path.lstrip("/")
        objects = await self._list_keys(search_prefix)
        results: list[FileInfo] = []

        for obj in objects:
            vpath = self._virtual_path(obj["Key"])
            rel_path = vpath[len(path) :].lstrip("/") if path != "/" else vpath[1:]

            if fnmatch.fnmatch(rel_path, pattern) or fnmatch.fnmatch(vpath, pattern):
                results.append(
                    {
                        "path": vpath,
                        "is_dir": False,
                        "size": obj.get("Size", 0),
                        "modified_at": obj["LastModified"].isoformat()
                        if "LastModified" in obj
                        else None,
                    }
                )

        results.sort(key=lambda x: x.get("path", ""))
        return results

    def upload_files(self, files: list[tuple[str, bytes]]) -> list[FileUploadResponse]:
        """Sync wrapper for aupload_files."""
        return run_async_safely(self.aupload_files(files))

    async def aupload_files(
        self, files: list[tuple[str, bytes]]
    ) -> list[FileUploadResponse]:
        """Upload multiple files."""
        responses: list[FileUploadResponse] = []
        async with self._client() as client:
            for path, content in files:
                try:
                    key = self._s3_key(path)
                    await client.put_object(
                        Bucket=self._bucket, Key=key, Body=content
                    )
                    responses.append(FileUploadResponse(path=path, error=None))
                except ClientError as e:
                    code = e.response["Error"]["Code"]
                    if code == "AccessDenied":
                        responses.append(
                            FileUploadResponse(path=path, error="permission_denied")
                        )
                    else:
                        responses.append(
                            FileUploadResponse(path=path, error="invalid_path")
                        )
                except Exception:
                    responses.append(
                        FileUploadResponse(path=path, error="invalid_path")
                    )
        return responses

    def download_files(self, paths: list[str]) -> list[FileDownloadResponse]:
        """Sync wrapper for adownload_files."""
        return run_async_safely(self.adownload_files(paths))

    async def adownload_files(self, paths: list[str]) -> list[FileDownloadResponse]:
        """Download multiple files."""
        responses: list[FileDownloadResponse] = []
        async with self._client() as client:
            for path in paths:
                try:
                    key = self._s3_key(path)
                    response = await client.get_object(Bucket=self._bucket, Key=key)
                    async with response["Body"] as stream:
                        content = await stream.read()
                    responses.append(
                        FileDownloadResponse(path=path, content=content, error=None)
                    )
                except ClientError as e:
                    code = e.response["Error"]["Code"]
                    if code == "NoSuchKey":
                        responses.append(
                            FileDownloadResponse(
                                path=path, content=None, error="file_not_found"
                            )
                        )
                    elif code == "AccessDenied":
                        responses.append(
                            FileDownloadResponse(
                                path=path, content=None, error="permission_denied"
                            )
                        )
                    else:
                        responses.append(
                            FileDownloadResponse(
                                path=path, content=None, error="invalid_path"
                            )
                        )
        return responses