"""Worker single-instance ownership and safe stale-worker cleanup."""

from __future__ import annotations

from collections.abc import Callable
from collections.abc import Mapping
from dataclasses import dataclass
import os
from pathlib import Path
import re
import shlex
import signal
import subprocess
import time
import uuid


PID_FILE_NAME = "worker.pid"
STALE_PID_FILE_NAME = "worker.stale.pid"
_CLAIM_ATTEMPTS = 12
_TERMINATE_GRACE_SECONDS = 1.0
_TERMINATE_POLL_SECONDS = 0.05


@dataclass(frozen=True)
class ProcessTermination:
    """The result of requesting one process to stop."""

    pid: int
    stopped: bool
    method: str
    error: str | None = None


@dataclass(frozen=True)
class WorkerInstanceClaim:
    """The PID-file ownership state for one worker process."""

    path: Path
    pid: int
    enabled: bool
    warnings: tuple[str, ...] = ()


@dataclass(frozen=True)
class WorkerProcess:
    """A process whose command line was available from the operating system."""

    pid: int
    command: str
    parent_command: str | None = None


@dataclass(frozen=True)
class StaleWorkerCleanup:
    """Summary returned by the fixed ``kill-stale-workers`` control action."""

    method: str
    terminations: tuple[ProcessTermination, ...]
    skipped_pids: tuple[int, ...]
    remaining_workers: int | None
    warnings: tuple[str, ...] = ()

    def report(self) -> str:
        """Return a compact, machine-readable-enough control response."""

        terminated = [result for result in self.terminations if result.stopped]
        failed = [result for result in self.terminations if not result.stopped]
        terminated_text = (
            ", ".join(f"{result.pid} ({result.method})" for result in terminated)
            if terminated
            else "none"
        )
        failed_text = (
            ", ".join(f"{result.pid} ({result.error or result.method})" for result in failed)
            if failed
            else "none"
        )
        skipped_text = ", ".join(str(pid) for pid in self.skipped_pids) or "none"
        lines = [
            "control completed: kill-stale-workers",
            f"method: {self.method}",
            f"terminated pids: {terminated_text}",
            f"failed pids: {failed_text}",
            f"skipped pids: {skipped_text}",
            "remaining workers: "
            f"{self.remaining_workers if self.remaining_workers is not None else 'unknown'}",
        ]
        lines.extend(f"warning: {warning}" for warning in self.warnings)
        return "\n".join(lines)


def worker_pid_path(home: Path) -> Path:
    """Return the one stable PID-file location for a local worker."""

    return home / PID_FILE_NAME


def stale_worker_pid_path(home: Path) -> Path:
    """Return the previous worker PID record retained after a failed takeover."""

    return home / STALE_PID_FILE_NAME


def single_instance_enabled(environ: Mapping[str, str] | None = None) -> bool:
    """Return whether the default worker single-instance protection is enabled."""

    environment = os.environ if environ is None else environ
    return environment.get("A2A_SINGLE_INSTANCE", "1").strip() != "0"


