client.py

import json
import re
import subprocess
from collections import defaultdict
from pathlib import Path
from textwrap import dedent
from typing import (
    Any,
    Callable,
    Literal,
    OrderedDict,
    Sequence,
    Set,
    TextIO,
    assert_never,
    cast,
)

from pydantic import BaseModel, Field, RootModel

from replit_river.codegen.format import reindent
from replit_river.codegen.typing import (
    ClassName,
    DictTypeExpr,
    FileContents,
    HandshakeType,
    ListTypeExpr,
    LiteralType,
    LiteralTypeExpr,
    ModuleName,
    NoneTypeExpr,
    OpenUnionTypeExpr,
    RenderedPath,
    TypeExpression,
    TypeName,
    UnionTypeExpr,
    extract_inner_type,
    normalize_special_chars,
    render_literal_type,
    render_type_expr,
)

_NON_ALNUM_RE = re.compile(r"[^a-zA-Z0-9_]+")

# Literal is here because HandshakeType can be Literal[None]
ROOT_FILE_HEADER = dedent(
    """\
# Code generated by river.codegen. DO NOT EDIT.
from pydantic import BaseModel
from typing import Literal

import replit_river as river

"""
)

SERVICE_FILE_HEADER = dedent(
    """\
# Code generated by river.codegen. DO NOT EDIT.
from collections.abc import AsyncIterable, AsyncIterator
from typing import Any
import datetime

from pydantic import TypeAdapter

from replit_river.error_schema import RiverError, RiverErrorTypeAdapter
import replit_river as river

"""
)

FILE_HEADER = dedent(
    """\
# Code generated by river.codegen. DO NOT EDIT.
from collections.abc import AsyncIterable, AsyncIterator
import datetime
from typing import (
    Any,
    Literal,
    Mapping,
    NotRequired,
    TypedDict,
)
from typing_extensions import Annotated

from pydantic import BaseModel, Field, TypeAdapter, WrapValidator
from replit_river.error_schema import RiverError
from replit_river.client import RiverUnknownError, translate_unknown_error, \
    RiverUnknownValue, translate_unknown_value

import replit_river as river

"""
)


class RiverConcreteType(BaseModel):
    type: str | None = Field(default=None)
    properties: dict[str, "RiverType"] = Field(default_factory=lambda: dict())
    required: Set[str] = Field(default=set())
    items: "RiverType | None" = Field(default=None)
    const: str | int | None = Field(default=None)
    patternProperties: dict[str, "RiverType"] = Field(default_factory=lambda: dict())
    id_: str | None = Field(default=None, alias="$id")
    ref: str | None = Field(default=None, alias="$ref")


class RiverUnionType(BaseModel):
    anyOf: list["RiverType"]


class RiverIntersectionType(BaseModel):
    allOf: list["RiverType"]


class RiverNotType(BaseModel):
    """This is used to represent void / never."""

    not_: Any = Field(..., alias="not")


RiverType = RiverConcreteType | RiverUnionType | RiverNotType | RiverIntersectionType


class RiverProcedure(BaseModel):
    init: RiverType | None = Field(default=None)
    input: RiverType | None = Field(default=None)
    output: RiverType
    errors: RiverType | None = Field(default=None)
    type: (
        Literal["rpc"] | Literal["stream"] | Literal["subscription"] | Literal["upload"]
    )
    description: str | None = Field(default=None)


class RiverService(BaseModel):
    procedures: dict[str, RiverProcedure]


class RiverSchema(BaseModel):
    services: dict[str, RiverService]
    handshakeSchema: RiverConcreteType | None = Field(default=None)


RiverSchemaFile = RootModel[RiverSchema]


def is_literal(tpe: RiverType) -> bool:
    if isinstance(tpe, RiverUnionType):
        return all(is_literal(t) for t in tpe.anyOf)
    elif isinstance(tpe, RiverConcreteType):
        return tpe.type not in set(["object", "array"])
    else:
        return False


