#!/usr/bin/env python3

from pydantic import BaseModel, TypeAdapter
from result import Result, Ok, Err
from typing import TypeVar, List, Literal

T = TypeVar("T")

def parse_response(ResponseType: T, json_response: dict) -> Result[T, str]:
    try:
        return Ok(TypeAdapter(ResponseType).validate_json(json_response))
    except Exception as e:
        return Err(str(e))

class GroupEntry(BaseModel):
    admins: List[str]
    blocked: bool
    id: str
    internal_id: str
    invite_link: str
    members: List[str]
    name: str
    pending_invites: List[str]
    pending_requests: List[str]

class IdentityEntry(BaseModel):
    added: str
    fingerprint: str
    number: str
    safety_number: str
    status: str

class TrustIdentityRequest(BaseModel):
    pass

class TrustAllKnownKeys(TrustIdentityRequest):
    trust_all_known_keys: Literal[True] = True

class TrustSafetyNumber(TrustIdentityRequest):
    verified_safety_number: str

class UsernameSetResponse(BaseModel):
    username: str
    username_link: str

GroupLinkPolicy = Literal["disabled", "enabled", "enabled-with-approval"]
GroupPermissionPolicy = Literal["only-admins", "every-member"]

class GroupPermissions(BaseModel):
    add_members: GroupPermissionPolicy
    edit_group: GroupPermissionPolicy

class CreateGroupRequest(BaseModel):
    description: str = ""
    expiration_time: int = 0 # Expiration time of messages
    group_link: GroupLinkPolicy = "disabled"
    members: List[str]
    name: str
    permissions: GroupPermissions = GroupPermissions(add_members="only-admins", edit_group="only-admins")

class CreateGroupResponse(BaseModel):
    id: str

def test_reaction_message():
    data = """{
        "envelope": {
            "source": "+4900000000001",
            "sourceNumber": "+4900000000001",
            "sourceUuid": "00000000-0000-0000-0000-000000000000",
            "sourceName": "lemoer",
            "sourceDevice": 1,
            "timestamp": 1734201022564,
            "dataMessage": {
                "timestamp": 1734201022564,
                "message": null,
                "expiresInSeconds": 0,
                "viewOnce": false,
                "reaction": {
                    "emoji": "👎",
                    "targetAuthor": "+4900000000001",
                    "targetAuthorNumber": "+4900000000001",
                    "targetAuthorUuid": "00000000-0000-0000-0000-000000000000",
                    "targetSentTimestamp": 1734201003509,
                    "isRemove": false
                }
            }
        },
        "account": "+4900000000002"
    }"""

    res = parse_response(Message, data)

    m = res.unwrap()

    assert m.envelope.source == "+4900000000001"
    assert m.envelope.sourceNumber == "+4900000000001"
    assert m.envelope.sourceUuid == "00000000-0000-0000-0000-000000000000"
    assert m.envelope.sourceName == "lemoer"
    assert m.envelope.sourceDevice == 1
    assert m.envelope.timestamp == 1734201022564
    assert m.envelope.dataMessage.timestamp == 1734201022564
    assert m.envelope.dataMessage.message is None
    assert m.envelope.dataMessage.expiresInSeconds == 0
    assert m.envelope.dataMessage.viewOnce is False
    assert m.envelope.dataMessage.reaction.emoji == "👎"
    assert m.envelope.dataMessage.reaction.targetAuthor == "+4900000000001"
    assert m.envelope.dataMessage.reaction.targetAuthorNumber == "+4900000000001"
    assert m.envelope.dataMessage.reaction.targetAuthorUuid == "00000000-0000-0000-0000-000000000000"
    assert m.envelope.dataMessage.reaction.targetSentTimestamp == 1734201003509
    assert m.envelope.dataMessage.reaction.isRemove is False
    assert m.account == "+4900000000002"

