# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/google/GDrive.ipynb.

# %% auto 0
__all__ = ['DEFAULT_GDRIVE_FIELDS', 'DEFAULT_GFOLDER_FIELDS', 'get_file_by_id', 'get_folder_contents_by_id',
           'walk_folder_and_flatten', 'walk_folder', 'GDrive_MimeType_Enum', 'GDrive_ExportError', 'export_file_as',
           'export_as_pptx', 'export_as_zip', 'ExportAs_Enum', 'GDrive_ServiceRequired', 'GDrive_File', 'GDrive_Folder']

# %% ../../nbs/google/GDrive.ipynb 2
from functools import partial

from dataclasses import dataclass, field
from enum import Enum
from typing import List, Any, Optional, Callable

import datetime as dt

from googleapiclient.errors import HttpError

from domolibrary_extensions.utils.utils import (
    upsert_folder,
    convert_str_to_date,
    download_zip,
    download_pptx,
)

import domolibrary_extensions.utils.execution as ce
import domolibrary_extensions.google.GAuth as ga

# %% ../../nbs/google/GDrive.ipynb 5
DEFAULT_GDRIVE_FIELDS = [
    "id",
    "webViewLink",
    "name",
    "mimeType",
    " modifiedTime",
    "parents",
]


async def get_file_by_id(
    service: ga.Resource,
    file_id: str,
    fields: List[
        str
    ] = None,  # fields to retrieve from the GDrive API (see Google Documentation for additional options)
    parent_ls: List[str] = None,
):
    parent_ls = parent_ls or []

    fields = fields or DEFAULT_GDRIVE_FIELDS

    file = service.files().get(fileId=file_id, fields=", ".join(fields)).execute()

    parent_ls += [
        parent for parent in file.get("parents", []) if parent not in parent_ls
    ]

    file.update({"parent_ls": parent_ls})

    return file

# %% ../../nbs/google/GDrive.ipynb 7
DEFAULT_GFOLDER_FIELDS = [
    "nextPageToken",
    "files(id,webViewLink, name, mimeType,modifiedTime, parents)",
]


async def get_folder_contents_by_id(
    service: ga.Resource,
    folder_id: str,
    fields: List[str] = None,
    parent_ls: List[str] = None,  # used by recursive query to build path to file
    just_folders: bool = False,
    return_raw: bool = False,
):

    fields = fields or DEFAULT_GFOLDER_FIELDS

    res = None
    file_ls = []
    page_token = None

    query = f"'{folder_id}' in parents"

    if just_folders:

        query += " and mimeType = 'application/vnd.google-apps.folder'"

    while True:

        res = (
            service.files()
            .list(
                q=query,
                pageSize=10,
                fields=",".join(fields),
                pageToken=page_token,
            )
            .execute()
        )

        file_ls.extend(res.get("files", []))

        page_token = res.get("nextPageToken")

        if not page_token:
            break

    if return_raw:
        res["files"] = file_ls
        return res

    parent_ls = parent_ls.copy() if parent_ls else []

    parent_ls = list(set(parent_ls + [folder_id]))

    for file in file_ls:

        file.update(
            {
                "parent_ls": parent_ls,
                "is_folder": file.get("mimeType")
                == "application/vnd.google-apps.folder",
            }
        )

    return file_ls

# %% ../../nbs/google/GDrive.ipynb 9
async def walk_folder_and_flatten(
    service: ga.Resource,
    folder_id: str,
    file_ls: List[dict] = None,
    is_recursive: bool = True,
    parent_ls: Optional[List[str]] = None,
    just_folders: bool = False,
):  # recursive function that retrieves all files in subfolders

    if file_ls is None:
        file_ls = []

    parent_ls = parent_ls.copy() if parent_ls else [folder_id]

    new_files = await get_folder_contents_by_id(
        service, folder_id, parent_ls=parent_ls, just_folders=just_folders
    )
    file_ls.extend(new_files)

    if new_files and is_recursive:
        await ce.gather_with_concurrency(
            *[
                walk_folder_and_flatten(
                    service=service,
                    folder_id=file_obj["id"],
                    file_ls=file_ls,
                    parent_ls=[
                        *parent_ls,
                        *[
                            parent
                            for parent in file_obj["parent_ls"]
                            if parent not in parent_ls
                        ],
                    ],
                    just_folders=just_folders,
                )
                for file_obj in new_files
                if file_obj.get("mimeType") == "application/vnd.google-apps.folder"
            ],
            n=5
        )

    return file_ls

