from construct import (
    Int8ul,
    Int32ul,
    Int64ul,
    Flag,
    PascalString,
    PrefixedArray,
    Struct as CStruct,
    Enum,
    Switch,
    this,
)

MessageInChat = CStruct(
    "role" / PascalString(Int32ul, "utf8"),
    "content" / PascalString(Int32ul, "utf8"),
)

ChatCompletionRequest = CStruct(
    "messages" / PrefixedArray(Int32ul, MessageInChat),
)

ChatCompletionCancellationRequest = CStruct()

RequestKind = Enum(Int8ul, chat=0, cancel=1)

Request = CStruct(
    "request_id" / Int64ul,
    "kind" / RequestKind,
    "payload" / Switch(this.kind, {
        "chat": ChatCompletionRequest,
        "cancel": ChatCompletionCancellationRequest,
    }),
)

ResponseChatCompletion = CStruct(
    "piece" / PascalString(Int32ul, "utf8"),
)

ResponseChatCompletionCancellation = CStruct()
ResponseChatCompletionEnd = CStruct()

ResponseKind = Enum(Int8ul, chat=0, cancel=1, end=2)

Response = CStruct(
    "request_id" / Int64ul,
    "kind" / ResponseKind,
    "payload" / Switch(this.kind, {
        "chat": ResponseChatCompletion,
        "cancel": ResponseChatCompletionCancellation,
        "end": ResponseChatCompletionEnd
    }),
)