"""
Worker Registry - Track running worker processes.

Stores worker state in JSON files for process monitoring.
Each worker writes a registration file that includes PID for validation.
"""

from __future__ import annotations

import json
import os
from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path


def get_registry_dir() -> Path:
    """
    Get the worker registry directory.

    Uses the project root from the active project, or cwd as fallback.
    """
    from iq_sdk.config import get_project_root

    project_root = get_project_root()
    return project_root / ".iq-workers"


@dataclass
class WorkerEntry:
    """A registered worker entry."""

    worker_id: str
    site_id: str
    tool_name: str  # Full name like "demo.my_tool.1.0"
    pid: int
    threads: int
    started_at: str
    host: str = ""

    def to_dict(self) -> dict:
        """Convert to dictionary."""
        return {
            "worker_id": self.worker_id,
            "site_id": self.site_id,
            "tool_name": self.tool_name,
            "pid": self.pid,
            "threads": self.threads,
            "started_at": self.started_at,
            "host": self.host,
        }

    @classmethod
    def from_dict(cls, data: dict) -> WorkerEntry:
        """Create from dictionary."""
        return cls(
            worker_id=data.get("worker_id", ""),
            site_id=data.get("site_id", ""),
            tool_name=data.get("tool_name", ""),
            pid=data.get("pid", 0),
            threads=data.get("threads", 1),
            started_at=data.get("started_at", ""),
            host=data.get("host", ""),
        )

    def get_filter_name(self) -> str:
        """
        Get the name to use for --tool filter when restarting.

        Extracts tool base name from full name (e.g., "demo.my_tool.1.0" -> "my_tool").
        """
        # Full name format: "project.tool_name.version" (e.g., "demo.my_tool.1.0")
        parts = self.tool_name.split(".")
        if len(parts) >= 2:
            return parts[1]  # Return tool_name part
        return self.tool_name

    def is_alive(self) -> bool:
        """Check if the worker process is still running."""
        try:
            # Send signal 0 to check if process exists
            os.kill(self.pid, 0)
            return True
        except (OSError, ProcessLookupError):
            return False