def test_simple_message():
    data = """{
        "envelope": {
            "source": "+4900000000001",
            "sourceNumber": "+4900000000001",
            "sourceUuid": "00000000-0000-0000-0000-000000000001",
            "sourceName": "Leo",
            "sourceDevice": 1,
            "timestamp": 1734300324644,
            "dataMessage": {
                "timestamp": 1734300324644,
                "message": "Test",
                "expiresInSeconds": 0,
                "viewOnce": false
            }
        },
        "account": "+4900000000002"
    }"""

    res = parse_response(Message, data)

    m = res.unwrap()

    assert m.envelope.source == "+4900000000001"
    assert m.envelope.sourceNumber == "+4900000000001"
    assert m.envelope.sourceUuid == "00000000-0000-0000-0000-000000000001"
    assert m.envelope.sourceName == "Leo"
    assert m.envelope.sourceDevice == 1
    assert m.envelope.timestamp == 1734300324644
    assert m.envelope.dataMessage.timestamp == 1734300324644
    assert m.envelope.dataMessage.message == "Test"
    assert m.envelope.dataMessage.expiresInSeconds == 0
    assert m.envelope.dataMessage.viewOnce is False
    assert m.account == "+4900000000002"

def test_reaction_removal_message():
    data = """{
        "envelope": {
            "source": "+4900000000001",
            "sourceNumber": "+4900000000001",
            "sourceUuid": "00000000-0000-0000-0000-000000000002",
            "sourceName": "Leo",
            "sourceDevice": 1,
            "timestamp": 1734300542349,
            "dataMessage": {
                "timestamp": 1734300542349,
                "message": null,
                "expiresInSeconds": 0,
                "viewOnce": false,
                "reaction": {
                    "emoji": "👍",
                    "targetAuthor": "+4900000000001",
                    "targetAuthorNumber": "+4900000000001",
                    "targetAuthorUuid": "00000000-0000-0000-0000-000000000002",
                    "targetSentTimestamp": 1734300324644,
                    "isRemove": true
                }
            }
        },
        "account": "+4900000000002"
    }"""

    res = parse_response(Message, data)

    assert res.is_ok()
    m = res.unwrap()

    assert isinstance(m.envelope, EnvelopeData)
    assert m.envelope.source == "+4900000000001"
    assert m.envelope.sourceNumber == "+4900000000001"
    assert m.envelope.sourceUuid == "00000000-0000-0000-0000-000000000002"
    assert m.envelope.sourceName == "Leo"
    assert m.envelope.sourceDevice == 1
    assert m.envelope.timestamp == 1734300542349
    assert isinstance(m.envelope.dataMessage, ReactionMessage)
    assert m.envelope.dataMessage.timestamp == 1734300542349
    assert m.envelope.dataMessage.message is None
    assert m.envelope.dataMessage.expiresInSeconds == 0
    assert m.envelope.dataMessage.viewOnce is False
    assert m.envelope.dataMessage.reaction.emoji == "👍"
    assert m.envelope.dataMessage.reaction.targetAuthor == "+4900000000001"
    assert m.envelope.dataMessage.reaction.targetAuthorNumber == "+4900000000001"
    assert m.envelope.dataMessage.reaction.targetAuthorUuid == "00000000-0000-0000-0000-000000000002"
    assert m.envelope.dataMessage.reaction.targetSentTimestamp == 1734300324644
    assert m.envelope.dataMessage.reaction.isRemove is True
    assert m.account == "+4900000000002"

def test_typing_started_message():
    data = """{
        "envelope": {
            "source": "+4900000000001",
            "sourceNumber": "+4900000000001",
            "sourceUuid": "00000000-0000-0000-0000-000000000000",
            "sourceName": "Leo",
            "sourceDevice": 1,
            "timestamp": 1734300998928,
            "typingMessage": {
                "action": "STARTED",
                "timestamp": 1734300998928
            }
        },
        "account": "+4900000000002"
    }"""

    res = parse_response(Message, data)

    assert res.is_ok()
    m = res.unwrap()

    assert isinstance(m.envelope, EnvelopeTyping)
    assert m.envelope.source == "+4900000000001"
    assert m.envelope.sourceNumber == "+4900000000001"
    assert m.envelope.sourceUuid == "00000000-0000-0000-0000-000000000000"
    assert m.envelope.sourceName == "Leo"
    assert m.envelope.sourceDevice == 1
    assert m.envelope.timestamp == 1734300998928
    assert isinstance(m.envelope.typingMessage, TypingMessage)
    assert isinstance(m.envelope.typingMessage, TypingStarted)
    assert m.envelope.typingMessage.action == "STARTED"
    assert m.envelope.typingMessage.timestamp == 1734300998928
    assert m.account == "+4900000000002"

