import base64
import functools
import logging
import os
import shutil
import stat
import tempfile
from collections.abc import Collection, Sequence
from datetime import datetime
from itertools import chain
from mimetypes import guess_extension, guess_type
from pathlib import Path
from urllib.parse import quote as urlquote
from uuid import uuid4
from xml.etree import ElementTree as ET

import aiohttp
import thumbhash
from PIL import Image, ImageOps
from slixmpp import JID, Iq, Message
from slixmpp.plugins.xep_0264.stanza import Thumbnail
from slixmpp.plugins.xep_0447.stanza import StatelessFileSharing

from ...db.avatar import avatar_cache
from ...db.models import Attachment
from ...util.types import (
    LegacyAttachment,
    LegacyMessageType,
    LegacyThreadType,
    MessageReference,
)
from ...util.util import fix_suffix
from .. import config
from .message_text import TextMessageMixin


class AttachmentMixin(TextMessageMixin):
    PRIVILEGED_UPLOAD = False

    @property
    def __is_component(self) -> bool:
        return self.session is NotImplemented

    async def __upload(
        self,
        file_path: Path,
        file_name: str | None = None,
        content_type: str | None = None,
    ) -> str:
        assert config.UPLOAD_SERVICE

        file_size = file_path.stat().st_size

        content_type = content_type or self.xmpp.plugin["xep_0363"].default_content_type
        iq_slot = await self.__request_upload_slot(
            config.UPLOAD_SERVICE,
            file_name or file_path.name,
            file_size,
            content_type,
        )
        slot = iq_slot["http_upload_slot"]
        headers = {
            "Content-Length": str(file_size),
            "Content-Type": content_type,
            **{header["name"]: header["value"] for header in slot["put"]["headers"]},
        }

        async with aiohttp.ClientSession() as http:
            with file_path.open("rb") as fp:
                async with http.put(
                    slot["put"]["url"], data=fp, headers=headers
                ) as response:
                    response.raise_for_status()
        return slot["get"]["url"]

    async def __request_upload_slot(
        self,
        upload_service: JID | str,
        filename: str,
        size: int,
        content_type: str,
    ) -> Iq:
        iq_request = self.xmpp.make_iq_get(
            ito=upload_service, ifrom=config.UPLOAD_REQUESTER or self.xmpp.boundjid
        )
        request = iq_request["http_upload_request"]
        request["filename"] = filename
        request["size"] = str(size)
        request["content-type"] = content_type
        if self.__is_component or not self.PRIVILEGED_UPLOAD:
            return await iq_request.send()

        assert self.session is not NotImplemented
        iq_request.set_from(self.session.user_jid)
        return await self.xmpp["xep_0356"].send_privileged_iq(iq_request)

    @staticmethod
    async def __no_upload(
        file_path: Path,
        file_name: str | None = None,
        legacy_file_id: str | int | None = None,
    ) -> tuple[Path, str]:
        file_id = str(uuid4()) if legacy_file_id is None else str(legacy_file_id)
        assert config.NO_UPLOAD_PATH is not None
        assert config.NO_UPLOAD_URL_PREFIX is not None
        destination_dir = Path(config.NO_UPLOAD_PATH) / file_id

        if destination_dir.exists():
            log.debug("Dest dir exists: %s", destination_dir)
            files = list(f for f in destination_dir.glob("**/*") if f.is_file())
            if len(files) == 1:
                log.debug(
                    "Found the legacy attachment '%s' at '%s'",
                    legacy_file_id,
                    files[0],
                )
                name = files[0].name
                uu = files[0].parent.name  # anti-obvious url trick, see below
                return files[0], "/".join([file_id, uu, name])
            else:
                log.warning(
                    (
                        "There are several or zero files in %s, "
                        "slidge doesn't know which one to pick among %s. "
                        "Removing the dir."
                    ),
                    destination_dir,
                    files,
                )
                shutil.rmtree(destination_dir)

        log.debug("Did not find a file in: %s", destination_dir)
        # let's use a UUID to avoid URLs being too obvious
        uu = str(uuid4())
        destination_dir = destination_dir / uu
        destination_dir.mkdir(parents=True)

        name = file_name or file_path.name
        destination = destination_dir / name
        method = config.NO_UPLOAD_METHOD
        if method == "copy":
            shutil.copy2(file_path, destination)
        elif method == "hardlink":
            os.link(file_path, destination)
        elif method == "symlink":
            os.symlink(file_path, destination, target_is_directory=True)
        elif method == "move":
            shutil.move(file_path, destination)
        else:
            raise RuntimeError("No upload method not recognized", method)

        if config.NO_UPLOAD_FILE_READ_OTHERS:
            log.debug("Changing perms of %s", destination)
            destination.chmod(destination.stat().st_mode | stat.S_IROTH)
        uploaded_url = "/".join([file_id, uu, name])

        return destination, uploaded_url

    async def __valid_url(self, url: str) -> bool:
        async with self.session.http.head(url) as r:
            return r.status < 400

    async def __get_stored(self, attachment: LegacyAttachment) -> Attachment:
        if attachment.legacy_file_id is not None and not self.__is_component:
            with self.xmpp.store.session() as orm:
                stored = (
                    orm.query(Attachment)
                    .filter_by(
                        legacy_file_id=str(attachment.legacy_file_id),
                        user_account_id=self.session.user_pk,
                    )
                    .one_or_none()
                )
                if stored is not None:
                    if not await self.__valid_url(stored.url):
                        stored.url = None  # type:ignore
                    return stored
        return Attachment(
            user_account_id=None if self.__is_component else self.session.user_pk,
            legacy_file_id=None
            if attachment.legacy_file_id is None
            else str(attachment.legacy_file_id),
            url=attachment.url if config.USE_ATTACHMENT_ORIGINAL_URLS else None,
        )

    async def __get_url(
        self, attachment: LegacyAttachment, stored: Attachment
    ) -> tuple[bool, Path | None, str]:
        file_name = attachment.name
        content_type = attachment.content_type
        file_path = attachment.path

        if file_name and len(file_name) > config.ATTACHMENT_MAXIMUM_FILE_NAME_LENGTH:
            log.debug("Trimming long filename: %s", file_name)
            base, ext = os.path.splitext(file_name)
            file_name = (
                base[: config.ATTACHMENT_MAXIMUM_FILE_NAME_LENGTH - len(ext)] + ext
            )

        if file_path is None:
            if file_name is None:
                file_name = str(uuid4())
                if content_type is not None:
                    ext = guess_extension(content_type, strict=False)  # type:ignore
                    if ext is not None:
                        file_name += ext
            temp_dir = Path(tempfile.mkdtemp())
            file_path = temp_dir / file_name
            if attachment.url:
                async with self.session.http.get(attachment.url) as r:
                    r.raise_for_status()
                    with file_path.open("wb") as f:
                        f.write(await r.read())

            elif attachment.stream is not None:
                data = attachment.stream.read()
                if data is None:
                    raise RuntimeError

                with file_path.open("wb") as f:
                    f.write(data)
            elif attachment.aio_stream is not None:
                # TODO: patch slixmpp to allow this as data source for
                #       upload_file() so we don't even have to write anything
                #       to disk.
                with file_path.open("wb") as f:
                    async for chunk in attachment.aio_stream:
                        f.write(chunk)
            elif attachment.data is not None:
                with file_path.open("wb") as f:
                    f.write(attachment.data)

            is_temp = not bool(config.NO_UPLOAD_PATH)
        else:
            is_temp = False

        assert isinstance(file_path, Path)
        if config.FIX_FILENAME_SUFFIX_MIME_TYPE:
            file_name, content_type = fix_suffix(file_path, content_type, file_name)
            attachment.content_type = content_type
            attachment.name = file_name

        if config.NO_UPLOAD_PATH:
            local_path, new_url = await self.__no_upload(
                file_path, file_name, stored.legacy_file_id
            )
            new_url = (config.NO_UPLOAD_URL_PREFIX or "") + "/" + urlquote(new_url)
        else:
            local_path = file_path
            new_url = await self.__upload(file_path, file_name, content_type)
        if stored.legacy_file_id and new_url is not None:
            stored.url = new_url

        if local_path is not None and local_path.stat().st_size == 0:
            raise RuntimeError("File size is 0")

        return is_temp, local_path, new_url

    async def __set_sims(
        self,
        msg: Message,
        uploaded_url: str,
        path: Path | None,
        attachment: LegacyAttachment,
        stored: Attachment,
    ) -> Thumbnail | None:
        if stored.sims is not None:
            ref = self.xmpp["xep_0372"].stanza.Reference(xml=ET.fromstring(stored.sims))
            msg.append(ref)
            if ref["sims"]["file"].get_plugin("thumbnail", check=True):
                return ref["sims"]["file"]["thumbnail"]
            else:
                return None

        if not path:
            return None

        ref = self.xmpp["xep_0385"].get_sims(
            path, [uploaded_url], attachment.content_type, attachment.caption
        )
        if attachment.name:
            ref["sims"]["file"]["name"] = attachment.name
        thumbnail = None
        if attachment.content_type is not None and attachment.content_type.startswith(
            "image"
        ):
            try:
                h, x, y = await self.xmpp.loop.run_in_executor(
                    avatar_cache._thread_pool, get_thumbhash, path
                )
            except Exception as e:
                log.debug("Could not generate a thumbhash", exc_info=e)
            else:
                thumbnail = ref["sims"]["file"]["thumbnail"]
                thumbnail["width"] = x
                thumbnail["height"] = y
                thumbnail["media-type"] = "image/thumbhash"
                thumbnail["uri"] = "data:image/thumbhash;base64," + urlquote(h)

        stored.sims = str(ref)
        msg.append(ref)

        return thumbnail

    def __set_sfs(
        self,
        msg: Message,
        uploaded_url: str,
        path: Path | None,
        attachment: LegacyAttachment,
        stored: Attachment,
        thumbnail: Thumbnail | None = None,
    ) -> None:
        if stored.sfs is not None:
            msg.append(StatelessFileSharing(xml=ET.fromstring(stored.sfs)))
            return

        if not path:
            return

        sfs = self.xmpp["xep_0447"].get_sfs(
            path, [uploaded_url], attachment.content_type, attachment.caption
        )
        if attachment.name:
            sfs["file"]["name"] = attachment.name
        if attachment.disposition:
            sfs["disposition"] = attachment.disposition
        else:
            del sfs["disposition"]
        if thumbnail is not None:
            sfs["file"].append(thumbnail)
        stored.sfs = str(sfs)
        msg.append(sfs)

    async def __set_sfs_and_sims_without_download(
        self, msg: Message, attachment: LegacyAttachment
    ) -> None:
        assert attachment.url is not None

        if not any(
            (
                attachment.content_type,
                attachment.name,
                attachment.disposition,
            )
        ):
            return

        sims = self.xmpp.plugin["xep_0385"].stanza.Sims()
        ref = self.xmpp["xep_0372"].stanza.Reference()

        ref["uri"] = attachment.url
        ref["type"] = "data"
        sims["sources"].append(ref)
        sims.enable("file")

        xep_0447_stanza = self.xmpp.plugin["xep_0447"].stanza
        sfs = xep_0447_stanza.StatelessFileSharing()
        url_data = xep_0447_stanza.UrlData()
        url_data["target"] = attachment.url
        sfs["sources"].append(url_data)
        sfs.enable("file")

        if attachment.content_type:
            sims["file"]["media-type"] = attachment.content_type
            sfs["file"]["media-type"] = attachment.content_type
        if attachment.caption:
            sims["file"]["desc"] = attachment.caption
            sfs["file"]["desc"] = attachment.caption
        if attachment.name:
            sims["file"]["name"] = attachment.name
            sfs["file"]["name"] = attachment.name
        if attachment.disposition:
            sfs["disposition"] = attachment.disposition

        msg.append(sims)
        msg.append(sfs)

    def __send_url(
        self,
        msg: Message,
        legacy_msg_id: LegacyMessageType,
        uploaded_url: str,
        caption: str | None = None,
        carbon: bool = False,
        when: datetime | None = None,
        correction: bool = False,
        **kwargs,
    ) -> list[Message]:
        msg["oob"]["url"] = uploaded_url
        msg["body"] = uploaded_url
        if msg.get_plugin("sfs", check=True):
            msg["fallback"].enable("body")
            msg["fallback"]["for"] = self.xmpp.plugin["xep_0447"].stanza.NAMESPACE
        if caption:
            if correction:
                msg["replace"]["id"] = self._replace_id(legacy_msg_id)
            else:
                self._set_msg_id(msg, legacy_msg_id)
            m1 = self._send(msg, carbon=carbon, correction=correction, **kwargs)
            m2 = self.send_text(
                caption, legacy_msg_id=None, when=when, carbon=carbon, **kwargs
            )
            return [m1, m2] if m2 else [m1]
        else:
            if correction:
                msg["replace"]["id"] = self._replace_id(legacy_msg_id)
            else:
                self._set_msg_id(msg, legacy_msg_id)
            return [self._send(msg, carbon=carbon, **kwargs)]

    def __get_base_message(
        self,
        legacy_msg_id: LegacyMessageType | None = None,
        reply_to: MessageReference | None = None,
        when: datetime | None = None,
        thread: LegacyThreadType | None = None,
        carbon: bool = False,
        correction: bool = False,
        mto: JID | None = None,
    ) -> Message:
        if correction:
            xmpp_ids = self._legacy_to_xmpp(legacy_msg_id)
            if xmpp_ids:
                original_xmpp_id = xmpp_ids[0]
                for xmpp_id in xmpp_ids:
                    if xmpp_id == original_xmpp_id:
                        continue
                    self.retract(xmpp_id, thread)

        if reply_to is not None and reply_to.body:
            # We cannot have a "quote fallback" for attachments since most (all?)
            # XMPP clients will only treat a message as an attachment if the
            # body is the URL and nothing else.
            reply_to_for_attachment: MessageReference | None = MessageReference(
                reply_to.legacy_id, reply_to.author
            )
        else:
            reply_to_for_attachment = reply_to

        return self._make_message(
            when=when,
            reply_to=reply_to_for_attachment,
            carbon=carbon,
            mto=mto,
            thread=thread,
        )

    async def send_file(
        self,
        attachment: LegacyAttachment | Path | str,
        legacy_msg_id: LegacyMessageType | None = None,
        *,
        reply_to: MessageReference | None = None,
        when: datetime | None = None,
        thread: LegacyThreadType | None = None,
        **kwargs,
    ) -> tuple[str | None, list[Message]]:
        """
        Send a single file from this :term:`XMPP Entity`.

        :param attachment: The file to send.
            Ideally, a :class:`.LegacyAttachment` with a unique ``legacy_file_id``
            attribute set, to optimise potential future reuses.
            It can also be:
            - a :class:`pathlib.Path` instance to point to a local file, or
            - a ``str``, representing a fetchable HTTP URL.
        :param legacy_msg_id: If you want to be able to transport read markers from the gateway
            user to the legacy network, specify this
        :param reply_to: Quote another message (:xep:`0461`)
        :param when: when the file was sent, for a "delay" tag (:xep:`0203`)
        :param thread:
        """
        coro = self.__send_file(
            attachment,
            legacy_msg_id,
            reply_to=reply_to,
            when=when,
            thread=thread,
            **kwargs,
        )
        if self.__is_component:
            return await coro
        elif not isinstance(attachment, LegacyAttachment):
            return await coro
        elif attachment.legacy_file_id is None:
            return await coro
        else:
            # prevents race conditions where we download the same thing several time
            # and end up attempting to insert it twice in the DB, raising an
            # IntegrityError.
            async with self.session.lock(("attachment", attachment.legacy_file_id)):
                return await coro

    async def __send_file(
        self,
        attachment: LegacyAttachment | Path | str,
        legacy_msg_id: LegacyMessageType | None = None,
        *,
        reply_to: MessageReference | None = None,
        when: datetime | None = None,
        thread: LegacyThreadType | None = None,
        **kwargs,
    ) -> tuple[str | None, list[Message]]:
        store_multi = kwargs.pop("store_multi", True)
        carbon = kwargs.pop("carbon", False)
        mto = kwargs.pop("mto", None)
        correction = kwargs.get("correction", False)

        msg = self.__get_base_message(
            legacy_msg_id, reply_to, when, thread, carbon, correction, mto
        )

        if isinstance(attachment, str):
            attachment = LegacyAttachment(url=attachment)
        elif isinstance(attachment, Path):
            attachment = LegacyAttachment(path=attachment)

        stored = await self.__get_stored(attachment)

        if attachment.content_type is None and (
            name := (attachment.name or attachment.url or attachment.path)
        ):
            attachment.content_type, _ = guess_type(name)

        if stored.url:
            is_temp = False
            local_path = None
            new_url = stored.url
        else:
            try:
                is_temp, local_path, new_url = await self.__get_url(attachment, stored)
            except Exception as e:
                log.error("Error with attachment: %s: %s", attachment, e)
                log.debug("", exc_info=e)
                msg["body"] = (
                    f"/me tried to send a file ({attachment.format_for_user()}), "
                    f"but something went wrong: {e}. "
                )
                self._set_msg_id(msg, legacy_msg_id)
                return None, [self._send(msg, **kwargs)]
        assert new_url is not None

        stored.url = new_url
        if config.USE_ATTACHMENT_ORIGINAL_URLS and attachment.url:
            await self.__set_sfs_and_sims_without_download(msg, attachment)
        else:
            thumbnail = await self.__set_sims(
                msg, new_url, local_path, attachment, stored
            )
            self.__set_sfs(msg, new_url, local_path, attachment, stored, thumbnail)

        if not self.__is_component:
            with self.xmpp.store.session(expire_on_commit=False) as orm:
                orm.add(stored)
                orm.commit()

        if is_temp and isinstance(local_path, Path):
            local_path.unlink()
            local_path.parent.rmdir()

        msgs = self.__send_url(
            msg, legacy_msg_id, new_url, attachment.caption, carbon, when, **kwargs
        )
        if not self.__is_component:
            if store_multi:
                self.__store_multi(legacy_msg_id, msgs)
        return new_url, msgs

    def __send_body(
        self,
        body: str | None = None,
        legacy_msg_id: LegacyMessageType | None = None,
        reply_to: MessageReference | None = None,
        when: datetime | None = None,
        thread: LegacyThreadType | None = None,
        **kwargs,
    ) -> Message | None:
        if body:
            return self.send_text(
                body,
                legacy_msg_id,
                reply_to=reply_to,
                when=when,
                thread=thread,
                **kwargs,
            )
        else:
            return None

    async def send_files(
        self,
        attachments: Collection[LegacyAttachment],
        legacy_msg_id: LegacyMessageType | None = None,
        body: str | None = None,
        *,
        reply_to: MessageReference | None = None,
        when: datetime | None = None,
        thread: LegacyThreadType | None = None,
        body_first: bool = False,
        correction: bool = False,
        correction_event_id: LegacyMessageType | None = None,
        **kwargs,
    ) -> None:
        # TODO: once the epic XEP-0385 vs XEP-0447 battle is over, pick
        #       one and stop sending several attachments this way
        # we attach the legacy_message ID to the last message we send, because
        # we don't want several messages with the same ID (especially for MUC MAM)
        if not attachments and not body:
            # ignoring empty message
            return
        body_msg_id = (
            legacy_msg_id if body_needs_msg_id(attachments, body, body_first) else None
        )
        send_body = functools.partial(
            self.__send_body,
            body=body,
            reply_to=reply_to,
            when=when,
            thread=thread,
            correction=correction,
            legacy_msg_id=body_msg_id,
            correction_event_id=correction_event_id,
            **kwargs,
        )
        all_msgs = []
        if body_first:
            all_msgs.append(send_body())
        for i, attachment in enumerate(attachments):
            if i == 0 and body_msg_id is None:
                legacy = legacy_msg_id
            else:
                legacy = None
            _url, msgs = await self.send_file(
                attachment,
                legacy,
                reply_to=reply_to,
                when=when,
                thread=thread,
                store_multi=False,
                **kwargs,
            )
            all_msgs.extend(msgs)
        if not body_first:
            all_msgs.append(send_body())
        self.__store_multi(legacy_msg_id, all_msgs)

    def __store_multi(
        self,
        legacy_msg_id: LegacyMessageType | None,
        all_msgs: Sequence[Message | None],
    ) -> None:
        if legacy_msg_id is None:
            return
        ids = []
        for msg in all_msgs:
            if not msg:
                continue
            if stanza_id := msg.get_plugin("stanza_id", check=True):
                ids.append(stanza_id["id"])
            else:
                ids.append(msg.get_id())
        with self.xmpp.store.session() as orm:
            self.xmpp.store.id_map.set_msg(
                orm, self._recipient_pk(), str(legacy_msg_id), ids, self.is_participant
            )
            orm.commit()


def body_needs_msg_id(
    attachments: Collection[LegacyAttachment], body: str | None, body_first: bool
) -> bool:
    if attachments:
        return bool(body and body_first)
    else:
        return True


def get_thumbhash(path: Path) -> tuple[str, int, int]:
    with path.open("rb") as fp:
        img = Image.open(fp)
        width, height = img.size
        img = img.convert("RGBA")
        if width > 100 or height > 100:
            img.thumbnail((100, 100))
    img = ImageOps.exif_transpose(img)
    rgba_2d = list(img.getdata())
    rgba = list(chain(*rgba_2d))
    ints = thumbhash.rgba_to_thumb_hash(img.width, img.height, rgba)
    return base64.b64encode(bytes(ints)).decode(), width, height


log = logging.getLogger(__name__)