@dataclass
class WorkerRegistry:
    """
    Registry for tracking running workers.

    Each worker process writes a JSON file with PID in the filename to support
    multiple instances of the same tool running in parallel.

    Filename format: {site_id}_{tool_name}_{pid}.json
    """

    registry_dir: Path = field(default_factory=get_registry_dir)

    def __post_init__(self) -> None:
        """Ensure registry directory exists."""
        self.registry_dir.mkdir(parents=True, exist_ok=True)

    def _get_entry_path(self, site_id: str, tool_name: str, pid: int | None = None) -> Path:
        """
        Get path for a worker entry file.

        Args:
            site_id: Site identifier
            tool_name: Tool name
            pid: Process ID (uses current PID if not specified)
        """
        # Sanitize tool_name for filename
        safe_name = tool_name.replace(".", "_").replace("/", "_")
        actual_pid = pid if pid is not None else os.getpid()
        return self.registry_dir / f"{site_id}_{safe_name}_{actual_pid}.json"

    def _is_worker_file(self, path: Path) -> bool:
        """Check if a file is a worker entry (not a monitor file)."""
        return path.suffix == ".json" and not path.name.startswith("_monitor")

    def register(
        self,
        worker_id: str,
        site_id: str,
        tool_name: str,
        threads: int,
        host: str = "",
    ) -> WorkerEntry:
        """
        Register a worker for a tool.

        Args:
            worker_id: Worker identifier
            site_id: Site identifier
            tool_name: Full tool name (e.g. "demo.my_tool.1.0")
            threads: Number of threads
            host: Host URL

        Returns:
            The registered WorkerEntry
        """
        pid = os.getpid()
        entry = WorkerEntry(
            worker_id=worker_id,
            site_id=site_id,
            tool_name=tool_name,
            pid=pid,
            threads=threads,
            started_at=datetime.now().isoformat(),
            host=host,
        )

        entry_path = self._get_entry_path(site_id, tool_name, pid)
        with open(entry_path, "w") as f:
            json.dump(entry.to_dict(), f, indent=2)

        return entry

    def unregister(self, site_id: str, tool_name: str, pid: int | None = None) -> None:
        """
        Unregister a worker.

        Args:
            site_id: Site identifier
            tool_name: Full tool name
            pid: Process ID (uses current PID if not specified)
        """
        actual_pid = pid if pid is not None else os.getpid()
        entry_path = self._get_entry_path(site_id, tool_name, actual_pid)
        if entry_path.exists():
            entry_path.unlink()

    def unregister_all(self, site_id: str | None = None) -> None:
        """
        Unregister all workers, optionally filtered by site.

        Args:
            site_id: Optional site ID to filter by
        """
        for entry_path in self.registry_dir.glob("*.json"):
            if not self._is_worker_file(entry_path):
                continue
            if site_id:
                # Check if file belongs to this site
                if not entry_path.name.startswith(f"{site_id}_"):
                    continue
            try:
                entry_path.unlink()
            except OSError:
                pass

    def get_entry(self, site_id: str, tool_name: str, pid: int | None = None) -> WorkerEntry | None:
        """
        Get a worker entry if it exists and is alive.

        Args:
            site_id: Site identifier
            tool_name: Full tool name
            pid: Process ID (uses current PID if not specified)

        Returns:
            WorkerEntry if exists and alive, None otherwise
        """
        actual_pid = pid if pid is not None else os.getpid()
        entry_path = self._get_entry_path(site_id, tool_name, actual_pid)
        if not entry_path.exists():
            return None

        try:
            with open(entry_path) as f:
                data = json.load(f)
            entry = WorkerEntry.from_dict(data)

            # Validate PID is still alive
            if not entry.is_alive():
                return None

            return entry
        except (json.JSONDecodeError, OSError):
            return None

    def has_active_monitors(self) -> bool:
        """
        Check if there are any active monitor processes.

        Used to decide whether to clean up stale worker entries
        (monitors handle restarts, so we shouldn't delete their tracking files).

        Returns:
            True if at least one monitor is alive
        """
        mon_reg = MonitorRegistry(self.registry_dir)
        return len(mon_reg.get_alive_monitors()) > 0

    def get_all_entries(
        self, site_id: str | None = None, cleanup_stale: bool = False
    ) -> list[WorkerEntry]:
        """
        Get all registered workers.

        Args:
            site_id: Optional site ID to filter by
            cleanup_stale: If True, remove entries for dead processes.
                          Should be False when monitors are active.

        Returns:
            List of WorkerEntry objects (both alive and dead if cleanup_stale=False)
        """
        entries = []

        for entry_path in self.registry_dir.glob("*.json"):
            if not self._is_worker_file(entry_path):
                continue

            try:
                with open(entry_path) as f:
                    data = json.load(f)
                entry = WorkerEntry.from_dict(data)

                # Filter by site if specified
                if site_id and entry.site_id != site_id:
                    continue

                if cleanup_stale and not entry.is_alive():
                    # Clean up stale entry
                    try:
                        entry_path.unlink()
                    except OSError:
                        pass
                    continue

                entries.append(entry)

            except (json.JSONDecodeError, OSError):
                # Remove corrupted files only if cleanup requested
                if cleanup_stale:
                    try:
                        entry_path.unlink()
                    except OSError:
                        pass

        return entries

    def get_alive_entries(self, site_id: str | None = None) -> list[WorkerEntry]:
        """
        Get only alive worker entries.

        Args:
            site_id: Optional site ID to filter by

        Returns:
            List of alive WorkerEntry objects
        """
        return [e for e in self.get_all_entries(site_id, cleanup_stale=False) if e.is_alive()]

    def cleanup_stale(self) -> int:
        """
        Remove all stale entries (dead processes).

        Only call this when no monitors are active!

        Returns:
            Number of entries cleaned up
        """
        cleaned = 0
        for entry_path in self.registry_dir.glob("*.json"):
            if not self._is_worker_file(entry_path):
                continue

            try:
                with open(entry_path) as f:
                    data = json.load(f)
                entry = WorkerEntry.from_dict(data)

                if not entry.is_alive():
                    entry_path.unlink()
                    cleaned += 1

            except (json.JSONDecodeError, OSError):
                try:
                    entry_path.unlink()
                    cleaned += 1
                except OSError:
                    pass

        return cleaned

    def get_all_entries_raw(self, site_id: str | None = None) -> list[tuple[Path, WorkerEntry]]:
        """
        Get all entries with their file paths (without cleaning up stale).

        Used by monitor to decide whether to restart dead workers.

        Args:
            site_id: Optional site ID to filter by

        Returns:
            List of (path, WorkerEntry) tuples
        """
        entries = []

        for entry_path in self.registry_dir.glob("*.json"):
            if not self._is_worker_file(entry_path):
                continue

            try:
                with open(entry_path) as f:
                    data = json.load(f)
                entry = WorkerEntry.from_dict(data)

                # Filter by site if specified
                if site_id and entry.site_id != site_id:
                    continue

                entries.append((entry_path, entry))

            except (json.JSONDecodeError, OSError):
                pass

        return entries

    def get_dead_entries(self, site_id: str | None = None) -> list[tuple[Path, WorkerEntry]]:
        """
        Get entries for dead workers (PID not running).

        Args:
            site_id: Optional site ID to filter by

        Returns:
            List of (path, WorkerEntry) tuples for dead workers
        """
        dead = []
        for path, entry in self.get_all_entries_raw(site_id):
            if not entry.is_alive():
                dead.append((path, entry))
        return dead

    def get_tool_stats(
        self, site_id: str, tool_name: str, exclude_pid: int | None = None
    ) -> tuple[list[int], int]:
        """
        Get PIDs and total thread count for a specific tool.

        Used for pull_task and leave requests to aggregate worker info.

        Args:
            site_id: Site identifier
            tool_name: Full tool name (e.g. "demo.my_tool.1.0")
            exclude_pid: Optional PID to exclude (for leave requests)

        Returns:
            Tuple of (list of alive PIDs, total thread count)
        """
        pids: list[int] = []
        total_threads = 0

        for entry in self.get_alive_entries(site_id):
            if entry.tool_name == tool_name:
                if exclude_pid is not None and entry.pid == exclude_pid:
                    continue
                pids.append(entry.pid)
                total_threads += entry.threads

        return pids, total_threads

    def kill_worker(self, entry: WorkerEntry, force: bool = False) -> bool:
        """
        Kill a worker process.

        Args:
            entry: Worker entry to kill
            force: If True, use SIGKILL (-9), otherwise SIGTERM

        Returns:
            True if signal was sent successfully
        """
        import signal

        try:
            sig = signal.SIGKILL if force else signal.SIGTERM
            os.kill(entry.pid, sig)
            return True
        except (OSError, ProcessLookupError):
            return False

    def remove_entry_file(self, entry: WorkerEntry) -> bool:
        """
        Remove an entry file for a specific worker.

        Args:
            entry: Worker entry whose file to remove

        Returns:
            True if file was removed
        """
        entry_path = self._get_entry_path(entry.site_id, entry.tool_name, entry.pid)
        if entry_path.exists():
            try:
                entry_path.unlink()
                return True
            except OSError:
                pass
        return False