def test_typing_stopped_message():
    data = """{
        "envelope": {
            "source": "+4900000000001",
            "sourceNumber": "+4900000000001",
            "sourceUuid": "00000000-0000-0000-0000-000000000000",
            "sourceName": "Leo",
            "sourceDevice": 1,
            "timestamp": 1734301001916,
            "typingMessage": {
                "action": "STOPPED",
                "timestamp": 1734301001916
            }
        },
        "account": "+4900000000002"
    }"""

    res = parse_response(Message, data)

    assert res.is_ok()
    m = res.unwrap()

    assert isinstance(m.envelope, EnvelopeTyping)
    assert m.envelope.source == "+4900000000001"
    assert m.envelope.sourceNumber == "+4900000000001"
    assert m.envelope.sourceUuid == "00000000-0000-0000-0000-000000000000"
    assert m.envelope.sourceName == "Leo"
    assert m.envelope.sourceDevice == 1
    assert m.envelope.timestamp == 1734301001916
    assert isinstance(m.envelope.typingMessage, TypingMessage)
    assert isinstance(m.envelope.typingMessage, TypingStopped)
    assert m.envelope.typingMessage.action == "STOPPED"
    assert m.envelope.typingMessage.timestamp == 1734301001916
    assert m.account == "+4900000000002"

def test_typing_message_in_group_chat():
    data = """{
        "envelope": {
            "source": "+4900000000001",
            "sourceNumber": "+4900000000001",
            "sourceUuid": "00000000-0000-0000-0000-000000000000",
            "sourceName": "Leo",
            "sourceDevice": 1,
            "timestamp": 1734301695337,
            "typingMessage": {
                "action": "STARTED",
                "timestamp": 1734301695337,
                "groupId": "nnJOsIQnEvJQ6tvddEkMh9VAyF+VGgwsMO1i5glGjpQ="
            }
        },
        "account": "+4900000000002"
    }"""

    res = parse_response(Message, data)

    assert res.is_ok()
    m = res.unwrap()

    assert isinstance(m.envelope, EnvelopeTyping)
    assert m.envelope.source == "+4900000000001"
    assert m.envelope.sourceNumber == "+4900000000001"
    assert m.envelope.sourceUuid == "00000000-0000-0000-0000-000000000000"
    assert m.envelope.sourceName == "Leo"
    assert m.envelope.sourceDevice == 1
    assert m.envelope.timestamp == 1734301695337
    assert isinstance(m.envelope.typingMessage, TypingMessage)
    assert isinstance(m.envelope.typingMessage, TypingStarted)
    assert m.envelope.typingMessage.action == "STARTED"
    assert m.envelope.typingMessage.timestamp == 1734301695337
    assert m.envelope.typingMessage.groupId == "nnJOsIQnEvJQ6tvddEkMh9VAyF+VGgwsMO1i5glGjpQ="
    assert m.account == "+4900000000002"

def test_group_message():
    data = """{
        "envelope": {
            "source": "+4900000000001",
            "sourceNumber": "+4900000000001",
            "sourceUuid": "00000000-0000-0000-0000-000000000000",
            "sourceName": "Leo",
            "sourceDevice": 1,
            "timestamp": 1734301695867,
            "dataMessage": {
                "timestamp": 1734301695867,
                "message": "Bla",
                "expiresInSeconds": 0,
                "viewOnce": false,
                "groupInfo": {
                    "groupId": "nnJOsIQnEvJQ6tvddEkMh9VAyF+VGgwsMO1i5glGjpQ=",
                    "type": "DELIVER"
                }
            }
        },
        "account": "+4900000000002"
    }"""

    res = parse_response(Message, data)

    assert res.is_ok()
    m = res.unwrap()

    assert isinstance(m.envelope, EnvelopeData)
    assert m.envelope.source == "+4900000000001"
    assert m.envelope.sourceNumber == "+4900000000001"
    assert m.envelope.sourceUuid == "00000000-0000-0000-0000-000000000000"
    assert m.envelope.sourceName == "Leo"
    assert m.envelope.sourceDevice == 1
    assert m.envelope.timestamp == 1734301695867
    assert isinstance(m.envelope.dataMessage, DataMessage)

