aion-package 0.9.1

//! Emission of the generated Python remote-worker module (tier = `RemotePython`).
//!
//! Generates `worker/worker.py`, the do-not-edit plumbing that serves a
//! package's `RemotePython` activities: it decodes each pushed task's JSON
//! input and routes it to the matching author-written handler in the
//! hand-written `handlers` module, exactly as the generated Gleam wrapper
//! references `activities.<name>` and the generated Rust worker references
//! `handlers::<name>`. The side-effecting bodies live in `worker/handlers.py`
//! (each a `async def <name>(request) -> DispatchOutcome`), so this file
//! carries no activity logic and is a pure, deterministic function of the
//! declarations — `aion generate --check` byte-compares it and a hand-edit is
//! a build failure (checklist C4).
//!
//! The handler registry follows declaration order (no map iteration), so
//! generation is deterministic. No activity retry/timeout/backoff is emitted
//! (ADR-001); the only values present are the worker's own connection
//! parameters. Every required value is read from the environment so the
//! do-not-edit file stays byte-stable across hosts, and a missing variable
//! fails loud (no invented default is supplied anywhere; ADR-001).

use std::fmt::Write as _;

use super::activity_model::ResolvedActivity;

/// Python do-not-edit banner (`#`-comment form; the Gleam `////` header is not
/// valid Python). `aion generate --check` byte-compares the whole file, so the
/// banner is the contract, not a parsed marker.
const PYTHON_HEADER: &str =
    "# Generated by aion generate — do not edit; regenerate from the activity declarations.";

/// Imports and module constants, fixed for every generated worker. The
/// `handlers` module is hand-written by the author and holds the activity
/// bodies.
const PREAMBLE: &str = r#"from __future__ import annotations

import asyncio
import json
import logging
import os
from collections.abc import Awaitable, Callable, Iterable

from aion_worker import (
    ActivityExecutionContext,
    ActivityTask,
    DispatchOutcome,
    Failed,
    GrpcWorkerSession,
    ReconnectConfig,
    WorkerConfig,
    connect_register_replay_and_serve,
)
from aion_worker.proto import common_pb2, worker_pb2

import handlers

JSON_CONTENT_TYPE = "application/json"
Handler = Callable[[dict[str, object]], Awaitable[DispatchOutcome]]
"#;

/// The dispatcher's `activity_types` + `dispatch` methods, fixed. `dispatch`
/// decodes the task input and routes to the registered handler, which returns
/// the `DispatchOutcome` itself.
const DISPATCH_BODY: &str = r#"
    def activity_types(self) -> Iterable[str]:
        return self._handlers.keys()

    async def dispatch(
        self, task: ActivityTask, context: ActivityExecutionContext
    ) -> DispatchOutcome:
        del context
        handler = self._handlers.get(task.activity_type)
        if handler is None:
            return worker_failure(f"unknown activity type: {task.activity_type}")
        try:
            request = decode_json_object(task.input)
            return await handler(request)
        except (KeyError, ValueError, json.JSONDecodeError, UnicodeDecodeError) as exc:
            return worker_failure(str(exc))
"#;

/// The module-level JSON-decode and failure helpers the plumbing uses to route
/// tasks. Result encoding (`Completed`/`json_payload`) lives in the author's
/// self-contained `handlers` module, not here.
const HELPERS: &str = r#"
def decode_json_object(payload: common_pb2.Payload) -> dict[str, object]:
    if payload.content_type != JSON_CONTENT_TYPE:
        raise ValueError(f"expected {JSON_CONTENT_TYPE} payload, got {payload.content_type!r}")
    value = json.loads(payload.bytes.decode("utf-8"))
    if not isinstance(value, dict):
        raise ValueError("expected JSON object input")
    return value


def worker_failure(message: str) -> DispatchOutcome:
    return Failed(
        worker_pb2.ActivityError(
            kind=worker_pb2.ACTIVITY_ERROR_KIND_TERMINAL,
            message=message,
        )
    )