# Monitor registry for tracking monitor processes
MONITOR_FILE_NAME = "_monitor.json"


@dataclass
class MonitorEntry:
    """Entry for monitor process."""

    pid: int
    partner_pid: int  # PID of the partner monitor (for dual-monitor pattern)
    started_at: str

    def to_dict(self) -> dict:
        return {
            "pid": self.pid,
            "partner_pid": self.partner_pid,
            "started_at": self.started_at,
        }

    @classmethod
    def from_dict(cls, data: dict) -> MonitorEntry:
        return cls(
            pid=data.get("pid", 0),
            partner_pid=data.get("partner_pid", 0),
            started_at=data.get("started_at", ""),
        )

    def is_alive(self) -> bool:
        """Check if this monitor process is running."""
        try:
            os.kill(self.pid, 0)
            return True
        except (OSError, ProcessLookupError):
            return False

    def is_partner_alive(self) -> bool:
        """Check if the partner monitor is running."""
        if self.partner_pid == 0:
            return False
        try:
            os.kill(self.partner_pid, 0)
            return True
        except (OSError, ProcessLookupError):
            return False


class MonitorRegistry:
    """Registry for monitor processes (dual-monitor pattern)."""

    def __init__(self, registry_dir: Path | None = None):
        self.registry_dir = registry_dir or get_registry_dir()
        self.registry_dir.mkdir(parents=True, exist_ok=True)

    def _get_monitor_path(self, index: int) -> Path:
        """Get path for monitor file (0 or 1 for dual monitors)."""
        return self.registry_dir / f"_monitor_{index}.json"

    def register(self, index: int, partner_pid: int = 0) -> MonitorEntry:
        """Register a monitor process."""
        entry = MonitorEntry(
            pid=os.getpid(),
            partner_pid=partner_pid,
            started_at=datetime.now().isoformat(),
        )
        path = self._get_monitor_path(index)
        with open(path, "w") as f:
            json.dump(entry.to_dict(), f, indent=2)
        return entry

    def update_partner(self, index: int, partner_pid: int) -> None:
        """Update the partner PID for a monitor."""
        path = self._get_monitor_path(index)
        if path.exists():
            try:
                with open(path) as f:
                    data = json.load(f)
                data["partner_pid"] = partner_pid
                with open(path, "w") as f:
                    json.dump(data, f, indent=2)
            except (json.JSONDecodeError, OSError):
                pass

    def unregister(self, index: int) -> None:
        """Unregister a monitor."""
        path = self._get_monitor_path(index)
        if path.exists():
            path.unlink()

    def get_monitors(self) -> list[MonitorEntry]:
        """Get all monitor entries."""
        entries = []
        for i in range(2):
            path = self._get_monitor_path(i)
            if path.exists():
                try:
                    with open(path) as f:
                        data = json.load(f)
                    entries.append(MonitorEntry.from_dict(data))
                except (json.JSONDecodeError, OSError):
                    pass
        return entries

    def get_alive_monitors(self) -> list[MonitorEntry]:
        """Get only alive monitors."""
        return [m for m in self.get_monitors() if m.is_alive()]

    def kill_monitors(self, force: bool = False) -> int:
        """Kill all monitor processes."""
        import signal

        killed = 0
        sig = signal.SIGKILL if force else signal.SIGTERM

        for m in self.get_monitors():
            if m.is_alive():
                try:
                    os.kill(m.pid, sig)
                    killed += 1
                except (OSError, ProcessLookupError):
                    pass

        # Clean up files
        for i in range(2):
            path = self._get_monitor_path(i)
            if path.exists():
                try:
                    path.unlink()
                except OSError:
                    pass

        return killed
