import http.client
import json
import logging
import os
from collections.abc import Mapping
from enum import StrEnum
from typing import Final, cast
from urllib.request import Request, urlopen

_logger = logging.getLogger(__name__)

_ENDPOINT_PATH: Final = "/task-protection/v1/state"
_REQUEST_TIMEOUT_SECONDS: Final = 5


# TASK_STOPPING_OR_STOPPED is returned by the agent endpoint during rolling
# deploys and scale-in, despite not being in the AWS docs' reason-code list.
# https://docs.aws.amazon.com/AmazonECS/latest/developerguide/api_failures_messages.html
class TaskProtectionResult(StrEnum):
    ENABLED = "enabled"
    FAILED = "failed"
    DEPLOYMENT_BLOCKED = "DEPLOYMENT_BLOCKED"
    TASK_STOPPING_OR_STOPPED = "TASK_STOPPING_OR_STOPPED"


def enable(*, expires_in_minutes: int) -> TaskProtectionResult:
    agent_uri = os.environ["ECS_AGENT_URI"]
    payload = {"ProtectionEnabled": True, "ExpiresInMinutes": expires_in_minutes}

    response = _put(agent_uri + _ENDPOINT_PATH, payload)
    if response is None:
        return TaskProtectionResult.FAILED

    if "failure" not in response:
        return TaskProtectionResult.ENABLED

    match response["failure"]:
        case {"Reason": TaskProtectionResult.DEPLOYMENT_BLOCKED}:
            return TaskProtectionResult.DEPLOYMENT_BLOCKED
        case {"Reason": TaskProtectionResult.TASK_STOPPING_OR_STOPPED}:
            return TaskProtectionResult.TASK_STOPPING_OR_STOPPED
        case failure:
            _logger.error("task protection enable rejected: %s", failure)
            return TaskProtectionResult.FAILED


def disable() -> bool:
    response = _put(os.environ["ECS_AGENT_URI"] + _ENDPOINT_PATH, {"ProtectionEnabled": False})
    if response is None:
        return False
    if "failure" in response:
        _logger.error("task protection disable rejected: %s", response["failure"])
        return False
    return True


def _put(url: str, payload: Mapping[str, object]) -> dict[str, object] | None:
    request = Request(
        url,
        data=json.dumps(payload).encode(),
        headers={"Content-Type": "application/json"},
        method="PUT",
    )

    try:
        with urlopen(request, timeout=_REQUEST_TIMEOUT_SECONDS) as response:
            return cast("dict[str, object]", json.loads(response.read()))
    except OSError, http.client.HTTPException, json.JSONDecodeError:
        _logger.exception("task protection agent call failed at %s", url)
        return None