"#;

/// Fail-loud environment-variable readers used by `worker_config`. Each
/// required value is read from the environment with no invented default
/// (ADR-001); a missing or malformed variable raises `SystemExit` with a clear
/// message instead of silently substituting a value.
const CONFIG_HELPERS: &str = r#"
def _require_env(name: str) -> str:
    value = os.environ.get(name)
    if value is None:
        raise SystemExit(f"required environment variable {name} is not set")
    return value


def _require_int(name: str) -> int:
    raw = _require_env(name)
    try:
        return int(raw)
    except ValueError:
        raise SystemExit(f"environment variable {name} is not a valid integer: {raw!r}") from None


def _require_float(name: str) -> float:
    raw = _require_env(name)
    try:
        return float(raw)
    except ValueError:
        raise SystemExit(f"environment variable {name} is not a valid number: {raw!r}") from None
"#;

/// The worker's connection config, fixed for every generated worker. Every
/// required field is read from the environment via the fail-loud helpers;
/// `namespace`/`subject` are omitted so they defer to the dataclass defaults
/// the library owns.
const WORKER_CONFIG: &str = r#"

def worker_config() -> WorkerConfig:
    return WorkerConfig(
        endpoint=_require_env("AION_WORKER_ENDPOINT"),
        task_queue=_require_env("AION_TASK_QUEUE"),
        identity=_require_env("AION_WORKER_IDENTITY"),
        max_concurrency=_require_int("AION_WORKER_CONCURRENCY"),
        reconnect=ReconnectConfig(
            initial_backoff_seconds=_require_float("AION_RECONNECT_INITIAL_BACKOFF_SECONDS"),
            max_backoff_seconds=_require_float("AION_RECONNECT_MAX_BACKOFF_SECONDS"),
            max_attempts=_require_int("AION_RECONNECT_MAX_ATTEMPTS"),
        ),
    )
"#;

/// Emits `worker/worker.py` serving `activities` (all `RemotePython`), routing
/// each engine activity name to `handlers.<name>` in declaration order.
pub(crate) fn emit(package_name: &str, activities: &[&ResolvedActivity]) -> String {
    let class_name = format!("{}Dispatcher", to_pascal(package_name));

    let mut out = String::new();
    out.push_str(PYTHON_HEADER);
    let _ = write!(
        out,
        "\n\"\"\"{package_name} Aion worker.\n\nServes the package's RemotePython activities against a running Aion server.\nThe activity bodies live in `handlers.py`; this module is generated plumbing.\n\"\"\"\n\n"
    );
    out.push_str(PREAMBLE);

    let _ = write!(out, "\n\nclass {class_name}:\n");
    let _ = writeln!(
        out,
        "    \"\"\"Routes the {package_name} activities to their handlers.\"\"\""
    );
    out.push_str(
        "\n    def __init__(self) -> None:\n        self._handlers: dict[str, Handler] = {\n",
    );
    for activity in activities {
        let _ = writeln!(
            out,
            "            \"{0}\": handlers.{0},",
            activity.declaration.name
        );
    }
    out.push_str("        }\n");
    out.push_str(DISPATCH_BODY);

    out.push('\n');
    out.push_str(HELPERS);

    out.push('\n');
    out.push_str(CONFIG_HELPERS);
    out.push_str(WORKER_CONFIG);

    let _ = write!(
        out,
        "\n\nasync def main() -> None:\n    logging.basicConfig(level=logging.INFO)\n    \
         config = worker_config()\n    dispatcher = {class_name}()\n    \
         await connect_register_replay_and_serve(\n        config=config,\n        \
         connect=lambda: GrpcWorkerSession.connect(config),\n        \
         dispatcher=dispatcher,\n    )\n\n\nif __name__ == \"__main__\":\n    asyncio.run(main())\n"
    );

    out
}

