"""
Site connection - Connection + tools for one site.

Manages Socket.IO connection lifecycle, tool subscriptions, wake-up handling,
and reconnection logic.
"""

from __future__ import annotations

import asyncio
import random
import sys
from collections.abc import Awaitable, Callable
from typing import TYPE_CHECKING, Any

from iq_sdk.client.connection import IQConnection
from iq_sdk.client.output_queue import OutputQueueManager
from iq_sdk.client.tool_worker import ToolConfig, ToolWorker

if TYPE_CHECKING:
    from iq_sdk.config import RetryConfig, SiteConfig
    from iq_sdk.worker_registry import WorkerRegistry


def _log(msg: str, *, prefix: str = "[Worker]", verbose_only: bool = False) -> None:
    """Print message to stderr."""
    print(f"{prefix} {msg}", file=sys.stderr)


class SiteConnection:
    """
    Manages connection and tool workers for a single site.

    Handles:
    - Socket.IO connection lifecycle
    - Tool subscriptions
    - Wake-up handling
    - Reconnection logic

    Timing parameters (alive_timeout, reconnect_interval) are fetched from
    the host via check_permissions and stored after the initial connection.
    """

    def __init__(
        self,
        site: SiteConfig,
        project_name: str,
        worker_id: str,
        tools: list[ToolConfig],
        registry: WorkerRegistry | None = None,
        project_path: str | None = None,
        retry_config: RetryConfig | None = None,
        debug: bool = False,
        verbose: bool = False,
    ):
        """
        Initialize site connection.

        Args:
            site: Site configuration
            project_name: Project namespace
            worker_id: Worker identifier
            tools: List of tool configurations
            registry: Worker registry for aggregating stats across processes
            project_path: Project root path for data directories
            retry_config: Retry configuration for output queue
            debug: Enable basic debug output
            verbose: Enable detailed verbose output (API calls, websocket events)
        """
        self.site = site
        self.project_name = project_name
        self.worker_id = worker_id
        self.tools = tools
        self.registry = registry
        self._debug = debug
        self._verbose = verbose

        self._conn: IQConnection | None = None
        self._tool_workers: dict[str, ToolWorker] = {}
        # These are set by check_permissions in start()
        self._alive_timeout: float = 0
        self._stopped = False
        self._api_client: Any = None

        # Wake-up wait tracking: tool_name -> asyncio.Event (set to cancel wait)
        self._wake_cancel_events: dict[str, asyncio.Event] = {}
        # Track active wait tasks to avoid duplicates
        self._active_wait_tasks: dict[str, asyncio.Task[None]] = {}

        # Output queue manager for retry
        self._output_queue: OutputQueueManager | None = None
        if project_path:
            self._output_queue = OutputQueueManager(
                base_path=project_path,
                site_id=site.id,
                retry_config=retry_config,
            )

    def _log(self, msg: str, *, verbose_only: bool = False) -> None:
        """Print log message if appropriate."""
        if verbose_only and not self._verbose:
            return
        if not self._debug and not self._verbose:
            return
        _log(msg)

    async def start(self) -> bool:
        """
        Start the site connection.

        Connects, verifies with ping/pong, subscribes to tools, and starts workers.

        Returns:
            True if started successfully
        """
        from iq_sdk.client import IQApiClient

        self._log(f"Connecting to {self.site.id} ({self.site.get_host()})...", verbose_only=True)

        # Step 1: Check permissions and get Socket.IO URL
        self._log(f"→ GET check_permissions (namespace: {self.project_name})", verbose_only=True)
        async with IQApiClient.from_site(self.site) as client:
            status = await client.check_status(namespace=self.project_name)
        self._log(
            f"← Response: reachable={status.reachable}, can_register={status.can_register_worker}",
            verbose_only=True,
        )

        if not status.reachable:
            self._log(f"✗ Site not reachable: {status.error}")
            return False

        if not status.can_register_worker:
            self._log(f"✗ No worker permission for {self.site.id}")
            return False

        if not status.socket_io_url or not status.socket_io_ns:
            self._log("✗ No Socket.IO URL/namespace")
            return False

        # Store timing configuration from server
        self._alive_timeout = float(status.alive_timeout)

        self._log(f"  alive_timeout={self._alive_timeout}s", verbose_only=True)

        # Step 2: Connect via Socket.IO
        self._log(
            f"→ WS connect to {status.socket_io_url} (ns: {status.socket_io_ns})", verbose_only=True
        )
        self._conn = IQConnection.from_site(
            self.site,
            status.socket_io_url,
            status.socket_io_ns,
            worker_id=self.worker_id,
        )

        connected = await self._conn.connect(debug=self._verbose)
        if not connected:
            self._log(f"✗ Connection failed: {self._conn.state.error}")
            return False
        self._log("← WS connected", verbose_only=True)

        # Step 3: Verify with ping/pong
        self._log("→ WS ping", verbose_only=True)
        pong_ok = await self._conn.ping()
        if not pong_ok:
            self._log(f"✗ Ping failed: {self._conn.state.error}")
            await self._conn.disconnect()
            return False
        self._log("← WS pong OK", verbose_only=True)

        # Step 4: Create API client for HTTP calls
        self._api_client = IQApiClient.from_site(self.site)
        await self._api_client.__aenter__()

        # Step 5: Create tool workers
        self._log(
            f"  Creating workers for {len(self.tools)} tools: {[t.name for t in self.tools]}",
            verbose_only=True,
        )
        for tool_config in self.tools:
            self._log(f"    → Worker for: {tool_config.name}", verbose_only=True)
            worker = ToolWorker(
                config=tool_config,
                worker_id=self.worker_id,
                api_client=self._api_client,
                site_id=self.site.id,
                registry=self.registry,
                output_queue=self._output_queue,
                debug=self._debug,
                verbose=self._verbose,
            )
            # Set callback to re-poll after task completion
            worker.on_task_complete = self._create_repoll_callback(worker)
            self._tool_workers[tool_config.name] = worker
        self._log(
            f"  Created {len(self._tool_workers)} workers: {list(self._tool_workers.keys())}",
            verbose_only=True,
        )

        # Step 6: Subscribe to all tools
        tool_names = [t.name for t in self.tools]
        self._log(f"→ WS subscribe_tools: {tool_names}", verbose_only=True)
        await self._conn.subscribe_tools(tool_names)
        self._log(f"← Subscribed to {len(tool_names)} tools", verbose_only=True)

        # Step 7: Register docinfo_update handler
        self._conn.on_docinfo_update(self._on_wake)

        return True

    async def initial_poll(self) -> None:
        """
        Perform initial poll for all tools.

        Should be called after start() succeeds and connection status is displayed.
        """
        self._log(f"→ Initial POLL for {len(self._tool_workers)} tools", verbose_only=True)
        for worker in self._tool_workers.values():
            tasks = await worker.poll()
            for task in tasks:
                asyncio.create_task(worker.execute_task(task))
        self._log("← Initial POLL complete", verbose_only=True)

    async def run(self) -> None:
        """
        Run the main worker loop.

        Handles:
        - Periodic POLL for alive status
        - Reconnection on disconnect
        - Graceful shutdown on CancelledError
        """
        try:
            while not self._stopped:
                if not self._conn or not self._conn.is_connected:
                    # Reconnect loop
                    self._log("Connection lost, attempting reconnect...")

                    self._reset_workers()

                    while not self._stopped:
                        try:
                            success = await self.start()
                            if success:
                                # Log reconnection success
                                tool_names = [t.name for t in self.tools]
                                _log(f"Reconnected to {self.site.get_host()}")
                                _log(f"  Tools: {', '.join(tool_names)}")
                                # Initial poll after reconnect
                                await self.initial_poll()
                                break
                        except asyncio.CancelledError:
                            # Shutdown requested during reconnect
                            raise
                        except Exception as e:
                            self._log(f"Reconnect error: {e}", verbose_only=True)

                        # Random delay between 3-10 seconds to avoid thundering herd
                        reconnect_delay = random.uniform(3.0, 10.0)
                        self._log(f"Reconnect failed, retrying in {reconnect_delay:.1f}s")

                        try:
                            await asyncio.sleep(reconnect_delay)
                        except asyncio.CancelledError:
                            raise

                    if self._stopped:
                        break

                # Wait for alive timeout or disconnect
                if not self._conn:
                    continue

                disconnect_event = asyncio.Event()
                self._conn.set_disconnect_callback(disconnect_event)

                try:
                    await asyncio.wait_for(
                        disconnect_event.wait(),
                        timeout=self._alive_timeout,
                    )
                    # Disconnected - loop will reconnect
                    self._log("← WS disconnected", verbose_only=True)
                except TimeoutError:
                    # Alive timeout - do POLL for all tools
                    self._log(f"→ Alive POLL ({self._alive_timeout}s timeout)", verbose_only=True)

                    for worker in self._tool_workers.values():
                        tasks = await worker.poll()
                        for task in tasks:
                            asyncio.create_task(worker.execute_task(task))

        except asyncio.CancelledError:
            # Graceful shutdown - don't propagate, just stop
            self._stopped = True

    async def _on_wake(self, data: Any) -> None:
        """
        Handle docinfo_update wake-up event.

        Events:
        - task_added: New task available, add task_id to pending
        - task_reserved: Task was taken by a worker, remove task_id from pending

        Args:
            data: Wake-up data containing event type, tool info, task_id, worker count
        """
        if not isinstance(data, dict):
            return

        event_type = data.get("event")
        tool_name = data.get("tool_name")
        task_id = data.get("task_id")
        idle_process_count = data.get("idle_process_count", 1)
        active_process_count = data.get("active_process_count", 1)
        running_task_count = data.get("running_task_count", 0)

        # Handle task_reserved: remove from pending (any tool)
        if event_type == "task_reserved" and task_id:
            for worker in self._tool_workers.values():
                worker.remove_pending_task(task_id)
            self._log(f"← WS task_reserved: {task_id}", verbose_only=True)
            return

        # For other events, we need a valid tool
        if not tool_name or tool_name not in self._tool_workers:
            self._log(f"← WS wake-up: {tool_name} (unknown tool, ignoring)", verbose_only=True)
            return

        worker = self._tool_workers[tool_name]

        # Handle task_added: add to pending and trigger poll logic
        if event_type == "task_added" and task_id:
            worker.add_pending_task(task_id)
            self._log(
                f"← WS task_added: {tool_name} ({task_id}, idle:{idle_process_count}, running:{running_task_count}, active:{active_process_count})",
                verbose_only=True,
            )
        else:
            self._log(
                f"← WS wake-up: {tool_name} (idle:{idle_process_count}, running:{running_task_count})",
                verbose_only=True,
            )

        # Check if worker is busy
        if worker.is_busy:
            self._log(
                f"  → Skip (busy: {worker._running_tasks}/{worker.config.thread_count})",
                verbose_only=True,
            )
            return

        # Cancel any existing wait for this tool
        if tool_name in self._wake_cancel_events:
            self._wake_cancel_events[tool_name].set()

        # Decide: poll immediately or wait?
        # Use idle_process_count and active_process_count to determine if we should poll now
        if worker.should_poll_immediately(idle_process_count, active_process_count):
            self._log(
                f"  → POLL {tool_name} (immediate, active:{active_process_count}, idle:{idle_process_count})",
                verbose_only=True,
            )
            await self._do_poll(worker)
        else:
            # Start delayed poll (or restart if already waiting)
            # Wait time scales with number of idle workers to spread load
            wait_time = worker.get_random_wait_time(idle_process_count)
            self._log(
                f"  → Wait {wait_time:.1f}s (active:{active_process_count}, idle:{idle_process_count})",
                verbose_only=True,
            )
            self._start_delayed_poll(worker, wait_time)

    def _start_delayed_poll(self, worker: ToolWorker, wait_time: float) -> None:
        """Start a delayed poll task for a worker."""
        tool_name = worker.config.name

        # Cancel existing task if any
        if tool_name in self._active_wait_tasks:
            task = self._active_wait_tasks[tool_name]
            if not task.done():
                task.cancel()

        # Create new cancel event
        cancel_event = asyncio.Event()
        self._wake_cancel_events[tool_name] = cancel_event

        # Start new wait task
        async def delayed_poll() -> None:
            try:
                # Wait for timeout or cancellation
                try:
                    await asyncio.wait_for(cancel_event.wait(), timeout=wait_time)
                    # Event was set = cancelled by new wake-up, don't poll
                    self._log(f"  ← Wait cancelled for {tool_name}", verbose_only=True)
                    return
                except TimeoutError:
                    # Timeout = wait completed normally
                    pass

                # Check if there are still pending tasks
                if not worker.has_pending_tasks():
                    self._log(f"  ← No pending tasks for {tool_name}, skip poll", verbose_only=True)
                    return

                if worker.is_busy:
                    self._log("  ← Worker busy, skip poll", verbose_only=True)
                    return

                # Poll for tasks
                self._log(f"  → POLL {tool_name} (after wait)", verbose_only=True)
                await self._do_poll(worker)

            except asyncio.CancelledError:
                pass
            finally:
                # Cleanup
                self._wake_cancel_events.pop(tool_name, None)
                self._active_wait_tasks.pop(tool_name, None)

        self._active_wait_tasks[tool_name] = asyncio.create_task(delayed_poll())

    async def _do_poll(self, worker: ToolWorker) -> None:
        """Execute poll and start tasks."""
        tasks = await worker.poll()
        if tasks:
            self._log(f"  ← Got {len(tasks)} task(s)", verbose_only=True)
            # Clear pending tasks since we're processing
            worker.clear_pending_tasks()
            for task in tasks:
                asyncio.create_task(worker.execute_task(task))
        else:
            self._log("  ← No task", verbose_only=True)
            worker.clear_pending_tasks()
            # No new tasks - process retry queue
            await self._process_retry_queue(worker)

    async def _process_retry_queue(self, worker: ToolWorker) -> None:
        """Process pending outputs from retry queue for a worker."""
        if not worker.output_queue:
            return

        retried = await worker.process_retry_queue()
        if retried > 0:
            self._log(f"  ← Retried {retried} queued output(s)", verbose_only=True)

    def _create_repoll_callback(self, worker: ToolWorker) -> Callable[[], Awaitable[None]]:
        """
        Create a callback that re-polls for tasks after one completes.

        This ensures workers keep processing until no more tasks are available.
        Also processes retry queue when no new tasks are available.
        """

        async def repoll() -> None:
            if self._stopped or not worker.available_threads:
                return

            self._log(f"  → Re-poll {worker.config.name}", verbose_only=True)
            tasks = await worker.poll()
            if tasks:
                self._log(f"  ← Got {len(tasks)} more task(s)", verbose_only=True)
                for task in tasks:
                    asyncio.create_task(worker.execute_task(task))
            else:
                # No more tasks - process retry queue before going idle
                await self._process_retry_queue(worker)

        return repoll

    def _reset_workers(self) -> None:
        """Reset all tool workers and cancel pending wait tasks."""
        # Cancel all pending wait tasks
        for task in self._active_wait_tasks.values():
            if not task.done():
                task.cancel()
        self._active_wait_tasks.clear()
        self._wake_cancel_events.clear()

        # Stop all workers
        for worker in self._tool_workers.values():
            worker.stop()
        self._tool_workers.clear()

    async def stop(self) -> None:
        """Stop the site connection and all workers."""
        self._stopped = True
        self._log(f"Stopping {self.site.id}...", verbose_only=True)

        # Notify host that workers are leaving (while API client is still active)
        # Use self.tools instead of _tool_workers, as workers may have been reset
        if self._api_client:
            await self._send_leave_for_all_tools()

        # Stop all workers
        for worker in self._tool_workers.values():
            try:
                worker.stop()
            except Exception:
                pass

        # Unsubscribe and disconnect
        if self._conn:
            try:
                if self._conn.is_connected:
                    tool_names = [t.name for t in self.tools]
                    self._log(f"→ WS unsubscribe_tools: {tool_names}", verbose_only=True)
                    await self._conn.unsubscribe_tools(tool_names)
                self._log("→ WS disconnect", verbose_only=True)
                await self._conn.disconnect()
            except Exception:
                pass

        # Close API client
        if self._api_client:
            try:
                await self._api_client.__aexit__(None, None, None)
            except Exception:
                pass

    async def _send_leave_for_all_tools(self) -> None:
        """Send leave request for all tools in this connection."""
        import os

        current_pid = os.getpid()

        for tool_config in self.tools:
            # Get remaining stats (excluding this process)
            if self.registry:
                pids, total_threads = self.registry.get_tool_stats(
                    self.site.id, tool_config.name, exclude_pid=current_pid
                )
                active_pids = ",".join(str(p) for p in pids) if pids else ""
            else:
                active_pids = ""
                total_threads = 0

            payload = {
                "tool_name": tool_config.name,
                "worker_id": self.worker_id,
                "thread_count": total_threads,
                "active_pids": active_pids,
            }

            self._log(f"→ POST leave ({tool_config.name})", verbose_only=True)

            try:
                response = await self._api_client.post(
                    "/api/v2/method/iq_core.iq_core.doctype.automation_worker.automation_worker.leave",
                    json=payload,
                )
                self._log(f"← {response.status_code}", verbose_only=True)
            except Exception as e:
                self._log(f"Leave request failed: {e}", verbose_only=True)