def encode_type(
    type: RiverType,
    prefix: TypeName,
    base_model: str,
    in_module: list[ModuleName],
    permit_unknown_members: bool,
    type_registry: dict[str, TypeName] | None = None,
) -> tuple[TypeExpression, list[ModuleName], list[FileContents], set[TypeName]]:
    # Registry to track $id -> TypeName mappings for resolving $ref
    if type_registry is None:
        type_registry = {}

    def _make_open_union_type_expr(one_of: list[TypeExpression]) -> OpenUnionTypeExpr:
        if base_model == "RiverError":
            return OpenUnionTypeExpr(
                UnionTypeExpr(one_of),
                fallback_type="RiverUnknownError",
                validator_function="translate_unknown_error",
            )
        else:
            return OpenUnionTypeExpr(
                UnionTypeExpr(one_of),
                fallback_type="RiverUnknownValue",
                validator_function="translate_unknown_value",
            )

    encoder_name: TypeName | None = None  # defining this up here to placate mypy
    chunks: list[FileContents] = []

    # Handle $ref - return a forward reference to the registered type
    if isinstance(type, RiverConcreteType) and type.ref is not None:
        ref_id = type.ref
        if ref_id in type_registry:
            # Use forward reference string for the type
            return (TypeName(f'"{type_registry[ref_id].value}"'), [], [], set())
        else:
            # Unknown ref, fall back to Any
            return (TypeName("Any"), [], [], set())

    if isinstance(type, RiverNotType):
        return (NoneTypeExpr(), [], [], set())
    elif isinstance(type, RiverUnionType):
        typeddict_encoder = list[str]()
        encoder_names: set[TypeName] = set()

        # First check if it's a discriminated union. Typebox currently doesn't have
        # a way of expressing the intention of having a discriminated union. So we
        # do a bit of detection if that is structurally true by checking that all the
        # types in the anyOf are objects, have properties, and have one property common
        # to all the alternatives that has a literal value.
        def flatten_union(tpe: RiverType) -> list[RiverType]:
            if isinstance(tpe, RiverUnionType):
                return [u for t in tpe.anyOf for u in flatten_union(t)]
            else:
                return [tpe]

        original_type = type

        type = RiverUnionType(anyOf=flatten_union(type))

        one_of_candidate_types: list[RiverConcreteType] = [
            t
            for _t in type.anyOf
            for t in (_t.anyOf if isinstance(_t, RiverUnionType) else [_t])
            if isinstance(t, RiverConcreteType)
            and t.type == "object"
            and t.properties
            and (not t.patternProperties or "^(.*)$" not in t.patternProperties)
        ]
        if len(type.anyOf) > 0 and len(type.anyOf) == len(one_of_candidate_types):
            # We have established that it is a union-of-objects. Now let's see if
            # there is a discriminator field common among all options.
            literal_fields = set[str]()
            for i, oneof_t in enumerate(one_of_candidate_types):
                lf = set[str](
                    name
                    for name, prop in oneof_t.properties.items()
                    if isinstance(prop, RiverConcreteType)
                    and prop.type in ("string", "number", "boolean")
                    and prop.const is not None
                )
                if i == 0:
                    literal_fields = lf
                else:
                    literal_fields.intersection_update(lf)
                if not literal_fields:
                    # There are no more candidates.
                    break
            if len(literal_fields) == 1:
                # Hooray! we found a discriminated union.
                discriminator_name = literal_fields.pop()
                one_of_pending = OrderedDict[str, tuple[str, list[RiverConcreteType]]]()

                for oneof_t in one_of_candidate_types:
                    discriminator_value = [
                        _NON_ALNUM_RE.sub("", str(prop.const))
                        for name, prop in sorted(
                            list(oneof_t.properties.items()), key=lambda kv: kv[0]
                        )
                        if isinstance(prop, RiverConcreteType)
                        and name == discriminator_name
                        and prop.const is not None
                    ].pop()
                    one_of_pending.setdefault(
                        f"{render_literal_type(prefix)}OneOf_{discriminator_value}",
                        (discriminator_value, []),
                    )[1].append(oneof_t)

                one_of: list[TypeExpression] = []
                if discriminator_name == "$kind":
                    discriminator_name = "kind"
                for pfx, (discriminator_value, oneof_ts) in one_of_pending.items():
                    if len(oneof_ts) > 1:
                        typeddict_encoder.append("(")
                        # Tricky bit. We need to derive our own discriminator based
                        # on known members. Be careful.

                        common_members = set[str]()
                        for i, oneof_t in enumerate(oneof_ts):
                            if i == 0:
                                common_members = set(oneof_t.properties.keys())
                            else:
                                common_members.intersection_update(
                                    oneof_t.properties.keys()
                                )

                        for i, oneof_t in enumerate(oneof_ts):
                            type_name, _, contents, _ = encode_type(
                                oneof_ts[i],
                                TypeName(f"{pfx}{i}"),
                                base_model,
                                in_module,
                                permit_unknown_members=permit_unknown_members,
                                type_registry=type_registry,
                            )
                            one_of.append(type_name)
                            chunks.extend(contents)
                            local_discriminators = set(
                                oneof_t.properties.keys()
                            ).difference(common_members)
                            encoder_name = TypeName(
                                f"encode_{render_literal_type(type_name)}"
                            )
                            if base_model == "TypedDict":
                                # "encoder_names" is only a TypedDict thing
                                encoder_names.add(encoder_name)
                            _field_name = render_literal_type(encoder_name)
                            typeddict_encoder.append(
                                f"""\
                                {_field_name}(x) # type: ignore[arg-type]
                                """.strip()
                            )
                            if local_discriminators:
                                local_discriminator = sorted(local_discriminators).pop()
                            else:
                                local_discriminator = "FIXME: Ambiguous discriminators"
                            typeddict_encoder.append(
                                f" if {repr(local_discriminator)} in x else "
                            )
                        typeddict_encoder.pop()  # Drop the last ternary
                        typeddict_encoder.append(")")
                    else:
                        oneof_t = oneof_ts[0]
                        type_name, _, contents, _ = encode_type(
                            oneof_t,
                            TypeName(pfx),
                            base_model,
                            in_module,
                            permit_unknown_members=permit_unknown_members,
                            type_registry=type_registry,
                        )
                        one_of.append(type_name)
                        chunks.extend(contents)
                        encoder_name = TypeName(
                            f"encode_{render_literal_type(type_name)}"
                        )
                        # TODO(dstewart): Figure out why uncommenting this breaks
                        #                 generated code
                        # encoder_names.add(encoder_name)
                        typeddict_encoder.append(
                            f"{render_literal_type(encoder_name)}(x)"
                        )
                    typeddict_encoder.append(
                        f"""
                            if x[{repr(discriminator_name)}]
                            == {repr(discriminator_value)}
                            else
                        """,
                    )
                union: TypeExpression
                if permit_unknown_members:
                    union = _make_open_union_type_expr(one_of)
                else:
                    union = UnionTypeExpr(one_of)
                chunks.append(
                    FileContents(
                        f"{render_literal_type(prefix)} = {render_type_expr(union)}"
                    )
                )
                chunks.append(FileContents(""))

                if base_model == "TypedDict":
                    encoder_name = TypeName(f"encode_{render_literal_type(prefix)}")
                    encoder_names.add(encoder_name)
                    _field_name = render_literal_type(encoder_name)
                    typeddict_encoder = typeddict_encoder[:-1]  # Drop the last ternary
                    chunks.append(
                        FileContents(
                            "\n".join(
                                [
                                    dedent(
                                        f"""\
                            def {_field_name}(
                                x: {repr(render_literal_type(prefix))},
                            ) -> Any:
                                return (
                                    {
                                            reindent(
                                                "                                    ",
                                                "\n".join(typeddict_encoder),
                                            )
                                        }
                                )
                                        """.rstrip()
                                    )
                                ]
                            )
                        )
                    )
                return (prefix, in_module, chunks, encoder_names)
            # End of stable union detection
        # Restore the non-flattened union type
        type = original_type
        any_of: list[TypeExpression] = []

        # Collect (type_check, encoder_expr) pairs for building ternary chain
        encoder_parts: list[tuple[str | None, str]] = []
        for i, t in enumerate(type.anyOf):
            type_name, _, contents, _ = encode_type(
                t,
                TypeName(f"{render_literal_type(prefix)}AnyOf_{i}"),
                base_model,
                in_module,
                permit_unknown_members=permit_unknown_members,
                type_registry=type_registry,
            )
            any_of.append(type_name)
            chunks.extend(contents)
            if isinstance(t, RiverConcreteType):
                if t.type == "string":
                    encoder_parts.append(("isinstance(x, str)", "x"))
                elif t.type == "array":
                    match type_name:
                        case ListTypeExpr(inner_type_name):
                            # Primitives don't need encoding
                            inner_type_str = render_literal_type(inner_type_name)
                            if inner_type_str in ("str", "int", "float", "bool", "Any"):
                                encoder_parts.append(("isinstance(x, list)", "list(x)"))
                            else:
                                encoder_parts.append(
                                    (
                                        "isinstance(x, list)",
                                        f"[encode_{inner_type_str}(y) for y in x]",
                                    )
                                )
                        case _:
                            encoder_parts.append(("isinstance(x, list)", "list(x)"))
                elif t.type == "object":
                    match type_name:
                        case TypeName(value):
                            encoder_parts.append(
                                ("isinstance(x, dict)", f"encode_{value}(x)")
                            )
                        case _:
                            encoder_parts.append(("isinstance(x, dict)", "dict(x)"))
                elif t.type in ("number", "integer"):
                    match type_name:
                        case LiteralTypeExpr(const):
                            encoder_parts.append((f"x == {repr(const)}", repr(const)))
                        case _:
                            encoder_parts.append(("isinstance(x, (int, float))", "x"))
                elif t.type == "boolean":
                    encoder_parts.append(("isinstance(x, bool)", "x"))
                elif t.type == "null" or t.type == "undefined":
                    encoder_parts.append(("x is None", "None"))
                else:
                    # Fallback for other types
                    match type_name:
                        case TypeName(value):
                            encoder_parts.append((None, f"encode_{value}(x)"))
                        case LiteralTypeExpr(const):
                            encoder_parts.append((None, repr(const)))
                        case NoneTypeExpr():
                            encoder_parts.append((None, "None"))
                        case _:
                            encoder_parts.append((None, "x"))

        # Build the ternary chain from encoder_parts
        typeddict_encoder = list[str]()
        for i, (type_check, encoder_expr) in enumerate(encoder_parts):
            is_last = i == len(encoder_parts) - 1
            if is_last or type_check is None:
                # Last item or no type check - just the expression
                typeddict_encoder.append(encoder_expr)
            else:
                # Add expression with type check
                typeddict_encoder.append(f"{encoder_expr} if {type_check} else")
        if permit_unknown_members:
            union = _make_open_union_type_expr(any_of)
        else:
            union = UnionTypeExpr(any_of)
        if is_literal(type):
            typeddict_encoder = ["x"]
        chunks.append(
            FileContents(f"{render_literal_type(prefix)} = {render_type_expr(union)}")
        )
        if base_model == "TypedDict":
            encoder_name = TypeName(f"encode_{render_literal_type(prefix)}")
            encoder_names.add(encoder_name)
            _field_name = render_literal_type(encoder_name)
            chunks.append(
                FileContents(
                    dedent(
                        f"""
                    def {_field_name}(x: {repr(render_literal_type(prefix))}) -> Any:
                        return (
                            {
                            reindent(
                                "                            ",
                                "\n".join(typeddict_encoder),
                            )
                        }
                        )
                        """
                    )
                )
            )
        return (prefix, in_module, chunks, encoder_names)
    elif isinstance(type, RiverIntersectionType):

        def extract_props(tpe: RiverType) -> list[dict[str, RiverType]]:
            if isinstance(tpe, RiverUnionType):
                return [t for p in tpe.anyOf for t in extract_props(p)]
            elif isinstance(tpe, RiverConcreteType):
                return [tpe.properties]
            elif isinstance(tpe, RiverIntersectionType):
                return [t for p in tpe.allOf for t in extract_props(p)]
            else:
                return []

        combined = {}
        for x in extract_props(type):
            combined.update(x)
        return encode_type(
            RiverConcreteType(type="object", properties=combined),
            prefix,
            base_model,
            in_module,
            permit_unknown_members=permit_unknown_members,
            type_registry=type_registry,
        )
    elif isinstance(type, RiverConcreteType):
        typeddict_encoder = list[str]()
        if type.type is None:
            # Handle the case where type is not specified
            typeddict_encoder.append("x")
            return (TypeName("Any"), [], [], set())
        elif type.type == "string":
            if type.const:
                typeddict_encoder.append(repr(type.const))
                return (LiteralTypeExpr(type.const), [], [], set())
            else:
                typeddict_encoder.append("x")
                return (TypeName("str"), [], [], set())
        elif type.type == "Uint8Array":
            typeddict_encoder.append("x.decode()")
            return (TypeName("bytes"), [], [], set())
        elif type.type == "number":
            if type.const is not None:
                # enums are represented as const number in the schema
                typeddict_encoder.append(f"{type.const}")
                return (LiteralTypeExpr(type.const), [], [], set())
            typeddict_encoder.append("x")
            return (TypeName("float"), [], [], set())
        elif type.type == "integer":
            if type.const is not None:
                # enums are represented as const number in the schema
                typeddict_encoder.append(f"{type.const}")
                return (LiteralTypeExpr(type.const), [], [], set())
            typeddict_encoder.append("x")
            return (TypeName("int"), [], [], set())
        elif type.type == "boolean":
            typeddict_encoder.append("x")
            return (TypeName("bool"), [], [], set())
        elif type.type == "null" or type.type == "undefined":
            typeddict_encoder.append("None")
            return (NoneTypeExpr(), [], [], set())
        elif type.type == "Date":
            typeddict_encoder.append("datetime.datetime")
            return (LiteralType("datetime.datetime"), [], [], set())
        elif type.type == "array" and type.items:
            type_name, module_info, type_chunks, encoder_names = encode_type(
                type.items,
                prefix,
                base_model,
                in_module,
                permit_unknown_members=permit_unknown_members,
                type_registry=type_registry,
            )
            typeddict_encoder.append("TODO: dstewart")
            return (ListTypeExpr(type_name), module_info, type_chunks, encoder_names)
        elif (
            type.type == "object"
            and type.patternProperties
            and "^(.*)$" in type.patternProperties
        ):
            type_name, module_info, type_chunks, encoder_names = encode_type(
                type.patternProperties["^(.*)$"],
                prefix,
                base_model,
                in_module,
                permit_unknown_members=permit_unknown_members,
                type_registry=type_registry,
            )
            return (DictTypeExpr(type_name), module_info, type_chunks, encoder_names)
        assert type.type == "object", type.type

        # Register $id for this type so $ref can resolve to it
        if type.id_ is not None:
            type_registry[type.id_] = prefix

        current_chunks: list[str] = [
            f"class {render_literal_type(prefix)}({base_model}):"
        ]
        # For the encoder path, do we need "x" to be bound?
        # lambda x: ... vs lambda _: {}
        needs_binding = False
        encoder_names = set()
        # Track effective field names to detect collisions after normalization
        # Maps effective name -> list of original field names
        effective_field_names: defaultdict[str, list[str]] = defaultdict(list)
        if type.properties:
            needs_binding = True
            typeddict_encoder.append("{")
            for (
                name,
                prop,
            ) in sorted(list(type.properties.items()), key=lambda xs: xs[0]):
                typeddict_encoder.append(f"{repr(name)}:")
                type_name, _, contents, _ = encode_type(
                    prop,
                    TypeName(prefix.value + name.title()),
                    "BaseModel" if base_model == "RiverError" else base_model,
                    in_module,
                    permit_unknown_members=permit_unknown_members,
                    type_registry=type_registry,
                )
                encoder_name = None
                chunks.extend(contents)
                if base_model == "TypedDict":
                    if isinstance(prop, RiverNotType):
                        typeddict_encoder.append("'not implemented'")
                    elif isinstance(prop, RiverUnionType):
                        encoder_name = TypeName(
                            f"encode_{render_literal_type(type_name)}"
                        )
                        encoder_names.add(encoder_name)
                        typeddict_encoder.append(
                            f"{render_literal_type(encoder_name)}(x[{repr(name)}])"
                        )
                        if name not in type.required:
                            typeddict_encoder.append(
                                f"if {repr(name)} in x and x[{repr(name)}] else None"
                            )
                    elif isinstance(prop, RiverIntersectionType):
                        encoder_name = TypeName(
                            f"encode_{render_literal_type(type_name)}"
                        )
                        encoder_names.add(encoder_name)
                        typeddict_encoder.append(f"{encoder_name}(x[{repr(name)}])")
                    elif isinstance(prop, RiverConcreteType):
                        if name == "$kind":
                            safe_name = "kind"
                        else:
                            # For TypedDict encoder, use normalized name to access
                            # the TypedDict field but the output dictionary key should
                            # use the original name
                            if base_model == "TypedDict":
                                safe_name = normalize_special_chars(name)
                            else:
                                safe_name = name
                        if prop.type == "object" and not prop.patternProperties:
                            encoder_name = TypeName(
                                f"encode_{render_literal_type(type_name)}"
                            )
                            encoder_names.add(encoder_name)
                            typeddict_encoder.append(
                                f"{render_literal_type(encoder_name)}(x[{repr(safe_name)}])"
                            )
                            if name not in prop.required:
                                typeddict_encoder.append(
                                    dedent(
                                        f"""
                                        if {repr(safe_name)} in x
                                        and x[{repr(safe_name)}] is not None
                                        else None
                                    """
                                    )
                                )
                        elif prop.type == "array":
                            items = cast(RiverConcreteType, prop).items
                            assert items, "Somehow items was none"
                            if is_literal(cast(RiverType, items)):
                                if name in prop.required:
                                    typeddict_encoder.append(f"x[{repr(name)}]")
                                else:
                                    typeddict_encoder.append(
                                        f"x.get({repr(safe_name)})"
                                    )
                            else:
                                match type_name:
                                    case ListTypeExpr(inner_type_name):
                                        encoder_name = TypeName(
                                            f"encode_{render_literal_type(inner_type_name)}"
                                        )
                                        encoder_names.add(encoder_name)
                                        typeddict_encoder.append(
                                            dedent(
                                                f"""\
                                            [
                                                {render_literal_type(encoder_name)}(y)
                                                for y in x[{repr(name)}]
                                            ]
                                            """.rstrip()
                                            )
                                        )
                                        if name not in prop.required:
                                            typeddict_encoder.append(
                                                dedent(
                                                    f"""
                                                    if {repr(name)} in x
                                                    and x[{repr(name)}] is not None
                                                    else None
                                                """
                                                )
                                            )
                        else:
                            if name in prop.required:
                                typeddict_encoder.append(f"x[{repr(safe_name)}]")
                            else:
                                typeddict_encoder.append(f"x.get({repr(safe_name)})")

                if name == "$kind":
                    # If the field is a literal, the Python type-checker will complain
                    # about the constructor not being able to specify a value for it:
                    # You can't put `$kind="ok"` in the ctor because `$` is not a valid
                    # character in an identifier, and putting `**{"$kind":"ok"}` makes
                    # it not recognize the `"ok"` as being `Literal["ok"]`, so we're
                    # stuck with an impossible-to-construct object.
                    field_value = "..."
                    match type_name:
                        case LiteralTypeExpr(literal_value):
                            field_value = repr(literal_value)
                    if name not in type.required:
                        type_name = UnionTypeExpr(
                            [
                                type_name,
                                NoneTypeExpr(),
                            ]
                        )
                        value = ""
                        if base_model != "TypedDict":
                            value = f"= {repr(None)}"
                    else:
                        value = ""
                        if base_model != "TypedDict":
                            value = f"= {field_value}"
                    # Track $kind -> "kind" mapping for collision detection
                    effective_field_names["kind"].append(name)

                    current_chunks.append(
                        f"  kind: Annotated[{render_type_expr(type_name)}, Field(alias={
                            repr(name)
                        })]{value}"
                    )
                else:
                    specialized_name = normalize_special_chars(name)
                    effective_name = name
                    extras = []
                    if name != specialized_name:
                        if base_model == "BaseModel":
                            # Pydantic doesn't allow leading underscores in field names
                            effective_name = specialized_name
                            extras.append(f"alias={repr(name)}")
                        elif base_model == "TypedDict":
                            # For TypedDict, we use the normalized name directly
                            # TypedDict doesn't support aliases, so we normalize
                            # the field name
                            effective_name = specialized_name
                        else:
                            # For RiverError (which extends BaseModel), use alias
                            # like BaseModel
                            effective_name = specialized_name
                            extras.append(f"alias={repr(name)}")

                    effective_field_names[effective_name].append(name)

                    if name not in type.required:
                        if base_model == "TypedDict":
                            current_chunks.append(
                                reindent(
                                    "  ",
                                    f"""\
                        {effective_name}: NotRequired[{
                                        render_type_expr(
                                            UnionTypeExpr([type_name, NoneTypeExpr()])
                                        )
                                    }]
                                """,
                                )
                            )
                        else:
                            extras.append("default=None")

                            current_chunks.append(
                                reindent(
                                    "  ",
                                    f"""\
                        {effective_name}: {
                                        render_type_expr(
                                            UnionTypeExpr(
                                                [
                                                    type_name,
                                                    NoneTypeExpr(),
                                                ]
                                            )
                                        )
                                    } = Field({", ".join(extras)})
                                """,
                                )
                            )
                    else:
                        extras_str = ""
                        if len(extras) != 0:
                            extras_str = f" = Field({', '.join(extras)})"

                        current_chunks.append(
                            f"  {effective_name}: {render_type_expr(type_name)}{extras_str}"  # noqa: E501
                        )
                typeddict_encoder.append(",")

            # Check for field name collisions after processing all fields
            for effective_name, original_names in effective_field_names.items():
                if len(original_names) > 1:
                    error_msg = (
                        f"Field name collision: fields {original_names} all normalize "
                        f"to the same effective name '{effective_name}'"
                    )

                    raise ValueError(error_msg)

            typeddict_encoder.append("}")
            # exclude_none
            typeddict_encoder = (
                ["{k: v for (k, v) in ("]
                + typeddict_encoder
                + [").items() if v is not None}"]
            )
        else:
            typeddict_encoder.append("{}")
            current_chunks.append("  pass")
        current_chunks.append("")

        if base_model == "TypedDict":
            binding = "x" if needs_binding else "_"
            encoder_name = TypeName(f"encode_{render_literal_type(prefix)}")
            encoder_names.add(encoder_name)
            _field_name = render_literal_type(encoder_name)
            current_chunks.insert(
                0,
                FileContents(
                    "\n".join(
                        [
                            dedent(
                                f"""\
                            def {_field_name}(
                                {binding}: {repr(render_literal_type(prefix))},
                            ) -> Any:
                                return (
                                {
                                    reindent(
                                        "                                ",
                                        "\n".join(typeddict_encoder),
                                    )
                                }
                            )
                            """
                            )
                        ]
                    )
                ),
            )
        chunks.append(FileContents("\n".join(current_chunks)))

    return (prefix, in_module, chunks, encoder_names)