/// `aion_order_saga` → `AionOrderSaga`.
fn to_pascal(snake: &str) -> String {
    let mut out = String::with_capacity(snake.len());
    for segment in snake.split('_').filter(|segment| !segment.is_empty()) {
        let mut chars = segment.chars();
        if let Some(first) = chars.next() {
            out.extend(first.to_uppercase());
            out.push_str(chars.as_str());
        }
    }
    out
}

#[cfg(test)]
mod tests {
    use std::path::PathBuf;

    use super::{emit, to_pascal};
    use crate::codegen::activity_model::{ResolvedActivity, ResolvedType};
    use crate::codegen::declaration::{ActivityDeclaration, Tier};
    use crate::codegen::model::{BoundaryType, GleamType};

    fn artifact() -> BoundaryType {
        BoundaryType {
            file: PathBuf::from("schemas/placeholder.json"),
            stem: "placeholder".to_owned(),
            root: GleamType::Named {
                type_name: "Placeholder".to_owned(),
                fn_prefix: "placeholder".to_owned(),
            },
            defs: Vec::new(),
        }
    }

    fn declaration(name: &str, input: &str, output: &str) -> ActivityDeclaration {
        ActivityDeclaration {
            name: name.to_owned(),
            tier: Tier::RemotePython,
            input_type: input.to_owned(),
            output_type: output.to_owned(),
        }
    }

    fn resolved<'a>(
        declaration: &'a ActivityDeclaration,
        art: &'a BoundaryType,
    ) -> ResolvedActivity<'a> {
        ResolvedActivity {
            declaration,
            input: ResolvedType {
                gleam_type: declaration.input_type.clone(),
                fn_prefix: declaration.input_type.to_lowercase(),
                boundary: art,
            },
            output: ResolvedType {
                gleam_type: declaration.output_type.clone(),
                fn_prefix: declaration.output_type.to_lowercase(),
                boundary: art,
            },
        }
    }

    #[test]
    fn pascal_case_joins_snake_segments() {
        assert_eq!(to_pascal("aion_order_saga"), "AionOrderSaga");
        assert_eq!(to_pascal("demo"), "Demo");
    }

    #[test]
    fn routes_each_activity_to_its_handler_in_order() {
        let art = artifact();
        let d1 = declaration("reserve_inventory", "OrderInput", "InventoryReservation");
        let d2 = declaration("charge_payment", "OrderInput", "PaymentReceipt");
        let activities = [resolved(&d1, &art), resolved(&d2, &art)];
        let refs: Vec<&ResolvedActivity> = activities.iter().collect();

        let module = emit("demo", &refs);

        assert!(module.starts_with(super::PYTHON_HEADER));
        assert!(module.contains("class DemoDispatcher:"));
        // Bodies live in the hand-written handlers module, not here.
        assert!(module.contains("import handlers\n"));
        assert!(module.contains("\"reserve_inventory\": handlers.reserve_inventory,"));
        assert!(module.contains("\"charge_payment\": handlers.charge_payment,"));
        // No inline activity body or stub is generated.
        assert!(!module.contains("NotImplementedError"));
        assert!(!module.contains("async def reserve_inventory"));
        // Declaration order preserved in the registry.
        let first = module.find("\"reserve_inventory\": handlers");
        let second = module.find("\"charge_payment\": handlers");
        assert!(first.is_some() && second.is_some() && first < second);
        // Fail-loud, ADR-001-pure config: every required value is read from the
        // environment through the require-helpers, with no invented default.
        assert!(module.contains("_require_env(\"AION_WORKER_ENDPOINT\")"));
        assert!(module.contains("_require_int(\"AION_WORKER_CONCURRENCY\")"));
        assert!(module.contains("_require_float(\"AION_RECONNECT_INITIAL_BACKOFF_SECONDS\")"));
        // No defaulted environment read survives anywhere (the violation form).
        assert!(!module.contains("os.environ.get(\""));
        // namespace/subject defer to the library's dataclass defaults.
        assert!(!module.contains("AION_NAMESPACE"));
        assert!(!module.contains("AION_SUBJECT"));
    }
}