import asyncio
import json
from dataclasses import asdict
from datetime import datetime
from typing import Any

from fastapi import WebSocket, WebSocketDisconnect

from vna_system.core.logging.logger import get_component_logger
from vna_system.core.processors.base_processor import ProcessedResult
from vna_system.core.processors.manager import ProcessorManager
from vna_system.core.processors.storage import DataStorage

logger = get_component_logger(__file__)


class ProcessorWebSocketHandler:
    """
    WebSocket hub for processor results.

    Responsibilities
    ----------------
    • Accept and manage client connections.
    • Handle client commands (e.g., recalculate, fetch history).
    • Broadcast `ProcessedResult` objects to all connected clients.
    • Bridge results produced on worker threads into the main asyncio loop.

    Threading model
    ---------------
    - Processor callbacks arrive on non-async worker threads.
    - We capture the main running event loop when the first client connects.
    - Cross-thread scheduling uses `asyncio.run_coroutine_threadsafe`.
    """

    def __init__(self, processor_manager: ProcessorManager, data_storage: DataStorage) -> None:
        self.processor_manager = processor_manager
        self.data_storage = data_storage

        self.active_connections: set[WebSocket] = set()

        # Main FastAPI/uvicorn event loop handle (set on first connection).
        self._loop: asyncio.AbstractEventLoop | None = None

        # Subscribe to processor results (thread-safe callback).
        self.processor_manager.add_result_callback(self._on_processor_result)
        logger.debug("ProcessorWebSocketHandler initialized")

    # --------------------------------------------------------------------- #
    # Connection lifecycle
    # --------------------------------------------------------------------- #
    async def handle_websocket_connection(self, websocket: WebSocket) -> None:
        """
        Accept a WebSocket, then serve client messages until disconnect.

        Stores a reference to the running event loop so worker threads can
        safely schedule broadcasts into it.
        """
        if self._loop is None:
            self._loop = asyncio.get_running_loop()
            logger.info("Main event loop captured for broadcasts")
            
        await websocket.accept()
        self.active_connections.add(websocket)
        logger.info("WebSocket connected", total_connections=len(self.active_connections))

        try:
            while True:
                raw = await websocket.receive_text()
                await self.handle_message(websocket, raw)
        except WebSocketDisconnect:
            await self.disconnect(websocket)
        except Exception as exc:  # noqa: BLE001
            logger.error("WebSocket error", error=repr(exc))
            await self.disconnect(websocket)

    async def disconnect(self, websocket: WebSocket) -> None:
        """Remove a connection and log the updated count."""
        if websocket in self.active_connections:
            self.active_connections.remove(websocket)
            logger.info("WebSocket disconnected", total_connections=len(self.active_connections))

    # --------------------------------------------------------------------- #
    # Inbound messages
    # --------------------------------------------------------------------- #
    async def handle_message(self, websocket: WebSocket, data: str) -> None:
        """Parse and route an inbound client message."""
        # Handle ping/pong messages first (they are not JSON)
        if data == "ping":
            await websocket.send_text("pong")
            return
        elif data == "pong":
            # Just acknowledge, no response needed
            return

        try:
            message = json.loads(data)
            mtype = message.get("type")

            if mtype == "recalculate":
                await self._handle_recalculate(websocket, message)
            elif mtype == "get_history":
                await self._handle_get_history(websocket, message)
            elif mtype == "load_history":
                await self._handle_load_history(websocket, message)
            elif mtype == "get_processor_state":
                await self._handle_get_processor_state(websocket, message)
            elif mtype == "delete_column":
                await self._handle_delete_column(websocket, message)
            elif mtype == "append_history":
                await self._handle_append_history(websocket, message)
            else:
                await self._send_error(websocket, f"Неизвестный тип сообщения: {mtype!r}")
        except json.JSONDecodeError as json_error:
            logger.error("JSON decode error", raw_data=data[:200], error=str(json_error))
            await self._send_error(
                websocket,
                "Неверный формат JSON",
                details=f"Ошибка парсинга JSON: {json_error}",
                source="websocket_handler",
                raw_data=data[:200] if len(data) <= 200 else f"{data[:200]}..."
            )
        except Exception as exc:  # noqa: BLE001
            logger.error("Error handling websocket message")
            await self._send_error(websocket, f"Внутренняя ошибка: {exc}")

    # --------------------------------------------------------------------- #
    # Client commands
    # --------------------------------------------------------------------- #
    async def _handle_recalculate(self, websocket: WebSocket, message: dict[str, Any]) -> None:
        """
        Recalculate a processor (optionally with config updates) and send the result back.
        """
        processor_id = message.get("processor_id")
        config_updates = message.get("config_updates")

        if not processor_id:
            await self._send_error(websocket, "Требуется processor_id")
            return

        try:
            result = self.processor_manager.recalculate_processor(processor_id, config_updates)
            if result:
                await websocket.send_text(json.dumps(self._result_to_message(processor_id, result)))
            else:
                await self._send_error(websocket, f"Нет результата от процессора {processor_id}")
        except Exception as exc:  # noqa: BLE001
            logger.error("Recalculation failed")
            await self._send_error(websocket, f"Пересчёт не удался: {exc}")

    async def _handle_get_history(self, websocket: WebSocket, message: dict[str, Any]) -> None:
        """
        Fetch recent results history for a given processor and send it to the client.
        """
        processor_id = message.get("processor_id")
        limit = int(message.get("limit", 10))

        if not processor_id:
            await self._send_error(websocket, "Требуется processor_id")
            return

        try:
            history = self.data_storage.get_results_history(processor_id, limit)
            response = {
                "type": "processor_history",
                "processor_id": processor_id,
                "history": [
                    {
                        "timestamp": r.timestamp,
                        "data": r.data,
                        "plotly_config": r.plotly_config,
                        "metadata": r.metadata,
                    }
                    for r in history
                ],
            }
            await websocket.send_text(json.dumps(response))
        except Exception as exc:  # noqa: BLE001
            logger.error("Error getting history")
            await self._send_error(websocket, f"Ошибка получения истории: {exc}")

    async def _handle_load_history(self, websocket: WebSocket, message: dict[str, Any]) -> None:
        """
        Load sweep history from JSON data into a processor and recalculate.

        Also restores processor configuration (UI parameters) if provided.
        """
        processor_id = message.get("processor_id")
        history_data = message.get("history_data")
        config = message.get("config")  # Optional processor configuration

        if not processor_id:
            await self._send_error(websocket, "Требуется processor_id")
            return

        if not history_data or not isinstance(history_data, list):
            await self._send_error(websocket, "Требуется history_data (список)")
            return

        try:
            result = self.processor_manager.load_processor_history(processor_id, history_data, config)
            if result:
                await websocket.send_text(json.dumps(self._result_to_message(processor_id, result)))
            else:
                await self._send_error(websocket, f"Нет результата от процессора {processor_id} после загрузки истории")
        except Exception as exc:  # noqa: BLE001
            logger.error("History load failed", processor_id=processor_id, error=repr(exc))
            await self._send_error(websocket, f"Загрузка истории не удалась: {exc}")

    async def _handle_append_history(self, websocket: WebSocket, message: dict[str, Any]) -> None:
        """
        Append sweep history from JSON data to existing processor history and recalculate.

        Does NOT update processor configuration - only appends history.
        """
        processor_id = message.get("processor_id")
        history_data = message.get("history_data")

        if not processor_id:
            await self._send_error(websocket, "Требуется processor_id")
            return

        if not history_data or not isinstance(history_data, list):
            await self._send_error(websocket, "Требуется history_data (список)")
            return

        try:
            result = self.processor_manager.append_processor_history(processor_id, history_data)
            if result:
                await websocket.send_text(json.dumps(self._result_to_message(processor_id, result)))
            else:
                await self._send_error(websocket, f"Нет результата от процессора {processor_id} после дополнения истории")
        except Exception as exc:  # noqa: BLE001
            logger.error("History append failed", processor_id=processor_id, error=repr(exc))
            await self._send_error(websocket, f"Дополнение истории не удалось: {exc}")

    async def _handle_get_processor_state(self, websocket: WebSocket, message: dict[str, Any]) -> None:
        """
        Fetch complete processor state including sweep history and current data.

        This should be called explicitly by the UI (e.g., via a download/export button)
        rather than being sent automatically with every result.
        """
        processor_id = message.get("processor_id")

        if not processor_id:
            await self._send_error(websocket, "Требуется processor_id")
            return

        try:
            response = self.processor_manager.build_processor_state(processor_id)
            await websocket.send_text(json.dumps(response))
        except Exception as exc:  # noqa: BLE001
            logger.error("Error getting processor state", processor_id=processor_id, error=repr(exc))
            await self._send_error(websocket, f"Ошибка получения состояния процессора: {exc}")

    async def _handle_delete_column(self, websocket: WebSocket, message: dict[str, Any]) -> None:
        """
        Delete a specific column (sweep) from the bscan processor plot history.
        """
        processor_id = message.get("processor_id")
        column_index = message.get("column_index")

        if not processor_id:
            await self._send_error(websocket, "Требуется processor_id")
            return

        if column_index is None or not isinstance(column_index, int):
            await self._send_error(websocket, "Требуется column_index (целое число)")
            return

        try:
            processor = self.processor_manager.get_processor(processor_id)
            if processor is None:
                await self._send_error(websocket, f"Процессор {processor_id} не найден")
                return

            # Check if processor has delete_column method (only BScanProcessor)
            if not hasattr(processor, "delete_column"):
                await self._send_error(
                    websocket,
                    f"Процессор {processor_id} не поддерживает удаление столбцов"
                )
                return

            # Execute deletion
            success = processor.delete_column(column_index)

            if not success:
                await self._send_error(
                    websocket,
                    f"Не удалось удалить столбец {column_index}"
                )
                return

            # Recalculate and send updated result to all clients
            result = processor.recalculate()
            if result:
                # Broadcast to all connected clients
                message_str = json.dumps(self._result_to_message(processor_id, result))
                await self._send_to_connections(message_str)
            else:
                await self._send_error(websocket, "Пересчёт после удаления не удался")

        except Exception as exc:  # noqa: BLE001
            logger.error("Column deletion failed", processor_id=processor_id, error=repr(exc))
            await self._send_error(websocket, f"Удаление столбца не удалось: {exc}")

    def _result_to_message(self, processor_id: str, result: ProcessedResult) -> dict[str, Any]:
        """
        Convert a `ProcessedResult` into a lightweight JSON-serializable message for broadcasting.

        Note: This excludes heavy data fields like 'data' to minimize websocket traffic.
        Use get_processor_state command to retrieve full data including sweep history.
        """
        return {
            "type": "processor_result",
            "processor_id": processor_id,
            "timestamp": result.timestamp,
            "plotly_config": result.plotly_config,
            "ui_parameters": [asdict(param) for param in result.ui_parameters],
            "metadata": result.metadata,
        }

    async def _send_error(
        self,
        websocket: WebSocket,
        message: str,
        *,
        details: str | None = None,
        source: str | None = None,
        code: str | None = None,
        raw_data: str | None = None
    ) -> None:
        """Send a standardized error payload to a single client."""
        try:
            payload = {
                "type": "error",
                "message": message,
                "timestamp": datetime.now().timestamp(),
            }

            # Add optional fields if provided
            if details:
                payload["details"] = details
            if source:
                payload["source"] = source
            if code:
                payload["code"] = code
            if raw_data:
                payload["raw_data"] = raw_data

            await websocket.send_text(json.dumps(payload))
        except Exception as exc:  # noqa: BLE001
            logger.error("Error sending error message", error=repr(exc))

    # --------------------------------------------------------------------- #
    # Result callback bridge (worker thread -> asyncio loop)
    # --------------------------------------------------------------------- #
    def _on_processor_result(self, processor_id: str, result: ProcessedResult) -> None:
        """
        Callback invoked on a worker thread when a processor produces a result.

        We:
        - Store the result synchronously in `DataStorage`.
        - Schedule a coroutine on the main event loop to broadcast to clients.
        """
        # Best-effort persistence
        # try:
        #     self.data_storage.store_result(processor_id, result)
        # except Exception:  # noqa: BLE001
        #     logger.error("Failed to store processor result")

        # Broadcast to clients
        self._broadcast_result_sync(processor_id, result)

    def _broadcast_result_sync(self, processor_id: str, result: ProcessedResult) -> None:
        """
        Thread-safe broadcast entrypoint from worker threads.

        Serializes once and schedules `_send_to_connections(...)` on the main loop.
        """
        if not self.active_connections:
            return

        loop = self._loop
        if loop is None or not loop.is_running():
            logger.debug("No running event loop available for broadcast; skipping")
            return

        try:
            message_str = json.dumps(self._result_to_message(processor_id, result))
        except Exception as exc:  # noqa: BLE001
            logger.error(f"Failed to serialize result for broadcast, {processor_id=}", error=repr(exc))
            return

        try:
            asyncio.run_coroutine_threadsafe(self._send_to_connections(message_str), loop)
        except Exception as exc:  # noqa: BLE001
            logger.error("Failed to schedule broadcast", error=repr(exc))

    async def _send_to_connections(self, message_str: str) -> None:
        """
        Broadcast a pre-serialized JSON string to all active connections.

        Removes connections that fail during send.
        """
        if not self.active_connections:
            return

        disconnected: list[WebSocket] = []
        for websocket in list(self.active_connections):  # snapshot
            try:
                await websocket.send_text(message_str)
            except Exception as exc:  # noqa: BLE001
                logger.error("Broadcast to client failed; marking for disconnect", error=repr(exc))
                disconnected.append(websocket)

        for websocket in disconnected:
            try:
                await self.disconnect(websocket)
            except Exception as exc:  # noqa: BLE001
                logger.error("Error during disconnect cleanup", error=repr(exc))
