"""ECS task protection wrapper.

Calls the local ECS agent's `/task-protection/v1/state` endpoint to opt the
running task out of scale-in / replacement while critical work is in flight.
"""

import functools
import json
import os
import urllib.error
import urllib.request
from collections.abc import Callable
from typing import Any


def _set_task_protection(log: Any, enabled: bool) -> bool:
    """Toggle ECS task protection via the agent endpoint.

    Returns True on success, or when running outside ECS (no agent URI).
    Returns False on HTTP/network failure so callers can skip protected work.
    """
    agent_uri = os.environ.get("ECS_AGENT_URI")
    if not agent_uri:
        return True
    req = urllib.request.Request(
        f"{agent_uri}/task-protection/v1/state",
        data=json.dumps({"ProtectionEnabled": enabled}).encode("utf-8"),
        method="PUT",
        headers={"Content-Type": "application/json"},
    )
    try:
        with urllib.request.urlopen(req, timeout=5) as resp:
            resp.read()
        return True
    except urllib.error.URLError as exc:
        log.error(f"Failed to set task protection enabled={enabled}: {exc}")
        return False


def with_task_protection(func: Callable[..., Any]) -> Callable[..., bool]:
    """Run `func(log, ...)` under ECS task protection; skip + log if unavailable.

    Returns True if the wrapped function ran, False if protection could not be
    acquired and the call was skipped.
    """

    @functools.wraps(func)
    def wrapper(log: Any, *args: Any, **kwargs: Any) -> bool:
        if not _set_task_protection(log, True):
            log.warning(f"Could not acquire task protection; skipping {func.__name__}")
            return False
        try:
            func(log, *args, **kwargs)
        finally:
            _set_task_protection(log, False)
        return True

    return wrapper