"""
WebSocket client for AiCippy.

Provides real-time bidirectional communication with the backend
via API Gateway WebSocket.
"""

from __future__ import annotations

import asyncio
import hashlib
import json
import math
import random
from contextlib import asynccontextmanager, suppress
from pathlib import Path
from typing import TYPE_CHECKING, Any

import websockets
from websockets.exceptions import ConnectionClosed, WebSocketException

from aicippy.auth.cognito import CognitoAuth
from aicippy.config import get_settings
from aicippy.exceptions import WebSocketDisconnectedError, WebSocketMessageError
from aicippy.utils.correlation import get_correlation_id
from aicippy.utils.logging import get_logger
from aicippy.websocket.models import (
    AgentUpdate,
    FileOperationType,
    FileTransferMessage,
    FileTransferProgress,
    MessageType,
    ToolOutput,
    WebSocketMessage,
)

if TYPE_CHECKING:
    from collections.abc import AsyncIterator, Callable

    from websockets.client import WebSocketClientProtocol

    from aicippy.platform.tenant_context import TenantContext

logger = get_logger(__name__)


class WebSocketClient:
    """
    WebSocket client for real-time communication.

    Handles connection management, authentication, reconnection,
    and message routing.
    """

    def __init__(
        self,
        tenant_ctx: TenantContext | None = None,
        on_reconnect_status: Callable[..., None] | None = None,
    ) -> None:
        """Initialize WebSocket client.

        Args:
            tenant_ctx: Optional tenant context for tenant-scoped
                connections. When provided, all outgoing messages are
                stamped with ``tenant_id`` and incoming messages are
                validated against it.
            on_reconnect_status: Optional callback invoked during reconnect
                attempts with status tuples such as
                ``("reconnecting", attempt, max_attempts)``,
                ``("reconnected",)``, or ``("reconnect_failed",)``.
        """
        self._settings = get_settings()
        self._auth = CognitoAuth()
        self._tenant_ctx = tenant_ctx
        self._on_reconnect_status = on_reconnect_status
        self._ws: WebSocketClientProtocol | None = None
        self._connected = False
        self._intentional_disconnect = False
        self._reconnect_task: asyncio.Task[None] | None = None
        self._receive_task: asyncio.Task[None] | None = None
        self._message_handlers: dict[MessageType, list[Callable[[dict[str, Any]], None]]] = {}
        self._pending_responses: dict[str, asyncio.Future[dict[str, Any]]] = {}
        self._connect_lock = asyncio.Lock()
        self._download_queue: asyncio.Queue[bytes] | None = None

    @property
    def is_connected(self) -> bool:
        """Check if WebSocket is connected."""
        return self._connected and self._ws is not None and not self._ws.closed

    async def connect(self) -> bool:
        """
        Establish WebSocket connection.

        Returns:
            True if connection successful, False otherwise.
        """
        async with self._connect_lock:
            if self.is_connected:
                return True

            # Reset intentional disconnect flag on new connection attempt
            self._intentional_disconnect = False

            # Get authentication token
            tokens = self._auth.get_current_tokens()
            if not tokens:
                logger.error("websocket_no_auth_token")
                return False

            url = self._settings.websocket_url
            logger.info("websocket_connecting", url=url)

            try:
                headers: dict[str, str] = {
                    "Authorization": f"Bearer {tokens.access_token}",
                }
                if self._tenant_ctx is not None:
                    headers["X-Tenant-ID"] = self._tenant_ctx.tenant_id

                self._ws = await websockets.connect(
                    url,
                    additional_headers=headers,
                    ping_interval=30,
                    ping_timeout=10,
                    close_timeout=5,
                )

                self._connected = True
                logger.info("websocket_connected")

                # Start receive loop
                self._receive_task = asyncio.create_task(self._receive_loop())

                # Authenticate
                await self._authenticate(tokens.access_token)

                return True

            except (WebSocketException, OSError) as e:
                logger.error("websocket_connection_failed", error=str(e))
                self._connected = False
                return False

    async def disconnect(self) -> None:
        """Close WebSocket connection."""
        self._intentional_disconnect = True
        self._connected = False

        async with self._connect_lock:
            await self._disconnect_inner()

    async def _disconnect_inner(self) -> None:
        """Inner disconnect logic (caller must hold _connect_lock or set flags)."""
        if self._receive_task:
            self._receive_task.cancel()
            with suppress(asyncio.CancelledError):
                await self._receive_task
            self._receive_task = None

        if self._reconnect_task:
            self._reconnect_task.cancel()
            with suppress(asyncio.CancelledError):
                await self._reconnect_task
            self._reconnect_task = None

        if self._ws:
            try:
                await self._ws.close()
            except Exception as e:
                logger.debug("websocket_close_error", error=str(e))
            finally:
                self._ws = None

        # Reject all pending response futures so callers are not left hanging
        for _correlation_id, future in self._pending_responses.items():
            if not future.done():
                future.set_exception(
                    WebSocketDisconnectedError(
                        message="WebSocket disconnected while awaiting response.",
                    )
                )
        self._pending_responses.clear()

        logger.info("websocket_disconnected")

    async def _authenticate(self, token: str) -> None:
        """
        Send authentication message.

        Args:
            token: Access token for authentication.
        """
        payload: dict[str, Any] = {"token": token}
        if self._tenant_ctx is not None:
            payload["tenant_id"] = self._tenant_ctx.tenant_id
            payload["user_id"] = self._tenant_ctx.user_id
        await self.send(
            MessageType.AUTHENTICATE,
            payload,
        )

    async def send(
        self,
        message_type: MessageType,
        payload: dict[str, Any],
        correlation_id: str | None = None,
    ) -> None:
        """
        Send a message through the WebSocket.

        Args:
            message_type: Type of message to send.
            payload: Message payload.
            correlation_id: Optional correlation ID for tracking.
        """
        if not self.is_connected:
            raise WebSocketDisconnectedError(
                message="WebSocket not connected. Call connect() first.",
            )

        message = WebSocketMessage(
            type=message_type,
            payload=payload,
            correlation_id=correlation_id or get_correlation_id(),
            tenant_id=self._tenant_ctx.tenant_id if self._tenant_ctx else None,
        )

        try:
            ws = self._ws
            if ws is None or ws.closed:
                raise WebSocketDisconnectedError(
                    message="WebSocket not available",
                )
            await ws.send(json.dumps(message.to_dict()))
            logger.debug("websocket_message_sent", type=message_type.value)
        except ConnectionClosed as exc:
            logger.warning("websocket_send_connection_closed")
            self._connected = False
            raise WebSocketDisconnectedError(
                message="WebSocket connection closed during send.",
            ) from exc

    async def send_and_wait(
        self,
        message_type: MessageType,
        payload: dict[str, Any],
        timeout: float = 30.0,
    ) -> dict[str, Any]:
        """
        Send a message and wait for response.

        Args:
            message_type: Type of message to send.
            payload: Message payload.
            timeout: Maximum time to wait for response.

        Returns:
            Response payload.
        """
        correlation_id = get_correlation_id() or str(id(payload))
        future: asyncio.Future[dict[str, Any]] = asyncio.get_running_loop().create_future()
        self._pending_responses[correlation_id] = future

        try:
            await self.send(message_type, payload, correlation_id)
            return await asyncio.wait_for(future, timeout=timeout)
        except TimeoutError as exc:
            logger.warning("websocket_response_timeout", correlation_id=correlation_id)
            raise TimeoutError("Response timeout") from exc
        finally:
            self._pending_responses.pop(correlation_id, None)

    async def _receive_loop(self) -> None:
        """Background loop to receive messages."""
        max_consecutive_errors = 10
        consecutive_errors = 0

        while self._connected and self._ws:
            try:
                message = await self._ws.recv()

                # B4: Route binary messages to download queue if active
                if isinstance(message, bytes) and self._download_queue is not None:
                    await self._download_queue.put(message)
                    consecutive_errors = 0
                    continue

                await self._handle_message(message)
                consecutive_errors = 0
            except ConnectionClosed:
                logger.warning("websocket_connection_closed")
                self._connected = False
                # B1: Do not reconnect on intentional disconnect
                if self._intentional_disconnect:
                    break
                # A1: Prevent double-reconnect race
                if self._reconnect_task is None or self._reconnect_task.done():
                    self._reconnect_task = asyncio.create_task(self._attempt_reconnect())
                break
            except asyncio.CancelledError:
                raise
            except Exception as e:
                consecutive_errors += 1
                logger.exception(
                    "websocket_receive_error",
                    error=str(e),
                    consecutive_errors=consecutive_errors,
                )
                # B5: Break on too many consecutive failures to avoid infinite spin
                if consecutive_errors >= max_consecutive_errors:
                    logger.error(
                        "websocket_receive_max_consecutive_errors",
                        max_errors=max_consecutive_errors,
                    )
                    self._connected = False
                    break
                continue

    async def _handle_message(self, raw_message: str | bytes) -> None:
        """
        Handle incoming WebSocket message.

        Args:
            raw_message: Raw message data from WebSocket.
        """
        try:
            if isinstance(raw_message, bytes):
                raw_message = raw_message.decode("utf-8")

            data = json.loads(raw_message)
            message = WebSocketMessage.from_dict(data)

            logger.debug("websocket_message_received", type=message.type.value)

            # Reject messages from a different tenant
            if (
                self._tenant_ctx is not None
                and message.tenant_id
                and message.tenant_id != self._tenant_ctx.tenant_id
            ):
                logger.error(
                    "websocket_tenant_mismatch",
                    expected=self._tenant_ctx.tenant_id,
                    received=message.tenant_id,
                )
                return

            # Check for pending response
            if message.correlation_id and message.correlation_id in self._pending_responses:
                future = self._pending_responses.pop(message.correlation_id)
                if not future.done():
                    try:
                        future.set_result(message.payload)
                    except asyncio.InvalidStateError:
                        logger.debug(
                            "websocket_future_already_resolved",
                            correlation_id=message.correlation_id,
                        )
                return

            # Route to registered handlers
            handlers = self._message_handlers.get(message.type, [])
            for handler in handlers:
                try:
                    handler(message.payload)
                except Exception as e:
                    logger.exception("websocket_handler_error", error=str(e))

        except json.JSONDecodeError as e:
            logger.warning("websocket_invalid_json", error=str(e))
        except (KeyError, ValueError) as e:
            logger.warning("websocket_invalid_message_format", error=str(e))

    def on_message(
        self,
        message_type: MessageType,
        handler: Callable[[dict[str, Any]], None],
    ) -> None:
        """
        Register a message handler.

        Args:
            message_type: Type of message to handle.
            handler: Callback function for the message.
        """
        if message_type not in self._message_handlers:
            self._message_handlers[message_type] = []
        self._message_handlers[message_type].append(handler)

    def remove_handler(
        self,
        message_type: MessageType,
        handler: Callable[[dict[str, Any]], None],
    ) -> None:
        """
        Remove a message handler.

        Args:
            message_type: Type of message.
            handler: Handler to remove.
        """
        if message_type in self._message_handlers:
            self._message_handlers[message_type] = [
                h for h in self._message_handlers[message_type] if h != handler
            ]

    async def _attempt_reconnect(self) -> None:
        """Attempt to reconnect with exponential backoff (capped at 60s)."""
        # B7: Reject all pending futures at the start of reconnect
        for _cid, future in self._pending_responses.items():
            if not future.done():
                future.set_exception(
                    WebSocketDisconnectedError(
                        message="WebSocket reconnecting; pending request rejected.",
                    )
                )
        self._pending_responses.clear()

        max_attempts = self._settings.websocket_reconnect_attempts
        base_delay = self._settings.websocket_reconnect_delay
        max_delay = 60.0  # Cap backoff at 60 seconds

        for attempt in range(1, max_attempts + 1):
            # A4: Add random jitter (decorrelated)
            delay = min(base_delay * (2 ** (attempt - 1)), max_delay)
            delay = random.uniform(delay * 0.5, delay)  # noqa: S311
            logger.info(
                "websocket_reconnecting",
                attempt=attempt,
                max_attempts=max_attempts,
                delay=round(delay, 2),
            )

            # A5: Notify via callback
            if self._on_reconnect_status:
                try:
                    self._on_reconnect_status("reconnecting", attempt, max_attempts)
                except Exception as exc:
                    logger.debug("websocket_reconnect_status_callback_error", error=str(exc))

            await asyncio.sleep(delay)

            # A3: Try to refresh tokens before reconnecting
            try:
                if hasattr(self._auth, "refresh_tokens"):
                    await self._auth.refresh_tokens()
            except Exception as exc:
                logger.warning(
                    "websocket_token_refresh_failed_before_reconnect",
                    error=str(exc),
                )

            if await self.connect():
                logger.info("websocket_reconnected", attempt=attempt)
                # A5: Notify success
                if self._on_reconnect_status:
                    try:
                        self._on_reconnect_status(
                            "reconnected",
                        )
                    except Exception:
                        logger.debug("websocket_reconnect_status_callback_error")
                return

        logger.error("websocket_reconnect_failed", attempts=max_attempts)

        # A5: Notify failure
        if self._on_reconnect_status:
            try:
                self._on_reconnect_status(
                    "reconnect_failed",
                )
            except Exception:
                logger.debug("websocket_reconnect_status_callback_error")

        # A7: Clean up after max retries exhausted
        await self.disconnect()

    @asynccontextmanager
    async def session(self) -> AsyncIterator[WebSocketClient]:
        """
        Context manager for WebSocket session.

        Usage:
            async with client.session():
                await client.send(...)
        """
        try:
            connected = await self.connect()
            if not connected:
                raise ConnectionError("Failed to connect WebSocket")
            yield self
        finally:
            await self.disconnect()

    async def stream_chat(
        self,
        message: str,
        on_token: Callable[[str], None] | None = None,
        on_agent_update: Callable[[AgentUpdate], None] | None = None,
        on_tool_output: Callable[[ToolOutput], None] | None = None,
    ) -> str:
        """
        Stream a chat response with callbacks for tokens and updates.

        Args:
            message: User message to send.
            on_token: Callback for each response token.
            on_agent_update: Callback for agent status updates.
            on_tool_output: Callback for tool execution outputs.

        Returns:
            Complete response text.
        """
        response_parts: list[str] = []
        complete = asyncio.Event()
        stream_error: list[str] = []

        def handle_response(payload: dict[str, Any]) -> None:
            token = payload.get("token", "")
            if token:
                response_parts.append(token)
                if on_token:
                    on_token(token)

            if payload.get("complete", False):
                complete.set()

        def handle_agent_update(payload: dict[str, Any]) -> None:
            if on_agent_update:
                on_agent_update(AgentUpdate.from_dict(payload))

        def handle_tool_output(payload: dict[str, Any]) -> None:
            if on_tool_output:
                on_tool_output(ToolOutput.from_dict(payload))

        def handle_error(payload: dict[str, Any]) -> None:
            error_msg = payload.get("message", payload.get("error", "Unknown server error"))
            stream_error.append(str(error_msg))
            complete.set()

        # Register handlers
        self.on_message(MessageType.RESPONSE, handle_response)
        self.on_message(MessageType.AGENT_UPDATE, handle_agent_update)
        self.on_message(MessageType.TOOL_OUTPUT, handle_tool_output)
        self.on_message(MessageType.ERROR, handle_error)

        try:
            chat_payload: dict[str, Any] = {"message": message}
            if self._tenant_ctx is not None:
                chat_payload["tenant_id"] = self._tenant_ctx.tenant_id
                chat_payload["user_id"] = self._tenant_ctx.user_id
            await self.send(MessageType.CHAT, chat_payload)
            try:
                await asyncio.wait_for(complete.wait(), timeout=300.0)
            except TimeoutError as exc:
                raise WebSocketMessageError(
                    message="Streaming response timed out after 300 seconds.",
                ) from exc
            if stream_error:
                raise WebSocketMessageError(
                    message=f"Server error during streaming: {stream_error[0]}",
                )
            return "".join(response_parts)
        finally:
            self.remove_handler(MessageType.RESPONSE, handle_response)
            self.remove_handler(MessageType.AGENT_UPDATE, handle_agent_update)
            self.remove_handler(MessageType.TOOL_OUTPUT, handle_tool_output)
            self.remove_handler(MessageType.ERROR, handle_error)

    # ========================================================================
    # File Operations
    # ========================================================================

    def _ensure_connected(self) -> None:
        """Verify WebSocket is connected, raise if not."""
        if not self.is_connected:
            raise WebSocketDisconnectedError(
                message="WebSocket is not connected. Call connect() first.",
            )

    @staticmethod
    def _compute_file_checksum(file_path: Path) -> str:
        """
        Compute SHA-256 checksum of a file.

        Args:
            file_path: Path to the file.

        Returns:
            Hex digest of the SHA-256 checksum.
        """
        sha256 = hashlib.sha256()
        with Path(file_path).open("rb") as f:
            while True:
                chunk = f.read(8192)
                if not chunk:
                    break
                sha256.update(chunk)
        return sha256.hexdigest()

    async def upload_file(
        self,
        file_path: str | Path,
        remote_path: str = "",
        chunk_size: int = 64 * 1024,  # 64KB chunks
        progress_callback: Callable[[FileTransferProgress], None] | None = None,
    ) -> bool:
        """
        Upload a file via WebSocket in chunks.

        Reads the local file, splits it into chunks, and sends each chunk
        as a binary WebSocket frame with a JSON header prefix. Progress is
        tracked and reported through the optional callback.

        Args:
            file_path: Local path to the file to upload.
            remote_path: Remote destination path. Defaults to the file name.
            chunk_size: Size of each chunk in bytes (default 64KB).
            progress_callback: Optional callback for progress updates.

        Returns:
            True if upload completed successfully, False otherwise.
        """
        self._ensure_connected()

        path = Path(file_path)
        if not path.exists():
            logger.error("upload_file_not_found", path=str(path))
            return False

        if not path.is_file():
            logger.error("upload_path_not_file", path=str(path))
            return False

        file_size = path.stat().st_size
        file_name = path.name
        dest_path = remote_path or file_name
        # Prefix remote path with tenant namespace for isolation
        if self._tenant_ctx is not None and not dest_path.startswith(
            f"{self._tenant_ctx.tenant_id}/"
        ):
            dest_path = f"{self._tenant_ctx.tenant_id}/{dest_path}"
        total_chunks = max(1, math.ceil(file_size / chunk_size))
        checksum = self._compute_file_checksum(path)

        progress = FileTransferProgress(
            file_name=file_name,
            total_bytes=file_size,
            transferred_bytes=0,
            status="pending",
        )

        logger.info(
            "upload_file_start",
            file_name=file_name,
            file_size=file_size,
            total_chunks=total_chunks,
        )

        try:
            with path.open("rb") as f:
                for chunk_index in range(total_chunks):
                    data = f.read(chunk_size)
                    if not data:
                        break

                    transfer_msg = FileTransferMessage(
                        operation=FileOperationType.UPLOAD,
                        file_path=dest_path,
                        file_name=file_name,
                        file_size=file_size,
                        chunk_index=chunk_index,
                        total_chunks=total_chunks,
                        data=data,
                        checksum=checksum if chunk_index == total_chunks - 1 else "",
                    )

                    # Build the wire frame: JSON header + null byte separator + binary data
                    header_json = json.dumps(transfer_msg.to_dict())
                    frame = header_json.encode("utf-8") + b"\x00" + data

                    try:
                        ws = self._ws
                        if ws is None or ws.closed:
                            raise WebSocketDisconnectedError(
                                message="WebSocket not available during upload",
                            )
                        await ws.send(frame)
                    except ConnectionClosed as exc:
                        logger.warning(
                            "upload_connection_lost",
                            chunk_index=chunk_index,
                            total_chunks=total_chunks,
                        )
                        self._connected = False
                        progress.status = "error"
                        progress.error = "Connection lost during upload"
                        if progress_callback:
                            progress_callback(progress)
                        raise WebSocketDisconnectedError(
                            message=(
                                "Connection lost during upload"
                                f" at chunk {chunk_index}"
                                f"/{total_chunks}"
                            ),
                        ) from exc

                    progress.transferred_bytes += len(data)
                    progress.status = "transferring"
                    if progress_callback:
                        progress_callback(progress)

                    logger.debug(
                        "upload_chunk_sent",
                        chunk=chunk_index + 1,
                        total=total_chunks,
                        progress_pct=f"{progress.progress_pct:.1f}%",
                    )

            progress.status = "complete"
            if progress_callback:
                progress_callback(progress)

            logger.info(
                "upload_file_complete",
                file_name=file_name,
                file_size=file_size,
                checksum=checksum,
            )
            return True

        except WebSocketDisconnectedError:
            raise
        except Exception as e:
            logger.exception("upload_file_error", file_name=file_name, error=str(e))
            progress.status = "error"
            progress.error = str(e)
            if progress_callback:
                progress_callback(progress)
            return False

    async def download_file(
        self,
        remote_path: str,
        local_path: str | Path,
        progress_callback: Callable[[FileTransferProgress], None] | None = None,
    ) -> bool:
        """
        Download a file via WebSocket.

        Sends a download request and receives file data in chunks.
        Each received chunk is written to the local file and progress
        is tracked. The final checksum is verified upon completion.

        Args:
            remote_path: Remote path of the file to download.
            local_path: Local destination path.
            progress_callback: Optional callback for progress updates.

        Returns:
            True if download completed successfully, False otherwise.
        """
        self._ensure_connected()

        dest = Path(local_path)
        dest.parent.mkdir(parents=True, exist_ok=True)
        file_name = dest.name

        progress = FileTransferProgress(
            file_name=file_name,
            total_bytes=0,
            transferred_bytes=0,
            status="pending",
        )

        logger.info("download_file_start", remote_path=remote_path, local_path=str(dest))

        download_complete = asyncio.Event()
        download_error: list[str] = []
        received_checksum: list[str] = []

        try:
            # Send download request
            correlation_id = get_correlation_id() or str(id(remote_path))
            await self.send(
                MessageType.COMMAND,
                {
                    "command": FileOperationType.DOWNLOAD.value,
                    "remote_path": remote_path,
                    "local_file_name": file_name,
                },
                correlation_id=correlation_id,
            )

            sha256 = hashlib.sha256()
            if self._ws is None:
                raise WebSocketDisconnectedError(
                    message="WebSocket not available for download",
                )

            # B4: Register a download queue so _receive_loop routes
            # binary messages here instead of stealing them via ws.recv()
            self._download_queue = asyncio.Queue()

            # Also register a handler for JSON completion/error signals
            download_json_queue: asyncio.Queue[dict[str, Any]] = asyncio.Queue()

            def _download_json_handler(payload: dict[str, Any]) -> None:
                download_json_queue.put_nowait(payload)

            self.on_message(MessageType.RESPONSE, _download_json_handler)
            self.on_message(MessageType.ERROR, _download_json_handler)

            with dest.open("wb") as f:
                while not download_complete.is_set():
                    try:
                        # Wait for binary data from the download queue
                        raw = await asyncio.wait_for(
                            self._download_queue.get(),
                            timeout=60.0,
                        )
                    except TimeoutError:
                        # Check if a JSON signal arrived instead
                        try:
                            data = download_json_queue.get_nowait()
                            if data.get("type") == "error" or data.get("error"):
                                download_error.append(
                                    data.get("message", data.get("error", "Unknown download error"))
                                )
                                break
                            if data.get("complete", False):
                                if data.get("checksum"):
                                    received_checksum.append(data["checksum"])
                                download_complete.set()
                                continue
                        except asyncio.QueueEmpty:
                            pass
                        download_error.append("Download timed out waiting for data")
                        break

                    # Handle binary frames (file chunks)
                    separator_idx = raw.find(b"\x00")
                    if separator_idx == -1:
                        # Pure binary data without header
                        chunk_data = raw
                        header: dict[str, Any] = {}
                    else:
                        header_bytes = raw[:separator_idx]
                        chunk_data = raw[separator_idx + 1 :]
                        header = json.loads(header_bytes.decode("utf-8"))

                    f.write(chunk_data)
                    sha256.update(chunk_data)

                    if progress.total_bytes == 0 and header.get("file_size", 0) > 0:
                        progress.total_bytes = header["file_size"]

                    progress.transferred_bytes += len(chunk_data)
                    progress.status = "transferring"
                    if progress_callback:
                        progress_callback(progress)

                    # Check for final chunk checksum
                    chunk_checksum = header.get("checksum", "")
                    if chunk_checksum:
                        received_checksum.append(chunk_checksum)
                        download_complete.set()

                    logger.debug(
                        "download_chunk_received",
                        bytes_received=progress.transferred_bytes,
                        progress_pct=f"{progress.progress_pct:.1f}%",
                    )

            # Clean up download queue and handlers
            self._download_queue = None
            self.remove_handler(MessageType.RESPONSE, _download_json_handler)
            self.remove_handler(MessageType.ERROR, _download_json_handler)

            if download_error:
                progress.status = "error"
                progress.error = download_error[0]
                if progress_callback:
                    progress_callback(progress)
                logger.error("download_file_error", error=download_error[0])
                return False

            # Verify checksum if provided
            local_checksum = sha256.hexdigest()
            if received_checksum and received_checksum[0] != local_checksum:
                progress.status = "error"
                progress.error = (
                    f"Checksum mismatch: expected {received_checksum[0]}, got {local_checksum}"
                )
                if progress_callback:
                    progress_callback(progress)
                logger.error(
                    "download_checksum_mismatch",
                    expected=received_checksum[0],
                    actual=local_checksum,
                )
                return False

            progress.status = "complete"
            if progress_callback:
                progress_callback(progress)

            logger.info(
                "download_file_complete",
                file_name=file_name,
                bytes_received=progress.transferred_bytes,
                checksum=local_checksum,
            )
            return True

        except WebSocketDisconnectedError:
            self._download_queue = None
            raise
        except Exception as e:
            self._download_queue = None
            logger.exception("download_file_error", remote_path=remote_path, error=str(e))
            progress.status = "error"
            progress.error = str(e)
            if progress_callback:
                progress_callback(progress)
            return False

    async def sync_files(
        self,
        local_dir: str | Path,
        remote_dir: str = "",
        patterns: list[str] | None = None,
        progress_callback: Callable[[FileTransferProgress], None] | None = None,
    ) -> list[FileTransferProgress]:
        """
        Sync local files to remote via WebSocket.

        Scans the local directory for files matching the given glob patterns
        (or all files if no patterns specified), computes checksums, compares
        with the remote file listing, and uploads any new or changed files.

        Args:
            local_dir: Local directory to sync.
            remote_dir: Remote destination directory.
            patterns: Optional glob patterns to filter files (e.g. ["*.py", "*.txt"]).
            progress_callback: Optional callback for progress updates per file.

        Returns:
            List of FileTransferProgress results, one per file attempted.
        """
        self._ensure_connected()

        local = Path(local_dir)
        if not local.is_dir():
            logger.error("sync_path_not_directory", path=str(local))
            return []

        # Collect local files matching patterns
        local_files: list[Path] = []
        if patterns:
            for pattern in patterns:
                local_files.extend(local.rglob(pattern))
        else:
            local_files = [p for p in local.rglob("*") if p.is_file()]

        if not local_files:
            logger.info("sync_no_files_found", path=str(local), patterns=patterns)
            return []

        # Compute local checksums
        local_checksums: dict[str, str] = {}
        for file_path in local_files:
            relative = file_path.relative_to(local)
            local_checksums[str(relative)] = self._compute_file_checksum(file_path)

        # Get remote file listing for comparison
        remote_files_info: dict[str, str] = {}
        try:
            remote_listing = await self.list_remote_files(remote_dir)
            for entry in remote_listing:
                name = entry.get("name", "")
                checksum = entry.get("checksum", "")
                if name and checksum:
                    remote_files_info[name] = checksum
        except Exception as e:
            logger.warning("sync_remote_list_failed", error=str(e))
            # Proceed with uploading all files if remote listing fails

        # Determine files that need uploading (new or changed)
        files_to_upload: list[Path] = []
        for file_path in local_files:
            relative = str(file_path.relative_to(local))
            remote_checksum = remote_files_info.get(relative, "")
            if remote_checksum != local_checksums.get(relative, ""):
                files_to_upload.append(file_path)

        logger.info(
            "sync_files_start",
            total_local=len(local_files),
            to_upload=len(files_to_upload),
        )

        results: list[FileTransferProgress] = []
        for file_path in files_to_upload:
            relative = str(file_path.relative_to(local))
            remote_file_path = str(Path(remote_dir) / relative) if remote_dir else relative

            file_progress = FileTransferProgress(
                file_name=relative,
                total_bytes=file_path.stat().st_size,
                status="pending",
            )

            def _make_sync_callback(
                fp: FileTransferProgress,
            ) -> Callable[[FileTransferProgress], None]:
                """Create a closure to forward per-file progress."""

                def _inner(chunk_progress: FileTransferProgress) -> None:
                    fp.transferred_bytes = chunk_progress.transferred_bytes
                    fp.status = chunk_progress.status
                    fp.error = chunk_progress.error
                    if progress_callback:
                        progress_callback(fp)

                return _inner

            success = await self.upload_file(
                file_path=file_path,
                remote_path=remote_file_path,
                progress_callback=_make_sync_callback(file_progress),
            )

            if not success and file_progress.status != "error":
                file_progress.status = "error"
                file_progress.error = "Upload returned failure"

            results.append(file_progress)

        logger.info(
            "sync_files_complete",
            total=len(files_to_upload),
            succeeded=sum(1 for r in results if r.status == "complete"),
            failed=sum(1 for r in results if r.status == "error"),
        )

        return results

    async def list_remote_files(
        self,
        remote_dir: str = "",
    ) -> list[dict[str, Any]]:
        """
        List files on the remote server.

        Sends a file list request and waits for the server to respond
        with a JSON list of file metadata.

        Args:
            remote_dir: Remote directory to list. Defaults to root.

        Returns:
            List of file info dicts, each containing at minimum
            'name', 'size', and optionally 'checksum'.
        """
        self._ensure_connected()

        logger.info("list_remote_files", remote_dir=remote_dir)

        response = await self.send_and_wait(
            MessageType.COMMAND,
            {
                "command": FileOperationType.LIST.value,
                "remote_dir": remote_dir,
            },
            timeout=30.0,
        )

        files: list[dict[str, Any]] = response.get("files", [])
        logger.info("list_remote_files_result", count=len(files))
        return files
