"""
Tool worker - Per-tool worker runtime.

Manages threads, task execution, and lifecycle for one tool.
"""

from __future__ import annotations

import os
import random
import sys
from collections.abc import Awaitable, Callable
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
    from iq_sdk.client.output_queue import OutputQueueManager
    from iq_sdk.tool import BaseTool
    from iq_sdk.worker_registry import WorkerRegistry


def _log(msg: str, *, prefix: str = "[Worker]", error: bool = False) -> None:
    """Print message to stderr."""
    if error:
        # Red color for errors
        print(f"{prefix} \033[91m{msg}\033[0m", file=sys.stderr)
    else:
        print(f"{prefix} {msg}", file=sys.stderr)


@dataclass
class ToolConfig:
    """Configuration for a tool worker."""

    name: str  # Full tool name (e.g. "demo.my_tool.1.0")
    tool_class: type[BaseTool]
    thread_count: int = 1
    max_runtime_seconds: int = 600


class ToolWorker:
    """
    Worker runtime for a single tool.

    Manages threads, task execution, and lifecycle for one tool.
    """

    def __init__(
        self,
        config: ToolConfig,
        worker_id: str,
        api_client: Any,  # IQApiClient
        site_id: str,
        registry: WorkerRegistry | None = None,
        output_queue: OutputQueueManager | None = None,
        debug: bool = False,
        verbose: bool = False,
    ):
        """
        Initialize tool worker.

        Args:
            config: Tool configuration
            worker_id: Worker identifier
            api_client: API client for HTTP calls
            site_id: Site identifier (for registry lookups)
            registry: Worker registry for aggregating stats across processes
            output_queue: Output queue manager for failed task retry
            debug: Enable basic debug output
            verbose: Enable detailed verbose output
        """
        self.config = config
        self.worker_id = worker_id
        self.api_client = api_client
        self.site_id = site_id
        self.registry = registry
        self.output_queue = output_queue
        self._debug = debug
        self._verbose = verbose

        self._running_tasks: int = 0
        self._stopped = False

        # Pending task IDs from wake-up events (task_added)
        # Removed when task_reserved events arrive
        self._pending_tasks: set[str] = set()

        # Callback for when a task completes - allows parent to trigger re-poll
        self.on_task_complete: Callable[[], Awaitable[None]] | None = None

    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)

    def _handle_poll_error(self, response: Any) -> None:
        """
        Handle poll error response.

        Parses the error and provides helpful guidance for known issues.
        Always logs errors (not verbose-only).
        """
        try:
            # Try to parse error details
            error_text = response.text[:500] if response.text else ""

            # Check for "Tool name not known" error
            if "Tool name" in error_text and "not known" in error_text:
                # Extract tool name from config
                tool_name = self.config.name
                # Extract base name for filter (e.g., "demo.my_tool.1.0" -> "my_tool")
                parts = tool_name.split(".")
                filter_name = parts[1] if len(parts) >= 2 else tool_name

                _log(f"✗ Tool '{tool_name}' not registered on host", error=True)
                _log(f"  → Run: iq tool apply -t {filter_name}")
                return

            # Generic error output
            _log(f"✗ Poll error ({response.status_code}): {error_text[:100]}", error=True)

        except Exception:
            _log(f"✗ Poll error ({response.status_code})", error=True)

    @property
    def available_threads(self) -> int:
        """Number of threads available for new tasks."""
        return max(0, self.config.thread_count - self._running_tasks)

    @property
    def is_busy(self) -> bool:
        """Check if all threads are occupied."""
        return self._running_tasks >= self.config.thread_count

    # -------------------------------------------------------------------------
    # Pending Task Tracking
    # -------------------------------------------------------------------------

    def add_pending_task(self, task_id: str) -> None:
        """Add a task ID from task_added event."""
        if task_id:
            self._pending_tasks.add(task_id)

    def remove_pending_task(self, task_id: str) -> None:
        """Remove a task ID when task_reserved event arrives."""
        self._pending_tasks.discard(task_id)

    def has_pending_tasks(self) -> bool:
        """Check if there are known pending tasks."""
        return len(self._pending_tasks) > 0

    def clear_pending_tasks(self) -> None:
        """Clear all pending task IDs."""
        self._pending_tasks.clear()

    # -------------------------------------------------------------------------
    # Polling Decision
    # -------------------------------------------------------------------------

    def should_poll_immediately(self, idle_worker_count: int, active_worker_count: int) -> bool:
        """
        Decide if this worker should poll immediately (without waiting).

        Logic:
        - If busy: never poll immediately (already handled in caller, but safety check)
        - If only one active worker exists: poll immediately (we're the only one)
        - If idle <= 1 and we're not busy: we ARE the idle worker, poll immediately
        - Otherwise: random chance based on idle count

        Args:
            idle_worker_count: Number of idle workers for this tool
            active_worker_count: Total number of active workers for this tool

        Returns:
            True if this worker should poll immediately
        """
        if self._stopped or self.is_busy:
            return False

        # Only one active worker (us): poll immediately
        if active_worker_count <= 1:
            return True

        # We're not busy and idle_count <= 1: we are THE idle worker
        # (busy workers already returned False above)
        if idle_worker_count <= 1:
            return True

        # Multiple idle workers: random selection with 1/N chance to poll immediately
        # This spreads load so not all idle workers poll at once
        roll = random.randint(1, idle_worker_count)
        return roll == 1

    def get_random_wait_time(self, idle_worker_count: int) -> float:
        """
        Get a random wait time before polling.

        Higher idle worker counts get longer potential wait times to spread load.

        Args:
            idle_worker_count: Number of idle workers for this tool

        Returns:
            Wait time in seconds (1-15s range)
        """
        # Base range: 1-15 seconds, scaling with idle count
        min_wait = 1.0
        max_wait = min(15.0, 1.0 + idle_worker_count * 2.0)
        return random.uniform(min_wait, max_wait)

    def _get_aggregated_stats(self) -> tuple[str, int]:
        """
        Get aggregated stats from all workers for this tool.

        Returns:
            Tuple of (active_pids as comma-separated string, total thread count)
        """
        if not self.registry:
            # Fallback: only this worker's info
            return str(os.getpid()), self.config.thread_count

        pids, total_threads = self.registry.get_tool_stats(self.site_id, self.config.name)

        active_pids = ",".join(str(p) for p in pids) if pids else ""
        return active_pids, total_threads

    async def poll(self) -> list[str]:
        """
        Poll for new tasks from backend.

        Registers worker visibility and optionally receives tasks.
        Sends aggregated stats (active_pids, total thread_count) from all
        workers for this tool.

        Returns:
            List of task IDs (may be empty if no tasks available)
        """
        if self._stopped:
            return []

        # Get aggregated stats from all workers for this tool
        active_pids, total_threads = self._get_aggregated_stats()

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

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

        try:
            response = await self.api_client.post(
                "/api/v2/method/iq_core.iq_core.doctype.automation_worker.automation_worker.pull_task",
                json=payload,
            )

            self._log(f"← {response.status_code}", verbose_only=True)

            if response.status_code == 200:
                data = response.json()
                message = data.get("data", {})
                tasks = message.get("tasks", [])
                if tasks:
                    self._log(f"  {len(tasks)} task(s) received", verbose_only=True)
                    return tasks
                else:
                    self._log("  No task available", verbose_only=True)
            else:
                # Always log errors (not just in verbose mode)
                self._handle_poll_error(response)
            return []

        except Exception as e:
            import traceback

            self._log(f"Poll failed for {self.config.name}: {e}")
            self._log(traceback.format_exc(), verbose_only=True)
            return []

    async def execute_task(self, task_id: str) -> None:
        """
        Execute a task: reserve, run tool, report result.

        Args:
            task_id: Task ID/name from poll
        """
        if self._stopped:
            return

        if not task_id:
            return

        # Check if we have available threads
        if not self.available_threads:
            self._log(f"  Skip {task_id} - no threads available", verbose_only=True)
            return

        self._running_tasks += 1

        try:
            # Step 1: Reserve the task (mark as Started)
            # The reserve response contains the inputs for the tool
            self._log(f"→ POST reserve ({task_id})", verbose_only=True)
            reserve_response = await self.api_client.post(
                f"/api/v2/document/Automation Task/{task_id}/method/reserve",
                json={
                    "worker_id": self.worker_id,
                },
            )

            if reserve_response.status_code != 200:
                self._log(
                    f"← {reserve_response.status_code} - reservation failed", verbose_only=True
                )
                _log(
                    f"✗ Task reservation failed ({task_id}): {reserve_response.text[:100]}",
                    error=True,
                )
                return

            self._log(f"← {reserve_response.status_code} - reserved", verbose_only=True)

            # Extract inputs from reserve response
            reserve_data = reserve_response.json()
            self._log(f"  Reserve response: {reserve_data}", verbose_only=True)
            task_input = reserve_data.get("data", {}).get("input", {})
            self._log(f"  Task input: {task_input}", verbose_only=True)

            # Step 2: Execute the tool
            # execute() expects {"input": {...}, "general": {...}, ...}
            self._log(f"Executing: {task_id}")

            tool_result = await self.config.tool_class.execute({"input": task_input})

            # Step 3: Complete the task
            status = "Processed" if tool_result.success else "Failed"
            self._log(f"→ POST processed ({task_id}, status={status})", verbose_only=True)

            # Serialize output for JSON
            if tool_result.success and tool_result.output:
                # Pydantic model -> dict
                output_data = (
                    tool_result.output.model_dump()
                    if hasattr(tool_result.output, "model_dump")
                    else tool_result.output
                )
            else:
                output_data = tool_result.error

            # Try to send the processed result
            processed_success = await self._send_processed(task_id, output_data, status)

            if processed_success:
                status_icon = "✓" if tool_result.success else "✗"
                self._log(f"{status_icon} Completed: {task_id}")
            else:
                # Queue for retry
                self._queue_failed_output(task_id, output_data, status)
                self._log(f"⏳ Queued for retry: {task_id}")

        except Exception as e:
            self._log(f"Task error: {e}")
        finally:
            self._running_tasks -= 1

            # Notify parent to check for more tasks
            if self.on_task_complete and not self._stopped:
                try:
                    await self.on_task_complete()
                except Exception:
                    pass  # Don't fail task completion due to callback error

    def stop(self) -> None:
        """Stop the worker (prevents new tasks)."""
        self._stopped = True

    async def leave(self) -> bool:
        """
        Notify host that this worker is leaving.

        Sends remaining PIDs and thread count (excluding this process).
        Should be called during graceful shutdown.

        Returns:
            True if leave request was successful
        """
        # Get stats excluding current process
        current_pid = os.getpid()

        if self.registry:
            pids, total_threads = self.registry.get_tool_stats(
                self.site_id, self.config.name, exclude_pid=current_pid
            )
            active_pids = ",".join(str(p) for p in pids) if pids else ""
        else:
            # No registry - report empty (this was the only worker)
            active_pids = ""
            total_threads = 0

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

        self._log(f"→ POST leave ({self.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)

            if response.status_code != 200:
                self._log(f"Leave failed: {response.text[:100]}", verbose_only=True)
                return False

            self._log(f"← {self.config.name} left", verbose_only=True)
            return True

        except Exception as e:
            self._log(f"Leave request failed: {e}", verbose_only=True)
            return False

    # -------------------------------------------------------------------------
    # Output Queue / Retry
    # -------------------------------------------------------------------------

    async def _send_processed(
        self,
        task_id: str,
        output_data: dict | str | None,
        status: str,
    ) -> bool:
        """
        Send processed result to host.

        Args:
            task_id: Task identifier
            output_data: Output data or error message
            status: Task status (Processed/Failed)

        Returns:
            True if successful, False otherwise
        """
        from iq_sdk.client.output_queue import send_processed

        success, error = await send_processed(self.api_client, task_id, output_data, status)

        if success:
            self._log("← 200", verbose_only=True)
        else:
            _log(f"✗ Task completion failed ({task_id}): {error}", error=True)

        return success

    def _queue_failed_output(
        self,
        task_id: str,
        output_data: dict | str | None,
        status: str,
    ) -> None:
        """Queue a failed output for later retry."""
        if not self.output_queue:
            _log(f"✗ Cannot queue {task_id}: no output queue configured", error=True)
            return

        # Ensure output_data is a dict
        if isinstance(output_data, str):
            output_data = {"error": output_data}
        elif output_data is None:
            output_data = {}

        self.output_queue.queue_output(
            task_id=task_id,
            tool_name=self.config.name,
            output=output_data,
            status=status,
        )

    async def process_retry_queue(self) -> int:
        """
        Process pending outputs from the retry queue.

        Sends all due outputs and updates their retry status.

        Returns:
            Number of successfully retried outputs
        """
        if not self.output_queue:
            return 0

        due_outputs = self.output_queue.get_pending_outputs(
            tool_filter=self.config.name,
            due_only=True,
        )

        if not due_outputs:
            return 0

        self._log(f"→ Retrying {len(due_outputs)} queued output(s)", verbose_only=True)

        success_count = 0

        for queued in due_outputs:
            self._log(
                f"  → Retry {queued.task_id} (attempt {queued.retry_count + 1})", verbose_only=True
            )

            success = await self._send_processed(
                queued.task_id,
                queued.output,
                queued.status,
            )

            if success:
                # Remove from queue
                self.output_queue.remove_output(queued)
                self._log(f"  ✓ Retry successful: {queued.task_id}")
                success_count += 1
            else:
                # Increment retry count (may move to failed if max reached)
                updated = self.output_queue.increment_retry(queued)
                if updated is None:
                    _log(f"✗ Max retries exceeded, moved to failed: {queued.task_id}", error=True)
                else:
                    self._log(
                        f"  ⏳ Retry {queued.task_id} scheduled for {updated.retry_at}",
                        verbose_only=True,
                    )

        return success_count
