Update Form Filling Session Context

from typing import Any, NotRequired, Optional, TypedDict, Union, cast

import requests

BASE_URL = "https://sdp.suki-stage.com"


class StaticFormEntry(TypedDict):
    # Suki Medical form template. Do not also send `schema` in this object.
    form_template_id: str


class DynamicFormEntry(TypedDict):
    # Partner-defined form. Do not also send `form_template_id` in this object.
    schema: dict[str, Any]
    id: NotRequired[str]
    name: NotRequired[str]
    type: NotRequired[str]


FormFillingMetadata = Union[StaticFormEntry, DynamicFormEntry]


class FormFillingContext(TypedDict):
    # Non-empty list. Mix static and dynamic entries when needed.
    values: list[FormFillingMetadata]


class FormFillingUpdateContextRequest(TypedDict, total=False):
    form_filling: FormFillingContext


class FormFillingReturnedContext(TypedDict, total=False):
    form_filling: FormFillingContext


class FormFillingUpdateContextResponse(TypedDict):
    context: FormFillingReturnedContext


class ApiHttpError(RuntimeError):
    def __init__(self, status: int, url: str, detail: str) -> None:
        super().__init__(f"HTTP {status} {url}: {detail}")
        self.status = status
        self.url = url


def _patch_json_expect(
    url: str,
    headers: dict[str, str],
    payload: dict[str, Any],
    expect_status: int,
) -> dict[str, Any]:
    r = requests.patch(url, json=payload, headers=headers, timeout=60)
    if r.status_code == expect_status:
        data = r.json()
        if isinstance(data, dict):
            return data
        raise ApiHttpError(expect_status, url, "response JSON was not an object")

    detail = ""
    try:
        err = r.json()
        if isinstance(err, dict) and isinstance(err.get("message"), str):
            detail = err["message"]
    except ValueError:
        detail = (r.text or "")[:500]
    raise ApiHttpError(r.status_code, url, detail or "(no body)")


def login_for_suki_token(partner_id: str, partner_token: str, *, provider_id: Optional[str] = None) -> str:
    url = f"{BASE_URL}/api/v1/auth/login"
    body: dict[str, Any] = {"partner_id": partner_id, "partner_token": partner_token}
    if provider_id is not None:
        body["provider_id"] = provider_id
    r = requests.post(url, json=body, headers={"Content-Type": "application/json"}, timeout=60)
    if r.status_code != 200:
        raise ApiHttpError(r.status_code, url, (r.text or "")[:500] or "(no body)")
    data = r.json()
    token = data.get("suki_token") if isinstance(data, dict) else None
    if not isinstance(token, str) or not token:
        raise ValueError(f"{url}: 200 response missing suki_token")
    return token


def update_form_filling_session_context(
    suki_token: str,
    ambient_session_id: str,
    body: FormFillingUpdateContextRequest,
) -> FormFillingUpdateContextResponse:
    """PATCH /api/v1/form-filling/session/{ambient_session_id}/context. HTTP 200 returns FormFillingUpdateContextResponse."""
    url = f"{BASE_URL}/api/v1/form-filling/session/{ambient_session_id}/context"
    headers = {"sdp_suki_token": suki_token, "sdp_provider_id": "<sdp_provider_id>", "Content-Type": "application/json"}
    data = _patch_json_expect(url, headers, dict(body), 200)
    ctx = data.get("context")
    if not isinstance(ctx, dict):
        raise ValueError(f"{url}: 200 response missing context")
    return cast(FormFillingUpdateContextResponse, {"context": cast(FormFillingReturnedContext, ctx)})


if __name__ == "__main__":
    try:
        token = login_for_suki_token("<partner_id>", "<partner_token>")

        # Example mixes one Suki template and one partner schema.
        # Each values[] entry must use either form_template_id or schema, never both.
        out = update_form_filling_session_context(
            token,
            ambient_session_id="<ambient_session_id>",
            body={
                "form_filling": {
                    "values": [
                        {"form_template_id": "019d4cdc-9319-7d81-ae2e-fd6de7f1b4f0"},
                        {
                            "id": "partner-custom-1",
                            "name": "Custom Partner Form",
                            "type": "NEURO_ASSESSMENT",
                            "schema": {
                                "type": "object",
                                "properties": {"gcs": {"type": "number"}},
                            },
                        },
                    ]
                }
            },
        )
        form_filling = out["context"].get("form_filling")
        print(form_filling)
    except (ApiHttpError, ValueError) as e:
        print(e)