def generate_common_client(
    client_name: str,
    handshake_type: HandshakeType,
    handshake_chunks: Sequence[str],
    modules: list[tuple[ModuleName, ClassName]],
    protocol_version: Literal["v1.1", "v2.0"],
) -> FileContents:
    client_module: str
    match protocol_version:
        case "v1.1":
            client_module = "river"
        case "v2.0":
            client_module = "river.v2"
    chunks: list[str] = [ROOT_FILE_HEADER]
    chunks.extend(
        [
            f"from .{model_name} import {class_name}"
            for model_name, class_name in sorted(modules, key=lambda kv: kv[1])
        ]
    )
    chunks.extend(handshake_chunks)
    chunks.extend(
        [
            dedent(
                f"""\
                class {client_name}:
                  def __init__(self, client: {client_module}.Client[{handshake_type}]):
                """.rstrip()
            )
        ]
    )
    for module_name, class_name in modules:
        chunks.append(
            f"    self.{module_name} = {class_name}(client)",
        )

    return FileContents("\n".join(chunks))


def render_library_call(
    schema_name: str,
    name: str,
    procedure: RiverProcedure,
    protocol_version: Literal["v1.1", "v2.0"],
    init_meta: tuple[RiverType, TypeExpression, str] | None,
    input_meta: tuple[RiverType, TypeExpression, str] | None,
    output_meta: tuple[RiverType, TypeExpression, str] | None,
    error_meta: tuple[RiverType, TypeExpression, str] | None,
) -> list[str]:
    """
    This method is only ever called from one place, but it's defensively establishing a
    namespace that lets us draw some new boundaries around the parameters, without the
    pollution from other intermediatae values.
    """
    current_chunks: list[str] = []

    binding: str
    if procedure.type == "rpc":
        match protocol_version:
            case "v1.1":
                assert input_meta, "rpc expects input to be required"
                _, tpe, render_method = input_meta
                binding = "input"
            case "v2.0":
                assert init_meta, "rpc expects init to be required"
                _, tpe, render_method = init_meta
                binding = "init"
            case other:
                assert_never(other)

        assert output_meta
        assert error_meta
        _, output_type, parse_output_method = output_meta
        _, _, parse_error_method = error_meta

        current_chunks.extend(
            [
                reindent(
                    "  ",
                    f"""\
        async def {name}(
          self,
          {binding}: {render_type_expr(tpe)},
          timeout: datetime.timedelta,
        ) -> {render_type_expr(output_type)}:
          return await self.client.send_rpc(
            {repr(schema_name)},
            {repr(name)},
            {binding},
            {reindent("                    ", render_method)},
            {reindent("                    ", parse_output_method)},
            {reindent("                    ", parse_error_method)},
            timeout,
          )
                """,
                )
            ]
        )
    elif procedure.type == "subscription":
        match protocol_version:
            case "v1.1":
                assert input_meta, "subscription expects input to be required"
                _, tpe, render_method = input_meta
                binding = "input"
            case "v2.0":
                assert init_meta, "subscription expects init to be required"
                _, tpe, render_method = init_meta
                binding = "init"
            case other:
                assert_never(other)

        assert output_meta
        assert error_meta
        _, output_type, parse_output_method = output_meta
        _, error_type, parse_error_method = error_meta
        error_type_name = extract_inner_type(error_type)

        output_or_error_type = UnionTypeExpr([output_type, error_type_name])

        output_or_error_type = UnionTypeExpr(
            [
                output_or_error_type,
                TypeName("RiverError"),
            ]
        )
        current_chunks.extend(
            [
                reindent(
                    "  ",
                    f"""\
        async def {name}(
          self,
          {binding}: {render_type_expr(tpe)},
        ) -> AsyncIterator[{render_type_expr(output_or_error_type)}]:
          return self.client.send_subscription(
            {repr(schema_name)},
            {repr(name)},
            {binding},
            {reindent("                    ", render_method)},
            {reindent("                    ", parse_output_method)},
            {reindent("                    ", parse_error_method)},
          )
              """,
                )
            ]
        )
    elif procedure.type == "upload":
        assert output_meta
        assert error_meta
        _, output_type, parse_output_method = output_meta
        _, error_type, parse_error_method = error_meta
        error_type_name = extract_inner_type(error_type)

        output_or_error_type = UnionTypeExpr([output_type, error_type_name])

        if init_meta and input_meta:
            _, init_type, render_init_method = init_meta
            _, input_type, render_input_method = input_meta
            current_chunks.extend(
                [
                    reindent(
                        "  ",
                        f"""\
        async def {name}(
          self,
          init: {render_type_expr(init_type)},
          inputStream: AsyncIterable[{render_type_expr(input_type)}],
        ) -> {render_type_expr(output_type)}:
          return await self.client.send_upload(
            {repr(schema_name)},
            {repr(name)},
            init,
            inputStream,
            {reindent("                    ", render_init_method)},
            {reindent("                    ", render_input_method)},
            {reindent("                    ", parse_output_method)},
            {reindent("                    ", parse_error_method)},
          )
                    """,
                    )
                ]
            )
        elif protocol_version == "v1.1":
            assert input_meta, "upload requires input to be defined"
            _, input_type, render_input_method = input_meta
            current_chunks.extend(
                [
                    reindent(
                        "  ",
                        f"""\
        async def {name}(
          self,
          inputStream: AsyncIterable[{render_type_expr(input_type)}],
        ) -> {  # TODO(dstewart) This should just be output_type
                            render_type_expr(output_or_error_type)
                        }:
          return await self.client.send_upload(
            {repr(schema_name)},
            {repr(name)},
            None,
            inputStream,
            None,
            {reindent("                    ", render_input_method)},
            {reindent("                    ", parse_output_method)},
            {reindent("                    ", parse_error_method)},
          )
                    """,
                    )
                ]
            )
        elif protocol_version == "v2.0":
            raise ValueError(
                "It is expected that protocol v2 uploads have both init and input "
                "defined, otherwise it's no different than rpc",
            )
        else:
            assert_never(protocol_version)
    elif procedure.type == "stream":
        assert output_meta
        assert error_meta
        _, output_type, parse_output_method = output_meta
        _, error_type, parse_error_method = error_meta
        error_type_name = extract_inner_type(error_type)

        output_or_error_type = UnionTypeExpr([output_type, error_type_name])

        output_or_error_type = UnionTypeExpr(
            [
                output_or_error_type,
                TypeName("RiverError"),
            ]
        )
        if init_meta and input_meta:
            _, init_type, render_init_method = init_meta
            _, input_type, render_input_method = input_meta
            current_chunks.extend(
                [
                    reindent(
                        "  ",
                        f"""\
        async def {name}(
          self,
          init: {render_type_expr(init_type)},
          inputStream: AsyncIterable[{render_type_expr(input_type)}],
        ) -> AsyncIterator[{render_type_expr(output_or_error_type)}]:
          return self.client.send_stream(
            {repr(schema_name)},
            {repr(name)},
            init,
            inputStream,
            {reindent("                    ", render_init_method)},
            {reindent("                    ", render_input_method)},
            {reindent("                    ", parse_output_method)},
            {reindent("                    ", parse_error_method)},
          )
                    """,
                    )
                ]
            )
        elif protocol_version == "v1.1":
            assert input_meta, "stream requires input to be defined"
            _, input_type, render_input_method = input_meta
            current_chunks.extend(
                [
                    reindent(
                        "  ",
                        f"""\
        async def {name}(
          self,
          inputStream: AsyncIterable[{render_type_expr(input_type)}],
        ) -> AsyncIterator[{render_type_expr(output_or_error_type)}]:
          return self.client.send_stream(
            {repr(schema_name)},
            {repr(name)},
            None,
            inputStream,
            None,
            {reindent("                    ", render_input_method)},
            {reindent("                    ", parse_output_method)},
            {reindent("                    ", parse_error_method)},
          )
                    """,
                    )
                ]
            )
        elif protocol_version == "v2.0":
            assert init_meta, "Protocol v2 requires init to be defined"
            _, init_type, render_init_method = init_meta
            current_chunks.extend(
                [
                    reindent(
                        "  ",
                        f"""\
        async def {name}(
          self,
          init: {render_type_expr(init_type)},
        ) -> AsyncIterator[{render_type_expr(output_or_error_type)}]:
          return self.client.send_stream(
            {repr(schema_name)},
            {repr(name)},
            init,
            None,
            {reindent("                    ", render_init_method)},
            None,
            {reindent("                    ", parse_output_method)},
            {reindent("                    ", parse_error_method)},
          )
                    """,
                    )
                ]
            )
        else:
            assert_never(protocol_version)

    current_chunks.append("")
    return current_chunks