def claim_worker_instance(
    home: Path,
    *,
    pid: int | None = None,
    environ: Mapping[str, str] | None = None,
    is_pid_running: Callable[[int], bool] | None = None,
    command_for_pid: Callable[[int], str | None] | None = None,
    terminate: Callable[[int], ProcessTermination] | None = None,
) -> WorkerInstanceClaim:
    """Record this worker PID after asking the previous recorded worker to leave.

    Creating the PID file with ``O_EXCL`` serializes normal simultaneous starts:
    the worker which observes another PID removes that record, stops the older
    worker, then retries the exclusive create.  If stopping an old PID fails,
    the new worker still claims the file and starts, while preserving a warning.
    """

    current_pid = os.getpid() if pid is None else pid
    path = worker_pid_path(home)
    if not single_instance_enabled(environ):
        return WorkerInstanceClaim(path=path, pid=current_pid, enabled=False)

    home.mkdir(parents=True, exist_ok=True)
    running = is_pid_running or pid_is_running
    command_lookup = command_for_pid or process_command
    terminator = terminate or terminate_process
    warnings: list[str] = []

    for _ in range(_CLAIM_ATTEMPTS):
        try:
            _create_pid_file(path, current_pid)
        except FileExistsError:
            previous_pid = read_worker_pid(path)
            if previous_pid is None:
                _remove_pid_file(path, expected_pid=None)
                continue
            if previous_pid == current_pid:
                return WorkerInstanceClaim(
                    path=path,
                    pid=current_pid,
                    enabled=True,
                    warnings=tuple(warnings),
                )
            if not running(previous_pid):
                _remove_pid_file(path, expected_pid=previous_pid)
                continue

            command = command_lookup(previous_pid)
            # The PID file is worker-specific evidence.  A matching normal
            # a2a-agent run command is enough for startup takeover, including
            # Windows wrappers whose command line does not expose A2A_HOME.
            if command is not None and not is_worker_command(command):
                warnings.append(
                    f"worker pid file pointed to unconfirmed pid={previous_pid}; "
                    "record was replaced without stopping it"
                )
                _remove_pid_file(path, expected_pid=previous_pid)
                continue

            stale_path = stale_worker_pid_path(home)
            stale_recorded = True
            try:
                _replace_pid_file(stale_path, previous_pid)
            except OSError as exc:
                stale_recorded = False
                warnings.append(
                    f"could not record previous worker pid={previous_pid}: {exc}"
                )
            result = terminator(previous_pid)
            if not result.stopped:
                warnings.append(
                    f"previous worker cleanup failed for pid={previous_pid}: "
                    f"{result.error or result.method}"
                )
            elif stale_recorded:
                _remove_pid_file(stale_path, expected_pid=previous_pid)
            _remove_pid_file(path, expected_pid=previous_pid)
            continue
        else:
            return WorkerInstanceClaim(
                path=path,
                pid=current_pid,
                enabled=True,
                warnings=tuple(warnings),
            )

    # A malformed or repeatedly replaced PID record must not prevent a new
    # worker from coming online.  This is intentionally a last-resort path.
    warnings.append("worker pid file changed repeatedly during startup; record replaced")
    _replace_pid_file(path, current_pid)
    return WorkerInstanceClaim(
        path=path,
        pid=current_pid,
        enabled=True,
        warnings=tuple(warnings),
    )


def release_worker_instance(claim: WorkerInstanceClaim) -> None:
    """Remove the PID file only when it still belongs to this worker."""

    if claim.enabled:
        _remove_pid_file(claim.path, expected_pid=claim.pid)


def owns_worker_pid(home: Path, *, pid: int | None = None) -> bool:
    """Return whether this process is the current worker.pid owner."""

    current_pid = os.getpid() if pid is None else pid
    return read_worker_pid(worker_pid_path(home)) == current_pid


def read_worker_pid(path: Path) -> int | None:
    """Read a valid positive PID from ``path`` without raising for stale data."""

    try:
        raw = path.read_text(encoding="utf-8").strip()
    except OSError:
        return None
    try:
        pid = int(raw)
    except ValueError:
        return None
    return pid if pid > 0 else None


def pid_is_running(pid: int) -> bool:
    """Check only whether a PID currently exists, without terminating it."""

    if pid <= 0:
        return False
    try:
        os.kill(pid, 0)
    except ProcessLookupError:
        return False
    except PermissionError:
        return True
    except OSError:
        return False
    return True


def process_command(
    pid: int,
    *,
    platform_name: str | None = None,
    run_command: Callable[[list[str]], subprocess.CompletedProcess[str] | None] | None = None,
) -> str | None:
    """Return a PID command line when the platform exposes one."""

    platform = os.name if platform_name is None else platform_name
    run = run_command or _run_command
    if platform == "nt":
        completed = run(
            [
                "wmic",
                "process",
                "where",
                f"ProcessId={pid}",
                "get",
                "CommandLine",
                "/value",
            ]
        )
        if completed is None or completed.returncode != 0:
            return None
        for line in completed.stdout.splitlines():
            if line.startswith("CommandLine="):
                value = line.partition("=")[2].strip()
                return value or None
        return None

    completed = run(["ps", "-p", str(pid), "-o", "command="])
    if completed is None or completed.returncode != 0:
        return None
    command = completed.stdout.strip()
    return command or None


