"""
Base worker class for iq-sdk.

Workers connect to sites via Socket.IO and process tasks
using registered tools.
"""

from __future__ import annotations

import asyncio
import logging
import signal
import sys
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass, field
from datetime import UTC, datetime
from typing import Any

import socketio

from iq_sdk.config import Settings, SiteConfig, WorkerConfig
from iq_sdk.tool import BaseTool, ToolResult

logger = logging.getLogger(__name__)


@dataclass
class WorkerStats:
    """Statistics for worker monitoring."""

    tasks_processed: int = 0
    tasks_failed: int = 0
    tasks_in_progress: int = 0
    total_processing_time_ms: float = 0.0
    started_at: datetime = field(default_factory=lambda: datetime.now(UTC))
    last_task_at: datetime | None = None


@dataclass
class SiteConnection:
    """Represents a connection to a site."""

    site: SiteConfig
    client: socketio.AsyncClient
    connected: bool = False


class BaseWorker:
    """
    Base class for Socket.IO connected workers.

    Workers:
    - Connect to one or more sites via Socket.IO
    - Authenticate with site credentials
    - Register available tools with the site
    - Receive and process tasks using registered tools
    - Support configurable thread pools and autoscaling

    Example:
        class MyWorker(BaseWorker):
            def __init__(self, settings: Settings):
                super().__init__(settings)
                self.register_tool(MyTool())

        worker = MyWorker(load_settings())
        await worker.run()
    """

    def __init__(self, settings: Settings):
        """
        Initialize worker with settings.

        Args:
            settings: Configuration settings
        """
        self.settings = settings
        self.config: WorkerConfig = settings.worker

        # Get all sites (default + additional)
        self._sites: list[SiteConfig] = [settings.site] + settings.sites

        self._tools: dict[str, BaseTool] = {}
        self._connections: list[SiteConnection] = []
        self._executor: ThreadPoolExecutor | None = None
        self._running = False
        self._stats = WorkerStats()
        self._tasks: set[asyncio.Task] = set()
        self._semaphore: asyncio.Semaphore | None = None

        # Autoscaling state
        self._current_threads = self.config.threads
        self._autoscale_task: asyncio.Task | None = None

    @property
    def name(self) -> str:
        """Worker name identifier."""
        return self.config.name

    @property
    def stats(self) -> WorkerStats:
        """Current worker statistics."""
        return self._stats

    @property
    def sites(self) -> list[SiteConfig]:
        """Get configured sites."""
        return self._sites

    def register_tool(self, tool: BaseTool) -> None:
        """
        Register a tool with this worker.

        Args:
            tool: Tool instance to register
        """
        if tool.name in self._tools:
            logger.warning(f"Tool '{tool.name}' already registered, replacing")
        self._tools[tool.name] = tool
        logger.info(f"Registered tool: {tool.name} v{tool.version}")

    def get_tool(self, name: str) -> BaseTool | None:
        """Get a registered tool by name."""
        return self._tools.get(name)

    async def connect(self) -> None:
        """Establish Socket.IO connections to all sites."""
        for site in self._sites:
            sio = socketio.AsyncClient(
                reconnection=site.reconnect,
                reconnection_attempts=site.reconnect_attempts,
            )

            conn = SiteConnection(site=site, client=sio)

            # Register event handlers
            sio.on("connect", self._make_connect_handler(conn))
            sio.on("disconnect", self._make_disconnect_handler(conn))
            sio.on("task", self._make_task_handler(conn))
            sio.on("ping", self._make_ping_handler(conn))

            logger.info(f"Connecting to site '{site.id}' at {site.get_host()}")
            try:
                # Build auth payload
                auth = site.get_auth()

                await sio.connect(
                    site.get_host(),
                    auth=auth if auth["api_key"] else None,
                )
                conn.connected = True
                self._connections.append(conn)
            except Exception as e:
                logger.error(f"Failed to connect to site '{site.id}': {e}")

    def _make_connect_handler(self, conn: SiteConnection):
        """Create connect handler for specific site."""

        async def handler():
            logger.info(f"Connected to site '{conn.site.id}' as worker '{self.name}'")
            conn.connected = True

            # Register with site
            await conn.client.emit(
                "register",
                {
                    "name": self.name,
                    "tools": [t.to_info() for t in self._tools.values()],
                    "threads": self._current_threads,
                    "env": conn.site.get_env(),
                },
            )

        return handler

    def _make_disconnect_handler(self, conn: SiteConnection):
        """Create disconnect handler for specific site."""

        async def handler():
            logger.warning(f"Disconnected from site '{conn.site.id}'")
            conn.connected = False

        return handler

    def _make_task_handler(self, conn: SiteConnection):
        """Create task handler for specific site."""

        async def handler(data: dict[str, Any]):
            await self._on_task(conn, data)

        return handler

    def _make_ping_handler(self, conn: SiteConnection):
        """Create ping handler for specific site."""

        async def handler():
            await conn.client.emit(
                "pong",
                {
                    "name": self.name,
                    "site": conn.site.id,
                    "stats": {
                        "tasks_processed": self._stats.tasks_processed,
                        "tasks_failed": self._stats.tasks_failed,
                        "tasks_in_progress": self._stats.tasks_in_progress,
                        "threads": self._current_threads,
                    },
                },
            )

        return handler

    async def disconnect(self) -> None:
        """Disconnect from all sites."""
        for conn in self._connections:
            if conn.client.connected:
                await conn.client.disconnect()
        self._connections.clear()
        logger.info("Disconnected from all sites")

    async def run(self) -> None:
        """
        Main worker loop.

        Connects to sites, sets up tools, and processes tasks until shutdown.
        """
        self._running = True
        self._semaphore = asyncio.Semaphore(self._current_threads)
        self._executor = ThreadPoolExecutor(max_workers=self.config.max_threads)

        # Setup signal handlers
        self._setup_signal_handlers()

        # Setup tools
        for tool in self._tools.values():
            await tool.setup()

        try:
            await self.connect()

            if not self._connections:
                logger.error("Failed to connect to any site")
                return

            connected_sites = [c.site.id for c in self._connections if c.connected]
            logger.info(f"Connected to sites: {', '.join(connected_sites)}")

            # Start autoscaling if enabled
            if self.config.autoscale:
                self._autoscale_task = asyncio.create_task(self._autoscale_loop())

            # Start heartbeat
            heartbeat_task = asyncio.create_task(self._heartbeat_loop())

            # Wait until shutdown
            while self._running:
                await asyncio.sleep(1)

            # Cleanup
            heartbeat_task.cancel()
            if self._autoscale_task:
                self._autoscale_task.cancel()

            # Wait for in-progress tasks
            if self._tasks:
                logger.info(f"Waiting for {len(self._tasks)} tasks to complete...")
                await asyncio.gather(*self._tasks, return_exceptions=True)

        finally:
            await self._shutdown()

    async def _shutdown(self) -> None:
        """Cleanup and shutdown worker."""
        logger.info("Shutting down worker...")

        # Teardown tools
        for tool in self._tools.values():
            await tool.teardown()

        # Disconnect
        await self.disconnect()

        # Shutdown executor
        if self._executor:
            self._executor.shutdown(wait=True)

        logger.info("Worker shutdown complete")

    def _setup_signal_handlers(self) -> None:
        """Setup graceful shutdown signal handlers."""

        def signal_handler(sig: int, frame: Any) -> None:
            logger.info(f"Received signal {sig}, initiating shutdown...")
            self._running = False

        if sys.platform != "win32":
            signal.signal(signal.SIGTERM, signal_handler)
            signal.signal(signal.SIGINT, signal_handler)

    async def _on_task(
        self,
        conn: SiteConnection,
        data: dict[str, Any],
    ) -> None:
        """
        Handle incoming task.

        Args:
            conn: Site connection that received the task
            data: Task data including tool name and input
        """
        task_id = data.get("task_id", "unknown")
        tool_name = data.get("tool", "")
        input_data = data.get("input", {})

        tool = self.get_tool(str(tool_name))
        if not tool:
            await self._send_result(
                conn,
                task_id,
                ToolResult(
                    success=False,
                    error=f"Unknown tool: {tool_name}",
                ),
            )
            return

        # Create task with semaphore for concurrency control
        task = asyncio.create_task(self._process_task(conn, task_id, tool, input_data))
        self._tasks.add(task)
        task.add_done_callback(self._tasks.discard)

    async def _process_task(
        self,
        conn: SiteConnection,
        task_id: str,
        tool: BaseTool,
        input_data: dict[str, Any],
    ) -> None:
        """Process a single task."""
        if self._semaphore is None:
            return
        async with self._semaphore:
            self._stats.tasks_in_progress += 1
            start_time = datetime.now(UTC)

            try:
                # Execute tool with input data
                result = await tool.execute(data={"input": input_data})

                # Update stats
                self._stats.tasks_processed += 1
                self._stats.last_task_at = datetime.now(UTC)
                duration = (datetime.now(UTC) - start_time).total_seconds() * 1000
                self._stats.total_processing_time_ms += duration
                result.duration_ms = duration

            except Exception as e:
                logger.exception(f"Task {task_id} failed: {e}")
                self._stats.tasks_failed += 1
                result = ToolResult(success=False, error=str(e))

            finally:
                self._stats.tasks_in_progress -= 1

            await self._send_result(conn, task_id, result)

    async def _send_result(
        self,
        conn: SiteConnection,
        task_id: str,
        result: ToolResult,
    ) -> None:
        """Send task result back to site."""
        await conn.client.emit(
            "task_result",
            {
                "task_id": task_id,
                "success": result.success,
                "output": result.output.model_dump() if result.output else None,
                "error": result.error,
                "duration_ms": result.duration_ms,
                "metadata": result.metadata,
            },
        )

    async def _heartbeat_loop(self) -> None:
        """Send periodic heartbeats to all sites."""
        while self._running:
            await asyncio.sleep(self.config.heartbeat_interval)
            for conn in self._connections:
                if conn.connected:
                    try:
                        await conn.client.emit(
                            "heartbeat",
                            {
                                "name": self.name,
                                "site": conn.site.id,
                                "threads": self._current_threads,
                                "in_progress": self._stats.tasks_in_progress,
                            },
                        )
                    except Exception as e:
                        logger.warning(f"Heartbeat to '{conn.site.id}' failed: {e}")

    async def _autoscale_loop(self) -> None:
        """Autoscaling loop for dynamic thread adjustment."""
        while self._running:
            await asyncio.sleep(self.config.autoscale_interval)
            await self._check_autoscale()

    async def _check_autoscale(self) -> None:
        """
        Check and adjust thread count based on load.

        Override this method for custom autoscaling logic.
        """
        if self._current_threads == 0:
            return

        utilization = self._stats.tasks_in_progress / self._current_threads

        if utilization > 0.8 and self._current_threads < self.config.max_threads:
            # Scale up
            new_threads = min(self._current_threads + 2, self.config.max_threads)
            await self._scale_to(new_threads)
        elif utilization < 0.2 and self._current_threads > self.config.min_threads:
            # Scale down
            new_threads = max(self._current_threads - 1, self.config.min_threads)
            await self._scale_to(new_threads)

    async def _scale_to(self, threads: int) -> None:
        """Scale to specified thread count."""
        if threads == self._current_threads:
            return

        logger.info(f"Scaling from {self._current_threads} to {threads} threads")
        self._current_threads = threads
        self._semaphore = asyncio.Semaphore(threads)

        # Notify all sites
        for conn in self._connections:
            if conn.connected:
                try:
                    await conn.client.emit(
                        "scale",
                        {"name": self.name, "threads": threads},
                    )
                except Exception as e:
                    logger.warning(f"Scale notification to '{conn.site.id}' failed: {e}")
