"""The playground's front door while no task is running, and the switch that
routes to it.

The load balancer has one listener rule for all paths, and this function
owns where that rule points: at the app's target group while the app has a
healthy task, at this function while it has none. The load balancer itself
cannot make that choice; a forward to an empty target group is a 503, so the
switch has to be thrown from outside, and everything that could change the
answer invokes this function.

Only a person wakes the service. A public hostname is never quiet: internet
scanners reach it every few minutes, and a task started for each of them
would run around the clock. So while no task is running this function is
the sign-in: it sends the browser to WorkOS, finishes the sign-in WorkOS
sends back, checks the organization against the same Parameter Store values
the app uses, and sets the same session cookie (``session.py``, packaged
beside this file). A completed sign-in, or a request that already carries a
valid session, starts a task and gets the startup page. Everything else gets
the sign-in page, a redirect to it, a refusal, or 401 for an API call, and
starts nothing.

- The startup page polls the health route and reloads on the first 200.
  Each poll reaching this function re-checks target health and throws the
  switch to the app once a task is healthy, but answers 503 regardless: a
  poll still arriving here means the load balancer has not moved yet, and
  the 200 that reloads the page comes from the app itself. The health route
  never wakes anything, and this function never redirects a request to
  itself.
- A sign-in here is recorded as one signed-in request on the idle metric,
  so the alarm and the scheduled check both see the person before the app
  is up to count them.
- The idle alarm's state change to ALARM (no signed-in request for the idle
  period) scales the service to zero and throws the switch back here.
- A scheduled check every few minutes does the same whenever the alarm is
  in ALARM while a task runs, unless there was a signed-in request in the
  last few minutes or the task is younger than the wake grace: a task woken
  while the alarm already sits in ALARM never causes a transition on its
  own, and a task just started for a person has not served them yet.
- A task of the service stopping, for any reason, throws the switch back
  here when no healthy task remains, so the next sign-in wakes it again
  instead of meeting a 503.

Every step is idempotent, so a burst of requests from one impatient browser
costs nothing extra, and a missed event is repaired by the next request.
"""

from __future__ import annotations

import json
import logging
import os
from datetime import datetime, timedelta, timezone
from pathlib import Path
from typing import Any, Optional
from urllib.parse import unquote, urlencode

import boto3

# Packaged beside this file from src/quber/playground/session.py; the
# developer host has it only inside the package.
import session  # pyright: ignore[reportMissingImports]

log = logging.getLogger()
log.setLevel(logging.INFO)

CLUSTER = os.environ["ECS_CLUSTER"]
SERVICE = os.environ["ECS_SERVICE"]
IDLE_ALARM = os.environ["IDLE_ALARM"]
ACTIVITY_NAMESPACE = os.environ["ACTIVITY_NAMESPACE"]
ACTIVITY_METRIC = "AuthenticatedRequests"
#: How far back the scheduled check looks for signed-in requests before it
#: trusts the alarm. The alarm lags: it evaluates five-minute periods, and a
#: person who signed in a moment ago has not moved it yet.
RECENT_MINUTES = 10
#: A task younger than this is left alone by the scheduled check. Only a
#: sign-in starts one, and the person who signed in has not reached the app
#: yet while it boots.
WAKE_GRACE_MINUTES = 5
RULE_ARN = os.environ["RULE_ARN"]
APP_TARGET_GROUP = os.environ["APP_TARGET_GROUP_ARN"]
WAKE_TARGET_GROUP = os.environ["WAKE_TARGET_GROUP_ARN"]
SIGN_IN_PARAMETERS = (
    os.environ["WORKOS_API_KEY_PARAMETER"],
    os.environ["WORKOS_CLIENT_ID_PARAMETER"],
    os.environ["WORKOS_ORGANIZATION_ID_PARAMETER"],
    os.environ["SESSION_SECRET_PARAMETER"],
)

ecs = boto3.client("ecs")
elb = boto3.client("elbv2")
cloudwatch = boto3.client("cloudwatch")
ssm = boto3.client("ssm")

# The startup page, packaged beside this file (deploy/lambda/startup.html):
# the mark and fonts it needs are inlined or fetched from Google Fonts, because
# nothing under the app's /static/ answers while the app is down.
PAGE = Path(__file__).with_name("startup.html").read_text(encoding="utf-8")
# The link-preview card, packaged as base64 text because the function's
# archive is built from text sources; the load balancer decodes it on the way
# out. Served without a session so a link scraper gets the image while the
# service is at zero, as the app serves it from its static files once up.
LINK_CARD = Path(__file__).with_name("link-card.png.b64").read_text(encoding="ascii").strip()
# The Qubera mark for the sign-in pages, inlined for the same reason.
MARK_SRC = "data:image/png;base64," + (
    Path(__file__).with_name("qubera-mark-160.png.b64").read_text(encoding="ascii").strip()
)

