from snek.system.service import BaseService


class DriveService(BaseService):

    mapper_name = "drive"

    EXTENSIONS_PICTURES = ["jpg", "jpeg", "png", "gif", "svg", "webp", "tiff"]
    EXTENSIONS_VIDEOS = [
        "mp4",
        "m4v",
        "mov",
        "wmv",
        "webm",
        "mkv",
        "mpg",
        "mpeg",
        "avi",
        "ogv",
        "ogg",
        "flv",
        "3gp",
        "3g2",
    ]
    EXTENSIONS_ARCHIVES = [
        "zip",
        "rar",
        "7z",
        "tar",
        "tar.gz",
        "tar.xz",
        "tar.bz2",
        "tar.lzma",
        "tar.lz",
    ]
    EXTENSIONS_AUDIO = [
        "mp3",
        "wav",
        "ogg",
        "flac",
        "m4a",
        "wma",
        "aac",
        "opus",
        "aiff",
        "au",
        "mid",
        "midi",
    ]
    EXTENSIONS_DOCS = [
        "pdf",
        "doc",
        "docx",
        "xls",
        "xlsx",
        "ppt",
        "pptx",
        "txt",
        "md",
        "json",
        "csv",
        "xml",
        "html",
        "css",
        "js",
        "py",
        "sql",
        "rs",
        "toml",
        "yml",
        "yaml",
        "ini",
        "conf",
        "config",
        "log",
        "csv",
        "tsv",
        "java",
        "cs",
        "csproj",
        "scss",
        "less",
        "sass",
        "json",
        "lock",
        "lock.json",
        "jsonl",
    ]

    async def get_drive_name_by_extension(self, extension):
        if extension.startswith("."):
            extension = extension[1:]
        if extension in self.EXTENSIONS_PICTURES:
            return "Pictures"
        if extension in self.EXTENSIONS_VIDEOS:
            return "Videos"
        if extension in self.EXTENSIONS_ARCHIVES:
            return "Archives"
        if extension in self.EXTENSIONS_AUDIO:
            return "Audio"
        if extension in self.EXTENSIONS_DOCS:
            return "Documents"
        return "My Drive"

    async def get_drive_by_extension(self, user_uid, extension):
        name = await self.get_drive_name_by_extension(extension)
        return await self.get_or_create(user_uid=user_uid, name=name)

    async def get_by_user(self, user_uid, name=None):
        kwargs = {"user_uid": user_uid}
        async for model in self.find(**kwargs):
            if not name:
                yield model
            elif model["name"] == name:
                yield model
            elif not model["name"] and name == "My Drive":
                model["name"] = "My Drive"
                await self.save(model)
                yield model

    async def get_or_create(self, user_uid, name=None, extensions=None):
        kwargs = {"user_uid": user_uid}
        if name:
            kwargs["name"] = name
        async for model in self.get_by_user(**kwargs):
            return model

        model = await self.new()
        model["user_uid"] = user_uid
        model["name"] = name
        await self.save(model)
        return model

    async def prepare_default_drives(self):
        async for drive_item in self.services.drive_item.find():
            extension = drive_item.extension
            drive = await self.get_drive_by_extension(drive_item["user_uid"], extension)
            if not drive_item["drive_uid"] == drive["uid"]:
                drive_item["drive_uid"] = drive["uid"]
                await self.services.drive_item.save(drive_item)

    async def prepare_default_drives_for_user(self, user_uid):
        await self.get_or_create(user_uid=user_uid, name="My Drive")
        await self.get_or_create(user_uid=user_uid, name="Shared Drive")
        await self.get_or_create(user_uid=user_uid, name="Pictures")
        await self.get_or_create(user_uid=user_uid, name="Videos")
        await self.get_or_create(user_uid=user_uid, name="Archives")
        await self.get_or_create(user_uid=user_uid, name="Documents")

    async def prepare_all(self):
        await self.prepare_default_drives()
        async for user in self.services.user.find():
            await self.prepare_default_drives_for_user(user["uid"])