def is_worker_command(command: str) -> bool:
    """Recognize only an actual a2a-agent ``run`` argv sequence.

    A substring search would mistake a commander command such as
    ``a2a-agent send device 'a2a-agent run'`` for a worker.  We inspect the
    executable and its immediate subcommand instead, without executing it.
    """

    try:
        tokens = [token.strip("\"'") for token in shlex.split(command, posix=False)]
    except ValueError:
        return False
    while tokens and _is_environment_assignment(tokens[0]):
        tokens.pop(0)
    if tokens[:1] == ["env"]:
        tokens.pop(0)
        while tokens and _is_environment_assignment(tokens[0]):
            tokens.pop(0)
    if not tokens:
        return False

    executable = tokens[0].replace("\\", "/").rsplit("/", maxsplit=1)[-1].lower()
    if executable in {"a2a-agent", "a2a-agent.exe", "a2a_agent", "a2a_agent.py"}:
        return len(tokens) >= 2 and tokens[1].lower() == "run"
    if executable.startswith("python"):
        return (
            len(tokens) >= 4
            and tokens[1] == "-m"
            and tokens[2].lower() in {"a2a_agent", "a2a_agent.cli"}
            and tokens[3].lower() == "run"
        )
    return False


def is_worker_command_for_home(
    command: str,
    home: Path,
    parent_command: str | None = None,
) -> bool:
    """Require command-line evidence that a run process belongs to this home.

    A commander uses the same ``a2a-agent run`` syntax, so it must not be
    treated as a worker solely from that shared command.  Confirm the current
    ``A2A_HOME`` or executable path, or (on Windows Startup launches) the
    parent ``start-worker.cmd`` path under this home.
    """

    if not is_worker_command(command):
        return False
    home_text = str(home.expanduser().resolve()).replace("\\", "/").lower().rstrip("/")
    command_text = command.replace("\\", "/").lower()
    declared_home = _declared_a2a_home(command_text)
    if declared_home is not None:
        return declared_home == home_text
    if home_text in command_text:
        return True
    if parent_command is None:
        return False
    parent_text = parent_command.replace("\\", "/").lower()
    return home_text in parent_text and "start-worker.cmd" in parent_text


def list_worker_processes(
    *,
    platform_name: str | None = None,
    run_command: Callable[[list[str]], subprocess.CompletedProcess[str] | None] | None = None,
) -> list[WorkerProcess] | None:
    """List a2a-agent ``run`` processes, or ``None`` if enumeration fails."""

    platform = os.name if platform_name is None else platform_name
    run = run_command or _run_command
    if platform == "nt":
        completed = run(
            [
                "wmic",
                "process",
                "get",
                "CommandLine,ParentProcessId,ProcessId",
                "/value",
            ]
        )
        if completed is None or completed.returncode != 0:
            return None
        return _parse_windows_processes(completed.stdout)

    completed = run(["ps", "-axeww", "-o", "pid=,command="])
    if completed is None or completed.returncode != 0:
        return None
    return _parse_posix_processes(completed.stdout)