_config: Optional[session.WorkOS] = None


def config() -> session.WorkOS:
    """The WorkOS key, client ID, organization and session secret, read once
    per container."""
    global _config
    if _config is None:
        found = ssm.get_parameters(Names=list(SIGN_IN_PARAMETERS), WithDecryption=True)["Parameters"]
        by_name = {p["Name"]: p["Value"] for p in found}
        key, client, organization, secret = (by_name[name] for name in SIGN_IN_PARAMETERS)
        _config = session.WorkOS(
            api_key=key, client_id=client, organization_id=organization, session_secret=secret
        )
    return _config


def app_healthy() -> bool:
    health = elb.describe_target_health(TargetGroupArn=APP_TARGET_GROUP)["TargetHealthDescriptions"]
    return any(t["TargetHealth"]["State"] == "healthy" for t in health)


def route_to(target_group: str) -> None:
    """Point the listener rule at a target group, if it is not there already."""
    rule = elb.describe_rules(RuleArns=[RULE_ARN])["Rules"][0]
    current = rule["Actions"][0].get("TargetGroupArn") if rule["Actions"] else None
    if current == target_group:
        return
    elb.modify_rule(RuleArn=RULE_ARN, Actions=[{"Type": "forward", "TargetGroupArn": target_group}])
    log.info("listener rule now forwards to %s", target_group.rsplit("/", 2)[-2])


def desired_count() -> int:
    return ecs.describe_services(cluster=CLUSTER, services=[SERVICE])["services"][0]["desiredCount"]


def set_desired(count: int) -> None:
    if desired_count() == count:
        return
    ecs.update_service(cluster=CLUSTER, service=SERVICE, desiredCount=count)
    log.info("%s desired count -> %d", SERVICE, count)


def youngest_task_minutes() -> Optional[float]:
    """Age of the service's newest task, or None when it has none."""
    arns = ecs.list_tasks(cluster=CLUSTER, serviceName=SERVICE, desiredStatus="RUNNING")["taskArns"]
    if not arns:
        return None
    tasks = ecs.describe_tasks(cluster=CLUSTER, tasks=arns)["tasks"]
    now = datetime.now(timezone.utc)
    return min((now - t["createdAt"]).total_seconds() for t in tasks) / 60


def record_sign_in() -> None:
    """One signed-in request on the idle metric, so the alarm and the scheduled
    check see the person before the app is up to count them."""
    cloudwatch.put_metric_data(
        Namespace=ACTIVITY_NAMESPACE,
        MetricData=[
            {
                "MetricName": ACTIVITY_METRIC,
                "Dimensions": [{"Name": "Service", "Value": SERVICE}],
                "Value": 1,
                "Unit": "Count",
            }
        ],
    )


def wake() -> None:
    set_desired(1)
    route_to(WAKE_TARGET_GROUP)
    record_sign_in()


def request_cookies(headers: dict[str, str]) -> dict[str, str]:
    cookies: dict[str, str] = {}
    for part in headers.get("cookie", "").split(";"):
        name, _, value = part.strip().partition("=")
        if name:
            cookies[name] = value
    return cookies


