# SPDX-FileCopyrightText: 2022 Chris Heywood
#
# SPDX-License-Identifier: GPL-3.0-or-later

import os
import shutil
import unittest
import warnings

from iotas.attachment import Attachment
from iotas.attachment_helpers import (
    copy_note_attachments,
    write_attachment_to_disk,
    get_attachments_on_disk_for_note,
    delete_attachments_for_note,
    get_attachment_disk_states,
    get_attachment_filesystem_uri,
    get_attachment_filesystem_path,
    attachment_exists_locally,
    get_attachment_file_object,
    flush_all_attachments_on_disk,
    set_attachment_filesystem_uris_on_tokens,
    trim_orphaned_images,
    AttachmentsCopyOutcome,
)
from iotas.markdown_helpers import (
    filter_image_tokens,
    parse_to_tokens,
)
from iotas.note import Note

import iotas.attachment_helpers


def get_attachments_dir() -> str:
    file_dir = os.path.dirname(__file__)
    return os.path.join(file_dir, os.pardir, "testing-tmp", "user-attachments")


iotas.attachment_helpers.get_attachments_dir = get_attachments_dir

warnings.filterwarnings("ignore", "version")


class Test(unittest.TestCase):

    def test_copy_note_attachments(self) -> None:
        self.__prepare_output_dir()
        note = Note()
        note.id = 1
        note.content = "![](image.png)\n\n![](image2.png)\n\n![](image2.png)"
        path = os.path.join(get_attachments_dir(), "1.image.png")
        self.assertTrue(self.__create_empty_file(path))
        path = os.path.join(get_attachments_dir(), "1.image2.png")
        self.assertTrue(self.__create_empty_file(path))

        out_dir = os.path.join(get_attachments_dir(), "copy-test")

        def refresh_out_dir():
            if os.path.exists(out_dir):
                shutil.rmtree(out_dir)
            os.makedirs(out_dir)

        refresh_out_dir()
        results = copy_note_attachments(note, out_dir, prefix_note_id=True)
        self.assertEqual(results.outcome, AttachmentsCopyOutcome.SUCCESS)
        self.assertEqual(len(os.listdir(out_dir)), 2)

        refresh_out_dir()
        note.content = "![](image.png)\n\n![](image2.png)\n\n![](image3.png)"
        results = copy_note_attachments(note, out_dir, prefix_note_id=True)
        self.assertEqual(results.outcome, AttachmentsCopyOutcome.HAD_MISSING)
        self.assertEqual(len(os.listdir(out_dir)), 2)

        refresh_out_dir()
        note.content = "[](image.png)"
        results = copy_note_attachments(note, out_dir, prefix_note_id=True)
        self.assertEqual(results.outcome, AttachmentsCopyOutcome.NONE)
        self.assertEqual(len(os.listdir(out_dir)), 0)

        self.__clean_output_dir()

    def test_write_attachment_to_disk(self) -> None:
        self.__prepare_output_dir()
        attachment = Attachment()
        attachment.note_id = 1
        attachment.path = "image.png"

        success = write_attachment_to_disk(attachment, b"bytes")
        self.assertTrue(success)
        path = os.path.join(get_attachments_dir(), "1.image.png")
        self.assertTrue(os.path.exists(path))
        with open(path, "rb") as f:
            out = f.read()
        self.assertEqual(out, b"bytes")

        self.__clean_output_dir()

    def test_get_attachments_on_disk_for_note(self) -> None:
        self.__prepare_output_dir()
        note = Note()
        note.id = 1
        self.assertEqual(len(get_attachments_on_disk_for_note(note)), 0)

        path = os.path.join(get_attachments_dir(), "1.image.png")
        self.assertTrue(self.__create_empty_file(path))
        path = os.path.join(get_attachments_dir(), "1.image2.png")
        self.assertTrue(self.__create_empty_file(path))
        path = os.path.join(get_attachments_dir(), "2.image.png")
        self.assertTrue(self.__create_empty_file(path))

        out = get_attachments_on_disk_for_note(note)
        self.assertEqual(len(out), 2)
        out_paths = {a.path for a in out}
        self.assertEqual(out_paths, {"image.png", "image2.png"})

        self.__clean_output_dir()

    def test_delete_attachments_for_note(self) -> None:
        self.__prepare_output_dir()
        note = Note()
        note.id = 1
        self.assertEqual(len(get_attachments_on_disk_for_note(note)), 0)

        path = os.path.join(get_attachments_dir(), "1.image.png")
        self.assertTrue(self.__create_empty_file(path))
        path = os.path.join(get_attachments_dir(), "1.image2.png")
        self.assertTrue(self.__create_empty_file(path))
        path = os.path.join(get_attachments_dir(), "2.image.png")
        self.assertTrue(self.__create_empty_file(path))

        delete_attachments_for_note(note)
        remaining = os.listdir(get_attachments_dir())
        self.assertEqual(len(remaining), 1)
        self.assertEqual(remaining[0], "2.image.png")

        self.__clean_output_dir()

    def test_get_attachment_disk_states(self) -> None:
        self.__prepare_output_dir()

        note = Note()
        note.id = 1
        note.content = "![](image.png)\n\n![](image2.png)\n\n![](image2.png)\n\n![](image3.png)"

        path = os.path.join(get_attachments_dir(), "1.image.png")
        self.assertTrue(self.__create_empty_file(path))
        path = os.path.join(get_attachments_dir(), "1.image2.png")
        self.assertTrue(self.__create_empty_file(path))

        _parser, tokens = parse_to_tokens(note, exporting=False, tex_support=False)
        states = get_attachment_disk_states(note, tokens)
        self.assertEqual(len(states.exists), 2)
        self.assertTrue("image.png" in states.exists)
        self.assertTrue("image2.png" in states.exists)
        self.assertEqual(len(states.missing), 1)
        self.assertTrue("image3.png" in states.missing)
        self.__clean_output_dir()

    def test_get_attachment_filesystem_uri(self) -> None:
        attachment = Attachment()
        attachment.note_id = 1
        attachment.path = "image.png"
        uri = get_attachment_filesystem_uri(attachment)
        expected = os.path.join(get_attachments_dir(), "1.image.png")
        self.assertEqual(uri, f"file://{expected}")

    def test_get_attachment_filesystem_path(self) -> None:
        attachment = Attachment()
        attachment.note_id = 1
        attachment.path = "image.png"
        uri = get_attachment_filesystem_path(attachment)
        expected = os.path.join(get_attachments_dir(), "1.image.png")
        self.assertEqual(uri, expected)

    def test_attachment_exists_locally(self) -> None:
        attachment = Attachment()
        attachment.note_id = 1
        attachment.path = "image.png"
        self.assertFalse(attachment_exists_locally(attachment))

        self.__prepare_output_dir()
        path = os.path.join(get_attachments_dir(), "1.image.png")
        self.assertTrue(self.__create_empty_file(path))
        self.assertTrue(attachment_exists_locally(attachment))
        self.__clean_output_dir()

    def test_get_attachment_file_object(self) -> None:
        attachment = Attachment()
        attachment.note_id = 1
        attachment.path = "image.png"
        self.__prepare_output_dir()
        path = os.path.join(get_attachments_dir(), "1.image.png")
        self.assertTrue(self.__create_empty_file(path))

        obj = get_attachment_file_object(attachment)
        self.assertIsNotNone(obj)
        if obj is not None:  # mypy
            self.assertEqual(obj.read(), b"")
        self.__clean_output_dir()

    def test_flush_all_attachments_on_disk(self) -> None:
        self.__prepare_output_dir()
        path = os.path.join(get_attachments_dir(), "1.image.png")
        self.assertTrue(self.__create_empty_file(path))
        path = os.path.join(get_attachments_dir(), "1.image2.png")
        self.assertTrue(self.__create_empty_file(path))
        flush_all_attachments_on_disk()
        self.assertEqual(len(os.listdir(get_attachments_dir())), 0)
        self.__clean_output_dir()

    def test_set_attachment_filesystem_uris_to_tokens(self) -> None:
        note = Note()
        note.id = 1
        note.content = "![](image.png)\n\n![](image2.png)\n\n![](image2.png)"
        _parser, tokens = parse_to_tokens(note, exporting=False, tex_support=False)

        attachment1 = Attachment()
        attachment1.note_id = 1
        attachment1.path = "image.png"
        attachment2 = Attachment()
        attachment2.note_id = 1
        attachment2.path = "image2.png"
        new_paths = {
            "image.png": attachment1,
            "image2.png": attachment2,
        }
        set_attachment_filesystem_uris_on_tokens(tokens, new_paths)
        image_tokens = filter_image_tokens(tokens)
        self.assertEqual(len(image_tokens), 3)

        dirpath = f"file://{os.path.join(get_attachments_dir())}/"
        path = dirpath + "1.image.png"
        self.assertEqual(image_tokens[0].attrs["src"], path)
        path = dirpath + "1.image2.png"
        self.assertEqual(image_tokens[1].attrs["src"], path)
        self.assertEqual(image_tokens[2].attrs["src"], path)

    def test_trim_orphaned_images(self) -> None:
        self.__prepare_output_dir()
        path = os.path.join(get_attachments_dir(), "1.image.png")
        self.assertTrue(self.__create_empty_file(path))
        path = os.path.join(get_attachments_dir(), "1.orphan1.png")
        self.assertTrue(self.__create_empty_file(path))
        path = os.path.join(get_attachments_dir(), "1.orphan2.png")
        self.assertTrue(self.__create_empty_file(path))
        path = os.path.join(get_attachments_dir(), "2.image.png")
        self.assertTrue(self.__create_empty_file(path))
        self.assertEqual(len(os.listdir(get_attachments_dir())), 4)

        note = Note()
        note.id = 1
        note.content = "![description](image.png)"
        trim_orphaned_images(note, on_thread=False)
        self.assertEqual(len(os.listdir(get_attachments_dir())), 2)

    def __prepare_output_dir(self) -> str:
        out_path = get_attachments_dir()
        if os.path.exists(out_path):
            shutil.rmtree(out_path)
        os.makedirs(out_path)
        return out_path

    def __clean_output_dir(self) -> None:
        out_path = get_attachments_dir()
        if os.path.exists(out_path):
            shutil.rmtree(out_path)

    def __create_empty_file(self, path: str) -> bool:
        try:
            with open(path, "w") as f:
                f.write("")
        except OSError:
            return False
        else:
            return True