# %% ../../nbs/google/GDrive.ipynb 11
async def walk_folder(
    service: ga.Resource,
    folder_id: str,
    folder_obj: dict = None,
    parent_ls: List[str] = None,
    is_recursive: bool = True,
):  # recursive function that retrieves all files in subfolders

    parent_ls = parent_ls.copy() if parent_ls else [folder_id]

    folder_obj = folder_obj or (
        await get_file_by_id(service=service, file_id=folder_id, parent_ls=parent_ls)
    )

    folder_obj["children"] = (
        await get_folder_contents_by_id(
            service, folder_id, just_folders=False, parent_ls=parent_ls
        )
    ) or []

    if not is_recursive:
        return folder_obj

    for child in folder_obj["children"]:
        if child and child.get("mimeType") == "application/vnd.google-apps.folder":
            await walk_folder(
                service=service,
                folder_obj=child,
                folder_id=child.get("id"),
                parent_ls=parent_ls,
                is_recursive=is_recursive,
            )

    return folder_obj

# %% ../../nbs/google/GDrive.ipynb 13
class GDrive_MimeType_Enum(Enum):
    """for translating Google Drive types into 'exports as'"""

    docx = "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
    pptx = "application/vnd.openxmlformats-officedocument.presentationml.presentation"

    odt = "application/vnd.oasis.opendocument.text"
    rtf = "application/rtf"
    pdf = "application/pdf"
    txt = "text/plain"

    zip_file = "application/zip"
    epub = "application/epub+zip"

    folder = "application/vnd.google-apps.folder"
    gdoc = "application/vnd.google-apps.document"
    gsheet = "application/vnd.google-apps.spreadsheet"
    gshortcut = "application/vnd.google-apps.shortcut"
    gslides = "application/vnd.google-apps.presentation"

    markdown = "text/markdown"

    # html = 'application/zip'
    # 'tex'  = 'application/zip'
    # 'html.zip'= 'application/zip'

# %% ../../nbs/google/GDrive.ipynb 14
class GDrive_ExportError(Exception):
    def __init__(self, doc_url):
        message = f"failure to download content for {doc_url}"
        super().__init__(message)


async def export_file_as(
    service: ga.Resource,
    file_id: str,
    export_mime_type: GDrive_MimeType_Enum,
    replace_folder: bool = True,
    export_fn: Callable = None,
    output_path: str = None,
):

    retry = 0
    content = None
    while retry < 3 and content is None:
        retry += 1
        try:
            content = (
                service.files()
                .export(fileId=file_id, mimeType=export_mime_type.value)
                .execute()
            )

        except HttpError as err:
            print(err)

        except Exception as e:
            print(e)

    if not content:
        raise GDrive_ExportError(file_id)

    if not export_fn:
        return content

    upsert_folder(output_path, replace_folder=replace_folder)

    return export_fn(content=content, output_folder=output_path)


async def export_as_pptx(service: ga.Resource, file_id: str, output_path: str):
    return await export_file_as(
        service=service,
        file_id=file_id,
        export_mime_type=GDrive_MimeType_Enum.pptx,
        export_fn=download_pptx,
        output_path=output_path,
    )


async def export_as_zip(service: ga.Resource, file_id: str, output_path: str):
    return await export_file_as(
        service=service,
        file_id=file_id,
        export_mime_type=GDrive_MimeType_Enum.zip_file,
        export_fn=download_zip,
        output_path=output_path,
    )


class ExportAs_Enum(Enum):
    pptx = partial(export_as_pptx)
    zip_file = partial(export_as_zip)

    @classmethod
    def _missing_(cls, value):
        value = value.lower()
        try:
            for member in cls:
                if member.name == value:
                    return member

        except:
            return cls.zip_file

        return cls.zip_file

# %% ../../nbs/google/GDrive.ipynb 18
class GDrive_ServiceRequired(Exception):
    def __init__(self, doc_url=None, doc_id=None):
        message = f"service (googleclientapi.discovery.Resource) required to download {doc_url or doc_id}"
        super().__init__(message)