def on_request(event: dict[str, Any]) -> dict[str, Any]:
    path = event.get("path", "/")
    # The load balancer hands query values over still URL-encoded.
    query = {unquote(k): unquote(v) for k, v in (event.get("queryStringParameters") or {}).items()}
    headers = {k.lower(): v for k, v in (event.get("headers") or {}).items()}
    healthy = app_healthy()
    if healthy:
        route_to(APP_TARGET_GROUP)
    if path == "/healthz":
        # Never 200 from here. The startup page reloads on the first 200, and
        # a request reaching this function means the load balancer has not
        # yet moved the switch; the app answers this route itself once it has.
        return respond(
            503, "application/json", json.dumps({"status": "switching" if healthy else "starting"})
        )
    if path == session.LINK_CARD_PATH:
        return {
            "statusCode": 200,
            "headers": {"Content-Type": "image/png", "Cache-Control": "public, max-age=86400"},
            "body": LINK_CARD,
            "isBase64Encoded": True,
        }
    sign_in = config()
    secure = headers.get("x-forwarded-proto") == "https"
    base = session.base_url(headers.get("x-forwarded-proto", "https"), headers.get("host", ""))
    cookie = request_cookies(headers).get(session.COOKIE)
    if path == "/login":
        return respond(
            200, "text/html; charset=utf-8", session.login_page(sign_in, base, query.get("next"), MARK_SRC)
        )
    if path == "/callback":
        outcome = session.complete_sign_in(sign_in, base, query.get("code"), query.get("state"), MARK_SRC)
        log.info("sign-in %s", "accepted" if outcome.token else f"not accepted ({outcome.status})")
        if outcome.token is None or outcome.location is None:
            return respond(outcome.status, "text/html; charset=utf-8", outcome.page or "")
        if healthy:
            record_sign_in()
        else:
            wake()
        return {
            "statusCode": 303,
            "headers": {
                "Location": outcome.location,
                "Set-Cookie": session.cookie_header(outcome.token, secure=secure),
                "Cache-Control": "no-store",
            },
            "body": "",
        }
    if path == "/logout":
        identity = session.signed_identity(sign_in.session_secret, cookie)
        location = session.logout_url(identity.sid, base) if identity and identity.sid else "/"
        return {
            "statusCode": 302,
            "headers": {
                "Location": location,
                "Set-Cookie": session.clear_cookie_header(secure=secure),
                "Cache-Control": "no-store",
            },
            "body": "",
        }
    if session.verify(sign_in.session_secret, sign_in.organization_id, cookie):
        if not healthy:
            wake()
        return respond(503, "text/html; charset=utf-8", PAGE)
    if path.startswith("/api/"):
        return respond(401, "application/json", json.dumps({"detail": "login required"}))
    target = path + ("?" + urlencode(query) if query else "")
    location = "/login?" + urlencode({"next": target})
    return {"statusCode": 302, "headers": {"Location": location, "Cache-Control": "no-store"}, "body": ""}


def on_alarm(event: dict[str, Any]) -> None:
    state = event["detail"]["state"]["value"]
    log.info("idle alarm state %s", state)
    if state == "ALARM":
        set_desired(0)
        route_to(WAKE_TARGET_GROUP)


def recent_activity() -> float:
    now = datetime.now(timezone.utc)
    stats = cloudwatch.get_metric_statistics(
        Namespace=ACTIVITY_NAMESPACE,
        MetricName=ACTIVITY_METRIC,
        Dimensions=[{"Name": "Service", "Value": SERVICE}],
        StartTime=now - timedelta(minutes=RECENT_MINUTES),
        EndTime=now,
        Period=60,
        Statistics=["Sum"],
    )
    return sum(point["Sum"] for point in stats["Datapoints"])


def on_schedule() -> None:
    state = cloudwatch.describe_alarms(AlarmNames=[IDLE_ALARM])["MetricAlarms"][0]["StateValue"]
    desired = desired_count()
    if state != "ALARM" or desired == 0:
        log.info("scheduled check: idle alarm %s, desired count %d, nothing to do", state, desired)
        return
    recent = recent_activity()
    age = youngest_task_minutes()
    log.info(
        "scheduled check: idle alarm ALARM, desired count %d, signed-in requests in the last %d min: %s, "
        "youngest task %s min old",
        desired,
        RECENT_MINUTES,
        recent,
        "none" if age is None else round(age, 1),
    )
    if recent > 0 or (age is not None and age < WAKE_GRACE_MINUTES):
        return
    set_desired(0)
    route_to(WAKE_TARGET_GROUP)


def on_task_stopped(event: dict[str, Any]) -> None:
    log.info("task stopped: %s", event["detail"].get("stoppedReason", "no reason given"))
    if not app_healthy():
        route_to(WAKE_TARGET_GROUP)


def respond(status: int, content_type: str, body: str) -> dict[str, Any]:
    return {
        "statusCode": status,
        "headers": {"Content-Type": content_type, "Cache-Control": "no-store"},
        "body": body,
    }


def handler(event: dict[str, Any], context: Any) -> Any:
    if "requestContext" in event and "elb" in event["requestContext"]:
        return on_request(event)
    if event.get("source") == "aws.ecs":
        return on_task_stopped(event)
    if event.get("source") == "aws.cloudwatch":
        return on_alarm(event)
    if event.get("source") == "aws.events" and event.get("detail-type") == "Scheduled Event":
        return on_schedule()
    log.warning("unrecognised event: %s", json.dumps(event)[:500])
    return None