def generate_individual_service(
    schema_name: str,
    schema: RiverService,
    input_base_class: Literal["TypedDict"] | Literal["BaseModel"],
    method_filter: set[str] | None,
    protocol_version: Literal["v1.1", "v2.0"],
) -> tuple[ModuleName, ClassName, dict[RenderedPath, FileContents]]:
    serdes: list[tuple[list[TypeName], list[ModuleName], list[FileContents]]] = []

    def _type_adapter_definition(
        type_adapter_name: TypeName,
        _type: TypeExpression,
        module_info: list[ModuleName],
    ) -> tuple[list[TypeName], list[ModuleName], list[FileContents]]:
        varname = render_type_expr(type_adapter_name)
        rendered_type_expr = render_type_expr(_type)
        return (
            [type_adapter_name],
            module_info,
            [
                FileContents(
                    dedent(f"""
                {varname}: TypeAdapter[{rendered_type_expr}] = (
                    TypeAdapter({rendered_type_expr})
                )
            """)
                )
            ],
        )

    client_module: str
    match protocol_version:
        case "v1.1":
            client_module = "river"
        case "v2.0":
            client_module = "river.v2"

    class_name = ClassName(f"{schema_name.title()}Service")
    current_chunks: list[str] = [
        dedent(
            f"""\
              class {class_name}:
                def __init__(self, client: {client_module}.Client[Any]):
                  self.client = client
            """
        ),
    ]
    for name, procedure in schema.procedures.items():
        if method_filter and (schema_name + "." + name) in method_filter:
            continue
        module_names = [ModuleName(name)]
        init_type: TypeExpression | None = None
        init_module_info: list[ModuleName] = []
        if procedure.init:
            init_type, init_module_info, init_chunks, encoder_names = encode_type(
                procedure.init,
                TypeName(f"{name.title()}Init"),
                input_base_class,
                module_names,
                permit_unknown_members=False,
            )
            serdes.append(
                (
                    [extract_inner_type(init_type), *encoder_names],
                    init_module_info,
                    init_chunks,
                )
            )
        input_type: TypeExpression | None = None
        input_module_info: list[ModuleName] = []
        if procedure.input:
            input_type, input_module_info, input_chunks, encoder_names = encode_type(
                procedure.input,
                TypeName(f"{name.title()}Input"),
                input_base_class,
                module_names,
                permit_unknown_members=False,
            )
            serdes.append(
                (
                    [extract_inner_type(input_type), *encoder_names],
                    input_module_info,
                    input_chunks,
                )
            )

        output_type, output_module_info, output_chunks, encoder_names = encode_type(
            procedure.output,
            TypeName(f"{name.title()}Output"),
            "BaseModel",
            module_names,
            permit_unknown_members=True,
        )
        output_type_name = extract_inner_type(output_type)
        serdes.append(
            (
                [output_type_name, *encoder_names],
                output_module_info,
                output_chunks,
            )
        )
        output_type_type_adapter_name = TypeName(
            f"{render_literal_type(output_type_name)}TypeAdapter"
        )
        serdes.append(
            _type_adapter_definition(
                output_type_type_adapter_name, output_type, output_module_info
            )
        )

        if procedure.errors:
            error_type, error_module_info, errors_chunks, encoder_names = encode_type(
                procedure.errors,
                TypeName(f"{name.title()}Errors"),
                "RiverError",
                module_names,
                permit_unknown_members=True,
            )
            if isinstance(error_type, NoneTypeExpr):
                error_type_name = TypeName("RiverError")
                error_type = error_type_name
            else:
                error_type_name = extract_inner_type(error_type)
                serdes.append(([error_type_name], error_module_info, errors_chunks))

        else:
            error_type_name = TypeName("RiverError")
            error_type = error_type_name

        error_type_type_adapter_name = TypeName(
            f"{render_literal_type(error_type_name)}TypeAdapter"
        )
        if error_type_type_adapter_name.value != "RiverErrorTypeAdapter":
            serdes.append(
                _type_adapter_definition(
                    error_type_type_adapter_name, error_type, output_module_info
                )
            )

        # NB: These strings must be indented to at least the same level of
        #     the function strings in the branches below, otherwise `dedent`
        #     will pick our indentation level for normalization, which will
        #     break the "def" indentation presuppositions.
        output_type_adapter = render_literal_type(output_type_type_adapter_name)
        parse_output_method = f"""\
                            lambda x: {output_type_adapter}
                                .validate_python(
                                    x # type: ignore[arg-type]
                                )
                            """
        error_type_adapter = render_literal_type(error_type_type_adapter_name)
        parse_error_method = f"""\
                            lambda x: {error_type_adapter}
                                .validate_python(
                                    x # type: ignore[arg-type]
                                )
                            """

        # Init renderer
        render_init_method: str | None = None
        if init_type and procedure.init is not None:
            if input_base_class == "TypedDict":
                if is_literal(procedure.init):
                    render_init_method = "lambda x: x"
                elif isinstance(
                    procedure.init, RiverConcreteType
                ) and procedure.init.type in ["array"]:
                    match init_type:
                        case ListTypeExpr(init_type_name):
                            render_init_method = f"lambda xs: [encode_{
                                render_literal_type(init_type_name)
                            }(x) for x in xs]"
                else:
                    render_init_method = f"encode_{render_literal_type(init_type)}"
            else:
                init_type_name = extract_inner_type(init_type)
                init_type_type_adapter_name = TypeName(
                    f"{init_type_name.value}TypeAdapter"
                )
                serdes.append(
                    _type_adapter_definition(
                        init_type_type_adapter_name, init_type, init_module_info
                    )
                )
                render_init_method = f"""\
                    lambda x: {render_type_expr(init_type_type_adapter_name)}
                        .validate_python(x)
                """

        # Input renderer
        render_input_method: str | None = None
        if input_type and procedure.input is not None:
            if input_base_class == "TypedDict":
                if is_literal(procedure.input):
                    render_input_method = "lambda x: x"
                elif isinstance(
                    procedure.input, RiverConcreteType
                ) and procedure.input.type in ["array"]:
                    match input_type:
                        case ListTypeExpr(list_type):
                            render_input_method = f"""\
                            lambda xs: [
                                encode_{render_literal_type(list_type)}(x) for x in xs
                            ]
                            """
                else:
                    render_input_method = f"encode_{render_literal_type(input_type)}"
            else:
                input_type_name = extract_inner_type(input_type)
                input_type_type_adapter = TypeName(
                    f"{render_literal_type(input_type_name)}TypeAdapter"
                )
                serdes.append(
                    _type_adapter_definition(
                        input_type_type_adapter, input_type, input_module_info
                    )
                )
                render_input_method = f"""\
                                lambda x: {render_type_expr(input_type_type_adapter)}
                                  .dump_python(
                                    x, # type: ignore[arg-type]
                                    by_alias=True,
                                    exclude_none=True,
                                  )
                                """
        if isinstance(output_type, NoneTypeExpr):
            parse_output_method = "lambda x: None"

        def combine_or_none(
            proc_type: RiverType | None,
            tpe: TypeExpression | None,
            serde_method: str | None,
        ) -> tuple[RiverType, TypeExpression, str] | None:
            if not proc_type and not tpe and not serde_method:
                return None
            if not proc_type or not tpe or not serde_method:
                raise ValueError(
                    f"Unable to convert {repr(proc_type)} into either"
                    f" tpe={tpe} or render_method={serde_method}"
                )
            return (proc_type, tpe, serde_method)

        current_chunks.extend(
            render_library_call(
                schema_name=schema_name,
                name=name,
                procedure=procedure,
                protocol_version=protocol_version,
                init_meta=combine_or_none(
                    procedure.init, init_type, render_init_method
                ),
                input_meta=combine_or_none(
                    procedure.input, input_type, render_input_method
                ),
                output_meta=combine_or_none(
                    procedure.output, output_type, parse_output_method
                ),
                error_meta=combine_or_none(
                    procedure.errors, error_type, parse_error_method
                ),
            )
        )

    emitted_files: dict[RenderedPath, FileContents] = {}

    imports: dict[str, set[TypeName]] = dict()

    in_root: list[FileContents] = []
    for names, module_info, contents in serdes:
        if not module_info:
            in_root.extend(contents)
            continue

        imports.setdefault(".".join(module_info), set()).update(names)

        leaf_module = module_info[-1] + ".py"
        file_path = RenderedPath(str(Path(schema_name, *module_info[:-1], leaf_module)))
        existing = emitted_files.get(file_path, FileContents(FILE_HEADER))
        emitted_files[file_path] = FileContents("\n".join([existing] + contents))

    def render_names(xs: set[TypeName]) -> str:
        return ", ".join(sorted(render_literal_type(x) for x in xs))

    rendered_imports = [
        f"from .{dotted_modules} import {render_names(names)}"
        for dotted_modules, names in imports.items()
    ]

    emitted_files[RenderedPath(str(Path(f"{schema_name}/__init__.py")))] = FileContents(
        "\n".join([SERVICE_FILE_HEADER] + rendered_imports + in_root + current_chunks)
    )
    return (
        ModuleName(schema_name),
        class_name,
        emitted_files,
    )


