#!/usr/bin/env python3
"""Apply Terraform-owned task env onto an ECS task definition JSON.

Reads describe-task-definition JSON on stdin. Writes register-task-definition
input on stdout. IMAGE, GIT_SHA, CONTAINER, and TASK_ENV_JSON must be set.
TASK_ENV_JSON is the SecureString at /seahaven-ap/deploy/task-environment.
GIT_SHA is owned by GitHub and always overwrites the map.
"""

from __future__ import annotations

import json
import os
import sys


def patch_task_definition(td: dict, *, image: str, sha: str, container_name: str, env_map: dict) -> dict:
    if not isinstance(env_map, dict) or not env_map:
        raise SystemExit("TASK_ENV_JSON must be a non-empty JSON object")
    owned = {str(key): str(value) for key, value in env_map.items()}
    owned.pop("GIT_SHA", None)
    owned["GIT_SHA"] = sha

    matched = False
    for container in td.get("containerDefinitions") or []:
        if container.get("name") != container_name:
            continue
        matched = True
        container["image"] = image
        container["environment"] = [{"name": key, "value": value} for key, value in owned.items()]
        container["stopTimeout"] = 60
        container.pop("command", None)
    if not matched:
        raise SystemExit(f"container {container_name!r} not found in task definition")
    return td


def main() -> None:
    image = os.environ["IMAGE"]
    sha = os.environ["GIT_SHA"]
    name = os.environ["CONTAINER"]
    env_map = json.loads(os.environ["TASK_ENV_JSON"])
    td = json.load(sys.stdin)
    for key in (
        "taskDefinitionArn",
        "revision",
        "status",
        "requiresAttributes",
        "compatibilities",
        "registeredAt",
        "registeredBy",
        "deregisteredAt",
    ):
        td.pop(key, None)
    json.dump(patch_task_definition(td, image=image, sha=sha, container_name=name, env_map=env_map), sys.stdout)


if __name__ == "__main__":
    main()