@dataclass
class GDrive_File:
    id: str
    name: str

    modified_time: dt.datetime
    url: str

    auth: ga.GoogleAuth = field(repr=False)

    mime_str: str

    service: Any = field(repr=False, default=None)

    gdrive_mime_type: GDrive_MimeType_Enum = None

    parent_ls: List[str] = field(default=None)

    # gdrive_parents: List[Any] = field(default=None, repr=False)

    file: Any = field(repr=False, default=None)

    def _parent_init(self):
        if not self.service:
            self.service = auth.generate_gdrive_service()

        try:
            self.gdrive_mime_type = GDrive_MimeType_Enum(self.mime_str)

        except Exception as e:
            print(e)

    def __post_init__(self):
        self._parent_init()

    @classmethod
    async def _from_json(
        cls,
        gdrive_obj: dict,
        auth: ga.GoogleAuth,
        service: ga.Resource = None,
        **kwargs,
    ):
        gfile = cls(
            auth=auth,
            service=service or auth.generate_gdrive_service(),
            file=gdrive_obj,
            mime_str=gdrive_obj["mimeType"],
            url=str(gdrive_obj["webViewLink"]),
            id=gdrive_obj["id"],
            name=gdrive_obj["name"],
            parent_ls=gdrive_obj["parent_ls"],
            modified_time=convert_str_to_date(gdrive_obj["modifiedTime"]),
            **kwargs,
        )

        # if gfile.parent_ls:
        #     gfile.gdrive_parents = await ce.gather_with_concurrency(
        #         *[
        #             GDrive_File.get_by_id(
        #                 id=parent, auth=gfile.auth, service=gfile.service
        #             )
        #             for parent in gfile.parent_ls
        #         ],
        #         n=5,
        #     )

        return gfile

    @classmethod
    async def _super_get_by_id(
        cls,
        id: str,
        fields: List[str] = None,
        auth: ga.GoogleAuth = None,
        service: ga.Resource = None,
        return_raw: bool = False,
    ):
        fields = fields or DEFAULT_GDRIVE_FIELDS

        service = service or auth.generate_gdrive_service()

        file = await get_file_by_id(file_id=id, service=service, fields=fields)

        if return_raw:
            return file

        return await cls._from_json(gdrive_obj=file, auth=auth, service=service)

    @classmethod
    async def get_by_id(
        cls,
        id: str,
        auth: ga.GoogleAuth = None,
        service: ga.Resource = None,
        return_raw: bool = False,
        fields: List[str] = None,
    ):

        return await cls._super_get_by_id(
            id=id, auth=auth, service=service, return_raw=return_raw, fields=fields
        )

    async def export(self, output_path: str, export_mime_type: GDrive_MimeType_Enum):

        if isinstance(export_mime_type, GDrive_MimeType_Enum):
            export_mime_type = export_mime_type.name

            return await ExportAs_Enum(export_mime_type).value(
                service=self.service, file_id=self.id, output_path=output_path
            )

        return await ExportAs_Enum(export_mime_type).value(
            service=self.service, file_id=self.id, output_path=output_path
        )

# %% ../../nbs/google/GDrive.ipynb 23
@dataclass
class GDrive_Folder(GDrive_File):

    gdrive_files: List[GDrive_File] = field(default=None)

    async def _get_child_from_json(self, child, is_recursive: bool = True):
        if not child.get("is_folder"):
            g_child = await GDrive_File._from_json(
                gdrive_obj=child, service=self.service, auth=self.auth
            )

        if child.get("is_folder"):
            g_child = await GDrive_Folder._from_json(
                gdrive_obj=child, service=self.service, auth=self.auth
            )

            if is_recursive:
                await g_child.get_files(is_recursive=is_recursive)

        if not self.gdrive_files:
            self.gdrive_files = []

        self.gdrive_files.append(g_child)

    async def get_files(self, is_recursive: bool = True, return_raw: bool = False):

        folder_obj = await walk_folder(
            service=self.service,
            folder_id=self.id,
            is_recursive=is_recursive,
        )

        if return_raw:
            return folder_obj

        await ce.gather_with_concurrency(
            *[
                self._get_child_from_json(child=child, is_recursive=is_recursive)
                for child in folder_obj.get("children")
            ],
            n=10
        )

        return self.gdrive_files