def generate_river_client_module(
    client_name: str,
    schema_root: RiverSchema,
    typed_dict_inputs: bool,
    method_filter: set[str] | None,
    protocol_version: Literal["v1.1", "v2.0"],
) -> dict[RenderedPath, FileContents]:
    files: dict[RenderedPath, FileContents] = {}

    # Negotiate handshake shape
    handshake_chunks: list[str] = []
    if schema_root.handshakeSchema is not None:
        _handshake_type, _, contents, _ = encode_type(
            schema_root.handshakeSchema,
            TypeName("HandshakeSchema"),
            "BaseModel",
            [],
            permit_unknown_members=False,
        )
        handshake_chunks.extend(contents)
        handshake_type = HandshakeType(render_type_expr(_handshake_type))
    else:
        handshake_type = HandshakeType("Literal[None]")

    modules: list[tuple[ModuleName, ClassName]] = []
    input_base_class: Literal["TypedDict"] | Literal["BaseModel"] = (
        "TypedDict" if typed_dict_inputs else "BaseModel"
    )
    for schema_name, schema in schema_root.services.items():
        module_name, class_name, emitted_files = generate_individual_service(
            schema_name,
            schema,
            input_base_class,
            method_filter,
            protocol_version,
        )
        if emitted_files:
            # Short-cut if we didn't actually emit anything
            files.update(emitted_files)
            modules.append((module_name, class_name))

    main_contents = generate_common_client(
        client_name,
        handshake_type,
        handshake_chunks,
        modules,
        protocol_version,
    )
    files[RenderedPath(str(Path("__init__.py")))] = main_contents

    return files


def schema_to_river_client_codegen(
    read_schema: Callable[[], TextIO],
    target_path: str,
    client_name: str,
    typed_dict_inputs: bool,
    file_opener: Callable[[Path], TextIO],
    method_filter: set[str] | None,
    protocol_version: Literal["v1.1", "v2.0"],
) -> None:
    """Generates the lines of a River module."""
    with read_schema() as f:
        schemas = RiverSchemaFile(json.load(f))
    for subpath, contents in generate_river_client_module(
        client_name,
        schemas.root,
        typed_dict_inputs,
        method_filter,
        protocol_version,
    ).items():
        module_path = Path(target_path).joinpath(subpath)
        module_path.parent.mkdir(mode=0o755, parents=True, exist_ok=True)
        with file_opener(module_path) as f:
            try:
                popen = subprocess.Popen(
                    ["ruff", "format", "-"],
                    stdin=subprocess.PIPE,
                    stdout=subprocess.PIPE,
                )
                stdout, _ = popen.communicate(contents.encode())
                if popen.returncode != 0:
                    f.write(contents)
                f.write(stdout.decode("utf-8"))
            except:
                f.write(contents)
                raise