"""
    Query and control an archive of messages stored on a server using
    XEP-0313: Message Archive Management(MAM).
"""

from __future__ import annotations

import logging
from datetime import datetime, timedelta, timezone
from hashlib import md5
from typing import (
    AsyncIterable,
    List,
    Optional,
    TYPE_CHECKING,
    TypedDict,
)

from slixmpp import JID, Message as SMessage
from slixmpp.exceptions import IqError, IqTimeout
from poezio.theming import get_theme, TextStyle
from poezio import tabs
from poezio import colors
from poezio.common import to_utc
from poezio.ui.types import (
    Message,
)

if TYPE_CHECKING:
    from poezio.core.core import Core


log = logging.getLogger(__name__)


class DiscoInfoException(Exception):
    pass


class MAMQueryException(Exception):
    def __init__(self) -> None:
        super().__init__("Query interrupted")


class NoMAMSupportException(Exception):
    pass


def make_line(
        tab: tabs.ChatTab,
        text: str,
        time: datetime,
        jid: JID,
        identifier: str = '',
        nick: str = ''
) -> Message:
    """Adds a textual entry in the TextBuffer"""

    # Convert to local timezone
    time = time.replace(tzinfo=timezone.utc).astimezone(tz=None)
    time = time.replace(tzinfo=None)
    color: TextStyle

    if isinstance(tab, tabs.MucTab):
        nick = jid.resource
        user = tab.get_user_by_name(nick)
        if user:
            color = user.color
        else:
            theme = get_theme()
            if theme.ccg_palette:
                fg_color = colors.ccg_text_to_color(theme.ccg_palette, nick)
                color = fg_color, -1
            else:
                mod = len(theme.LIST_COLOR_NICKNAMES)
                nick_pos = int(md5(nick.encode('utf-8')).hexdigest(), 16) % mod
                color = theme.LIST_COLOR_NICKNAMES[nick_pos]
    else:
        if jid.bare == tab.core.xmpp.boundjid.bare:
            if not nick:
                nick = tab.core.own_nick
            color = get_theme().COLOR_OWN_NICK
        else:
            color = get_theme().COLOR_REMOTE_USER
            if not nick:
                nick = tab.get_nick()
    return Message(
        txt=text,
        identifier=identifier,
        time=time,
        nickname=nick,
        nick_color=color,
        history=True,
        user=None,
    )


class IterateArgs(TypedDict, total=False):
    reverse: bool
    rsm: dict[str, int]
    total: int
    start: str | None
    end: str | None
    jid: JID
    with_jid: JID


async def get_mam_iterator(
        core: 'Core',
        groupchat: bool,
        remote_jid: JID,
        amount: int,
        reverse: bool = True,
        start: Optional[str] = None,
        end: Optional[str] = None,
        before: Optional[str] = None,
) -> AsyncIterable[SMessage]:
    """Get an async iterator for this mam query"""
    try:
        query_jid = remote_jid if groupchat else JID(core.xmpp.boundjid.bare)
        iq = await core.xmpp.plugin['xep_0030'].get_info(jid=query_jid)
    except (IqError, IqTimeout) as exc:
        raise DiscoInfoException() from exc
    if 'urn:xmpp:mam:2' not in iq['disco_info'].get_features():
        raise NoMAMSupportException()

    args: IterateArgs = {
        'reverse': reverse,
        'rsm': {
            'max': 30,
        },
    }

    if groupchat:
        args['jid'] = remote_jid
    else:
        args['with_jid'] = remote_jid

    args['total'] = amount
    args['start'] = start
    args['end'] = end
    return core.xmpp['xep_0313'].iterate(**args)  # type: ignore


class MessageDict(TypedDict):
    time: datetime
    jid: JID
    text: str
    identifier: str


def _parse_message(msg: SMessage) -> MessageDict:
    """Parse info inside a MAM forwarded message"""
    forwarded = msg['mam_result']['forwarded']
    message = forwarded['stanza']
    result: MessageDict = {
        'time': forwarded['delay']['stamp'],
        'jid': message['from'],
        'text': message['body'],
        'identifier': message['id']
    }
    return result


def _ignore_private_message(stanza: SMessage, filter_jid: Optional[JID]) -> bool:
    """Returns True if a MUC-PM should be ignored, as prosody returns
    all PMs within the same room.
    """
    if filter_jid is None:
        return False
    sent = stanza['from'].bare != filter_jid.bare
    if sent and stanza['to'].full != filter_jid.full:
        return True
    elif not sent and stanza['from'].full != filter_jid.full:
        return True
    return False


async def retrieve_messages(tab: tabs.ChatTab,
                            results: AsyncIterable[SMessage],
                            amount: int = 100) -> List[Message]:
    """Run the MAM query and put messages in order"""
    msg_count = 0
    to_add: list[Message] = []
    tab_is_private = isinstance(tab, tabs.PrivateTab)
    filter_jid = None
    if tab_is_private:
        filter_jid = tab.jid
    try:
        async for msg in results:
            stanza = msg['mam_result']['forwarded']['stanza']
            if stanza.xml.find('{%s}%s' % ('jabber:client', 'body')) is not None:
                if _ignore_private_message(stanza, filter_jid):
                    continue
                args = _parse_message(msg)
                to_add.append(make_line(tab, **args))
                msg_count += 1
                if msg_count == amount:
                    to_add.reverse()
                    return to_add
    except (IqError, IqTimeout) as exc:
        log.debug('Unable to complete MAM query: %s', exc, exc_info=True)
        raise MAMQueryException() from exc
    to_add.reverse()
    return to_add


async def fetch_history(tab: tabs.ChatTab,
                        start: Optional[datetime] = None,
                        end: Optional[datetime] = None,
                        amount: int = 100) -> List[Message]:
    remote_jid = tab.jid
    if not end:
        for msg in tab._text_buffer.messages:
            if isinstance(msg, Message):
                end = msg.time
                end -= timedelta(microseconds=1)
                break
    if end is None:
        end = datetime.now() + timedelta(seconds=1)
    end = to_utc(end)
    end_str = datetime.strftime(end, '%Y-%m-%dT%H:%M:%SZ')

    start_str = None
    if start is not None:
        start = to_utc(start)
        start_str = datetime.strftime(start, '%Y-%m-%dT%H:%M:%SZ')

    mam_iterator = await get_mam_iterator(
        core=tab.core,
        groupchat=isinstance(tab, tabs.MucTab),
        remote_jid=remote_jid,
        amount=amount,
        end=end_str,
        start=start_str,
        reverse=True,
    )
    return await retrieve_messages(tab, mam_iterator, amount)