def cleanup_stale_workers(
    home: Path,
    *,
    current_pid: int | None = None,
    list_processes: Callable[[], list[WorkerProcess] | None] | None = None,
    command_for_pid: Callable[[int], str | None] | None = None,
    terminate: Callable[[int], ProcessTermination] | None = None,
) -> StaleWorkerCleanup:
    """Stop other verified worker processes while leaving this worker untouched."""

    own_pid = os.getpid() if current_pid is None else current_pid
    processes_provider = list_processes or list_worker_processes
    command_lookup = command_for_pid or process_command
    terminator = terminate or terminate_process
    recorded_pid = read_worker_pid(worker_pid_path(home))

    # A duplicate subscriber can receive the same MQTT control task.  Only
    # the process that currently owns worker.pid may perform cleanup; an older
    # ghost reports its state but cannot race the owner into killing it.
    if recorded_pid is not None and recorded_pid != own_pid:
        listed = processes_provider()
        return StaleWorkerCleanup(
            method="pid-file-owner",
            terminations=(),
            skipped_pids=(),
            remaining_workers=_remaining_worker_count(
                listed,
                home,
                own_pid,
                known_owner_pid=recorded_pid,
            ),
            warnings=(
                f"cleanup skipped: worker.pid is owned by pid={recorded_pid}",
            ),
        )

    processes = processes_provider()
    method = "process-list"
    skipped: list[int] = []
    fallback_stale_pid: int | None = None

    if processes is None:
        # On platforms without command-line enumeration, PID-file fallback is
        # deliberately conservative: unlike startup takeover, remote control
        # never stops an unconfirmed process.
        method = "pid-file"
        processes = []
        fallback_stale_pid = read_worker_pid(stale_worker_pid_path(home))
        if fallback_stale_pid is not None and fallback_stale_pid != own_pid:
            command = command_lookup(fallback_stale_pid)
            if command is not None and is_worker_command(command):
                processes.append(WorkerProcess(pid=fallback_stale_pid, command=command))
            else:
                skipped.append(fallback_stale_pid)

    candidates: list[WorkerProcess] = []
    for process in processes:
        if process.pid == own_pid or not is_worker_command(process.command):
            continue
        if process.pid == fallback_stale_pid or is_worker_command_for_home(
            process.command,
            home,
            process.parent_command,
        ):
            candidates.append(process)
        else:
            skipped.append(process.pid)
    terminations = tuple(terminator(process.pid) for process in candidates)
    if fallback_stale_pid is not None and any(
        result.pid == fallback_stale_pid and result.stopped for result in terminations
    ):
        _remove_pid_file(stale_worker_pid_path(home), expected_pid=fallback_stale_pid)

    remaining_processes = processes_provider()
    if remaining_processes is None:
        remaining_workers = None
    else:
        remaining_workers = _remaining_worker_count(remaining_processes, home, own_pid)

    warning_lines = [
        f"pid={result.pid} was not stopped: {result.error or result.method}"
        for result in terminations
        if not result.stopped
    ]
    if remaining_workers is None:
        warning_lines.append("worker process re-enumeration was unavailable")
    return StaleWorkerCleanup(
        method=method,
        terminations=terminations,
        skipped_pids=tuple(skipped),
        remaining_workers=remaining_workers,
        warnings=tuple(warning_lines),
    )


def terminate_process(
    pid: int,
    *,
    platform_name: str | None = None,
    is_pid_running: Callable[[int], bool] | None = None,
    send_signal: Callable[[int, int], None] | None = None,
    run_command: Callable[[list[str]], subprocess.CompletedProcess[str] | None] | None = None,
    sleep: Callable[[float], None] | None = None,
    grace_seconds: float = _TERMINATE_GRACE_SECONDS,
) -> ProcessTermination:
    """Stop one process with the platform's standard, non-shell mechanism."""

    platform = os.name if platform_name is None else platform_name
    running = is_pid_running or pid_is_running
    if not running(pid):
        return ProcessTermination(pid=pid, stopped=True, method="already-exited")

    if platform == "nt":
        run = run_command or _run_command
        completed = run(["taskkill", "/PID", str(pid), "/F"])
        if completed is None:
            return ProcessTermination(
                pid=pid,
                stopped=False,
                method="taskkill",
                error="taskkill could not be started",
            )
        if completed.returncode != 0:
            error = (completed.stderr or completed.stdout).strip() or "taskkill failed"
            return ProcessTermination(pid=pid, stopped=False, method="taskkill", error=error)
        still_running = running(pid)
        return ProcessTermination(
            pid=pid,
            stopped=not still_running,
            method="taskkill",
            error=None if not still_running else "process still running after taskkill",
        )

    signal_sender = send_signal or os.kill
    sleeper = sleep or time.sleep
    try:
        signal_sender(pid, signal.SIGTERM)
    except OSError as exc:
        return ProcessTermination(pid=pid, stopped=False, method="SIGTERM", error=str(exc))

    deadline = time.monotonic() + max(0.0, grace_seconds)
    while running(pid) and time.monotonic() < deadline:
        sleeper(_TERMINATE_POLL_SECONDS)
    if not running(pid):
        return ProcessTermination(pid=pid, stopped=True, method="SIGTERM")

    try:
        signal_sender(pid, signal.SIGKILL)
    except OSError as exc:
        return ProcessTermination(
            pid=pid,
            stopped=False,
            method="SIGTERM+SIGKILL",
            error=str(exc),
        )
    still_running = running(pid)
    return ProcessTermination(
        pid=pid,
        stopped=not still_running,
        method="SIGTERM+SIGKILL",
        error=None if not still_running else "process still running after SIGKILL",
    )