def test_delete_message_in_group_chat():
    data = """{
        "envelope": {
            "source": "+4900000000001",
            "sourceNumber": "+4900000000001",
            "sourceUuid": "00000000-0000-0000-0000-000000000000",
            "sourceName": "Leo",
            "sourceDevice": 1,
            "timestamp": 1734302431760,
            "dataMessage": {
                "timestamp": 1734302431760,
                "message": null,
                "expiresInSeconds": 0,
                "viewOnce": false,
                "remoteDelete": {
                    "timestamp": 1734302411689
                },
                "groupInfo": {
                    "groupId": "nnJOsIQnEvJQ6tvddEkMh9VAyF+VGgwsMO1i5glGjpQ=",
                    "type": "DELIVER"
                }
            }
        },
        "account": "+4900000000002"
    }"""

    res = parse_response(Message, data)

    assert res.is_ok()
    m = res.unwrap()

    assert isinstance(m.envelope, EnvelopeData)
    assert m.envelope.source == "+4900000000001"
    assert m.envelope.sourceNumber == "+4900000000001"
    assert m.envelope.sourceUuid == "00000000-0000-0000-0000-000000000000"
    assert m.envelope.sourceName == "Leo"
    assert m.envelope.sourceDevice == 1
    assert m.envelope.timestamp == 1734302431760
    assert isinstance(m.envelope.dataMessage, DeleteMessage)
    assert m.envelope.dataMessage.timestamp == 1734302431760
    assert m.envelope.dataMessage.message is None
    assert m.envelope.dataMessage.expiresInSeconds == 0
    assert m.envelope.dataMessage.viewOnce is False
    assert m.envelope.dataMessage.remoteDelete.timestamp == 1734302411689
    assert m.envelope.dataMessage.groupInfo.groupId == "nnJOsIQnEvJQ6tvddEkMh9VAyF+VGgwsMO1i5glGjpQ="
    assert m.envelope.dataMessage.groupInfo.type == "DELIVER"
    assert m.account == "+4900000000002"

def test_reaction_message_in_group_chat():
    data = """{
        "envelope": {
            "source": "+4900000000001",
            "sourceNumber": "+4900000000001",
            "sourceUuid": "00000000-0000-0000-0000-000000000000",
            "sourceName": "Leo",
            "sourceDevice": 1,
            "timestamp": 1734302940233,
            "dataMessage": {
                "timestamp": 1734302940233,
                "message": null,
                "expiresInSeconds": 0,
                "viewOnce": false,
                "reaction": {
                    "emoji": "👍",
                    "targetAuthor": "+4900000000001",
                    "targetAuthorNumber": "+4900000000001",
                    "targetAuthorUuid": "00000000-0000-0000-0000-000000000000",
                    "targetSentTimestamp": 1733878988615,
                    "isRemove": false
                },
                "groupInfo": {
                    "groupId": "nnJOsIQnEvJQ6tvddEkMh9VAyF+VGgwsMO1i5glGjpQ=",
                    "type": "DELIVER"
                }
            }
        },
        "account": "+4900000000002"
    }"""

    res = parse_response(Message, data)

    assert res.is_ok()
    m = res.unwrap()

    assert isinstance(m.envelope, EnvelopeData)
    assert m.envelope.source == "+4900000000001"
    assert m.envelope.sourceNumber == "+4900000000001"
    assert m.envelope.sourceUuid == "00000000-0000-0000-0000-000000000000"
    assert m.envelope.sourceName == "Leo"
    assert m.envelope.sourceDevice == 1
    assert m.envelope.timestamp == 1734302940233
    assert isinstance(m.envelope.dataMessage, ReactionMessage)
    assert m.envelope.dataMessage.timestamp == 1734302940233
    assert m.envelope.dataMessage.message is None
    assert m.envelope.dataMessage.expiresInSeconds == 0
    assert m.envelope.dataMessage.viewOnce is False
    assert m.envelope.dataMessage.reaction.emoji == "👍"
    assert m.envelope.dataMessage.reaction.targetAuthor == "+4900000000001"
    assert m.envelope.dataMessage.reaction.targetAuthorNumber == "+4900000000001"
    assert m.envelope.dataMessage.reaction.targetAuthorUuid == "00000000-0000-0000-0000-000000000000"
    assert m.envelope.dataMessage.reaction.targetSentTimestamp == 1733878988615
    assert m.envelope.dataMessage.reaction.isRemove is False
    assert m.envelope.dataMessage.groupInfo.groupId == "nnJOsIQnEvJQ6tvddEkMh9VAyF+VGgwsMO1i5glGjpQ="
    assert m.envelope.dataMessage.groupInfo.type == "DELIVER"

