Seed Form Filling Session Context

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

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 FormFillingSessionContext(TypedDict, total=False):
    # Session context. To generate filled forms, include `form_filling` with a non-empty `values` array.
    # The request body MUST include `form_filling` when you expect filled form outputs.
    form_filling: FormFillingContext


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 _post_json_expect(
    url: str,
    headers: dict[str, str],
    payload: dict[str, Any],
    expect_status: int,
) -> Optional[dict[str, Any]]:
    r = requests.post(url, json=payload, headers=headers, timeout=60)
    if r.status_code == expect_status:
        # OpenAPI does not define a response body for 200 here, so treat it as optional.
        if not r.text:
            return None
        try:
            data = r.json()
        except ValueError:
            return None
        return data if isinstance(data, dict) else None

    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 seed_form_filling_session_context(
    suki_token: str,
    ambient_session_id: str,
    body: Optional[FormFillingSessionContext] = None,
) -> None:
    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"}
    _post_json_expect(url, headers, dict(body or {}), 200)


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

        # The request body MUST include `form_filling` with a non-empty `values` array to generate filled forms.
        # Each values[] entry must use either form_template_id or schema, never both.
        seed_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"}},
                            },
                        },
                    ]
                }
            },
        )
        print("Context seeded.")
    except (ApiHttpError, ValueError) as e:
        print(e)