def _create_pid_file(path: Path, pid: int) -> None:
    # Write a complete PID to a private file first, then create the public
    # name with link(2).  Unlike O_EXCL followed by write, readers can never
    # observe a newly-created but still-empty PID file during a startup race.
    temporary = path.with_name(f".{path.name}.{pid}.{uuid.uuid4().hex}.tmp")
    descriptor = os.open(temporary, os.O_CREAT | os.O_EXCL | os.O_WRONLY, 0o600)
    try:
        os.write(descriptor, f"{pid}\n".encode("ascii"))
    finally:
        os.close(descriptor)
    try:
        os.link(temporary, path)
    finally:
        temporary.unlink(missing_ok=True)


def _replace_pid_file(path: Path, pid: int) -> None:
    temporary = path.with_name(f".{path.name}.{pid}.{uuid.uuid4().hex}.tmp")
    temporary.write_text(f"{pid}\n", encoding="ascii")
    os.replace(temporary, path)


def _remove_pid_file(path: Path, *, expected_pid: int | None) -> None:
    if expected_pid is not None and read_worker_pid(path) != expected_pid:
        return
    try:
        path.unlink()
    except FileNotFoundError:
        return
    except OSError:
        return


def _run_command(args: list[str]) -> subprocess.CompletedProcess[str] | None:
    try:
        return subprocess.run(args, capture_output=True, text=True, check=False)
    except OSError:
        return None


def _parse_posix_processes(output: str) -> list[WorkerProcess]:
    processes: list[WorkerProcess] = []
    for line in output.splitlines():
        fields = line.strip().split(maxsplit=1)
        if len(fields) != 2:
            continue
        try:
            pid = int(fields[0])
        except ValueError:
            continue
        command = fields[1]
        if is_worker_command(command):
            processes.append(WorkerProcess(pid=pid, command=command))
    return processes


def _parse_windows_processes(output: str) -> list[WorkerProcess]:
    records: list[tuple[int, str, int | None]] = []
    command: str | None = None
    pid: int | None = None
    parent_pid: int | None = None

    def add_current() -> None:
        nonlocal command, pid, parent_pid
        if pid is not None and command is not None:
            records.append((pid, command, parent_pid))
        command = None
        pid = None
        parent_pid = None

    for line in output.splitlines() + [""]:
        line = line.strip()
        if not line:
            add_current()
            continue
        key, separator, value = line.partition("=")
        if not separator:
            continue
        if key == "CommandLine":
            command = value.strip()
        elif key == "ProcessId":
            try:
                pid = int(value.strip())
            except ValueError:
                pid = None
        elif key == "ParentProcessId":
            try:
                parent_pid = int(value.strip())
            except ValueError:
                parent_pid = None

    commands_by_pid = {record_pid: record_command for record_pid, record_command, _ in records}
    return [
        WorkerProcess(
            pid=record_pid,
            command=record_command,
            parent_command=commands_by_pid.get(record_parent_pid),
        )
        for record_pid, record_command, record_parent_pid in records
        if is_worker_command(record_command)
    ]


def _declared_a2a_home(command: str) -> str | None:
    match = re.search(r"(?:^|\s)a2a_home=(?:\"([^\"]+)\"|'([^']+)'|([^\s]+))", command)
    if match is None:
        return None
    value = next(part for part in match.groups() if part is not None)
    return value.replace("\\", "/").rstrip("/")


def _is_environment_assignment(token: str) -> bool:
    return re.fullmatch(r"[A-Za-z_][A-Za-z0-9_]*=.*", token) is not None


def _remaining_worker_count(
    processes: list[WorkerProcess] | None,
    home: Path,
    own_pid: int,
    *,
    known_owner_pid: int | None = None,
) -> int | None:
    if processes is None:
        return None
    listed_workers = sum(
        1
        for process in processes
        if process.pid == own_pid
        or process.pid == known_owner_pid
        or is_worker_command_for_home(process.command, home, process.parent_command)
    )
    return max(1, listed_workers)