# Reformat, replace my number and uuids with placeholder
def test_delete_message_in_normal_chat():
    data = """{
        "envelope": {
            "source": "+4900000000001",
            "sourceNumber": "+4900000000001",
            "sourceUuid": "00000000-0000-0000-0000-000000000000",
            "sourceName": "Leo",
            "sourceDevice": 1,
            "timestamp": 1734304012251,
            "dataMessage": {
                "timestamp": 1734304012251,
                "message": null,
                "expiresInSeconds": 0,
                "viewOnce": false,
                "remoteDelete": {
                    "timestamp": 1734300999933
                }
            }
        },
        "account": "+4900000000002"
    }"""

    res = parse_response(Message, data)

    assert res.is_ok()
    m = res.unwrap()

    assert isinstance(m.envelope, EnvelopeData)
    assert m.envelope.source == "+4900000000001"
    assert m.envelope.sourceNumber == "+4900000000001"
    assert m.envelope.sourceUuid == "00000000-0000-0000-0000-000000000000"
    assert m.envelope.sourceName == "Leo"
    assert m.envelope.sourceDevice == 1
    assert m.envelope.timestamp == 1734304012251
    assert isinstance(m.envelope.dataMessage, DeleteMessage)
    assert m.envelope.dataMessage.timestamp == 1734304012251
    assert m.envelope.dataMessage.message is None
    assert m.envelope.dataMessage.expiresInSeconds == 0
    assert m.envelope.dataMessage.viewOnce is False
    assert m.envelope.dataMessage.remoteDelete.timestamp == 1734300999933
    assert m.account == "+4900000000002"

class Reaction(BaseModel):
    emoji: str
    targetAuthor: str
    targetAuthorNumber: str
    targetAuthorUuid: str
    targetSentTimestamp: int
    isRemove: bool

class TypingMessageBase(BaseModel):
    timestamp: int
    groupId: str | None = None

class TypingStarted(TypingMessageBase):
    action: Literal["STARTED"]

class TypingStopped(TypingMessageBase):
    action: Literal["STOPPED"]

TypingMessage = TypingStarted | TypingStopped

class GroupInfo(BaseModel):
    groupId: str
    type: Literal["DELIVER"]

class BaseDataMessage(BaseModel):
    timestamp: int
    expiresInSeconds: int
    viewOnce: bool
    groupInfo: GroupInfo | None = None

class ReactionMessage(BaseDataMessage):
    reaction: Reaction
    message: Literal[None] = None

class DeleteEntry(BaseModel):
    timestamp: int

class DeleteMessage(BaseDataMessage):
    remoteDelete: DeleteEntry
    message: Literal[None] = None

class DataMessage(BaseDataMessage):
    message: str

class EditMessage(BaseModel):
    targetSentTimestamp: int
    dataMessage: DataMessage

class Envelope(BaseModel):
    source: str
    sourceNumber: str
    sourceUuid: str
    sourceName: str
    sourceDevice: int
    timestamp: int

class EnvelopeData(Envelope):
    dataMessage: DataMessage | ReactionMessage | DeleteMessage

class EnvelopeEdit(Envelope):
    editMessage: EditMessage

class EnvelopeTyping(Envelope):
    typingMessage: TypingMessage

class Message(BaseModel):
    envelope: EnvelopeData | EnvelopeEdit | EnvelopeTyping
    account: str