"""
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 os
from collections.abc import AsyncIterator, Callable
from contextlib import asynccontextmanager
from pathlib import Path
from typing import Any

import websockets
from websockets.client import WebSocketClientProtocol
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.utils.retry import async_retry
from aicippy.websocket.models import (
    AgentUpdate,
    FileOperationType,
    FileTransferMessage,
    FileTransferProgress,
    MessageType,
    ToolOutput,
    WebSocketMessage,
)

logger = get_logger(__name__)


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

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

    def __init__(self) -> None:
        """Initialize WebSocket client."""
        self._settings = get_settings()
        self._auth = CognitoAuth()
        self._ws: WebSocketClientProtocol | None = None
        self._connected = 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()

    @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

            # 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:
                self._ws = await websockets.connect(
                    url,
                    additional_headers={
                        "Authorization": f"Bearer {tokens.access_token}",
                    },
                    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 as e:
                logger.error("websocket_connection_failed", error=str(e))
                self._connected = False
                return False

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

        if self._receive_task:
            self._receive_task.cancel()
            try:
                await self._receive_task
            except asyncio.CancelledError:
                pass
            self._receive_task = None

        if self._reconnect_task:
            self._reconnect_task.cancel()
            try:
                await self._reconnect_task
            except asyncio.CancelledError:
                pass
            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.
        """
        await self.send(
            MessageType.AUTHENTICATE,
            {"token": token},
        )

    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(),
        )

        try:
            await self._ws.send(json.dumps(message.to_dict()))  # type: ignore
            logger.debug("websocket_message_sent", type=message_type.value)
        except ConnectionClosed:
            logger.warning("websocket_send_connection_closed")
            self._connected = False
            raise WebSocketDisconnectedError(
                message="WebSocket connection closed during send.",
            )

    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 asyncio.TimeoutError:
            logger.warning("websocket_response_timeout", correlation_id=correlation_id)
            raise TimeoutError("Response timeout")
        finally:
            self._pending_responses.pop(correlation_id, None)

    async def _receive_loop(self) -> None:
        """Background loop to receive messages."""
        while self._connected and self._ws:
            try:
                message = await self._ws.recv()
                await self._handle_message(message)
            except ConnectionClosed:
                logger.warning("websocket_connection_closed")
                self._connected = False
                # Schedule reconnect as a separate task so disconnect() can cancel it
                self._reconnect_task = asyncio.create_task(self._attempt_reconnect())
                break
            except asyncio.CancelledError:
                raise
            except Exception as e:
                logger.exception("websocket_receive_error", error=str(e))
                # Continue loop for transient errors; break on too many consecutive failures
                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)

            # 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)."""
        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):
            delay = min(base_delay * (2 ** (attempt - 1)), max_delay)
            logger.info(
                "websocket_reconnecting",
                attempt=attempt,
                max_attempts=max_attempts,
                delay=delay,
            )

            await asyncio.sleep(delay)

            if await self.connect():
                logger.info("websocket_reconnected", attempt=attempt)
                return

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

    @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:
            await self.send(MessageType.CHAT, {"message": message})
            try:
                await asyncio.wait_for(complete.wait(), timeout=300.0)
            except asyncio.TimeoutError:
                raise WebSocketMessageError(
                    message="Streaming response timed out after 300 seconds.",
                )
            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 open(file_path, "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
        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 open(path, "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:
                        await self._ws.send(frame)  # type: ignore[union-attr]
                    except ConnectionClosed:
                        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=f"Connection lost during upload at chunk {chunk_index}/{total_chunks}",
                        )

                    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()

            with open(dest, "wb") as f:
                while not download_complete.is_set():
                    try:
                        raw = await asyncio.wait_for(
                            self._ws.recv(),  # type: ignore[union-attr]
                            timeout=60.0,
                        )
                    except asyncio.TimeoutError:
                        download_error.append("Download timed out waiting for data")
                        break
                    except ConnectionClosed:
                        self._connected = False
                        download_error.append("Connection lost during download")
                        break

                    # Handle binary frames (file chunks)
                    if isinstance(raw, bytes):
                        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}%",
                        )
                    else:
                        # JSON text message (could be completion or error signal)
                        data = json.loads(raw)
                        if data.get("type") == "error":
                            download_error.append(data.get("message", "Unknown download error"))
                            break
                        if data.get("complete", False):
                            if data.get("checksum"):
                                received_checksum.append(data["checksum"])
                            download_complete.set()

            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:
            raise
        except Exception as e:
            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 = os.path.join(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
