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

# %% auto 0
__all__ = ['DEFAULT_GDRIVE_FIELDS', 'DEFAULT_GFOLDER_FIELDS', 'GDrive_MimeType_Enum', 'GoogleAPI_RouteError', 'get_file_by_id',
           'get_file_content_as_markdown', 'get_folder_contents_by_id', 'walk_folder_and_flatten', 'walk_folder',
           'create_gfile', 'replace_gfile', 'delete_gfile', 'export_file_as', 'export_as_pptx', 'export_as_markdown',
           'export_as_zip', 'export_metadata_to_file']

# %% ../../../nbs/google/routes/drive.ipynb 2
from googleapiclient.http import MediaFileUpload
from ...client import LibraryError, RouteError

# %% ../../../nbs/google/routes/drive.ipynb 3
from domolibrary_extensions.utils.execution import auto_retry_async

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

import domolibrary_extensions.utils.execution as dece


import domolibrary_extensions.google.GAuth as ga

from googleapiclient.errors import HttpError

import json
import csv
from pathlib import Path
from enum import Enum
from typing import List, Optional, Callable, Dict

# %% ../../../nbs/google/routes/drive.ipynb 7
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'


class GoogleAPI_RouteError(RouteError):
    """Custom error for Google API operations"""

    def __init__(self, message: str = None, e: Exception = None):

        if isinstance(e, Exception):
            # Extract meaningful info from HttpError
            error_details = e.error_details if hasattr(e, "error_details") else str(e)

            super().__init__(f"{message or ''} {error_details}")

        else:
            super().__init__(message)

# %% ../../../nbs/google/routes/drive.ipynb 8
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 = None

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

    except HttpError as e:
        raise GoogleAPI_RouteError("unable to retrieve file", e=e) from e

    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/routes/drive.ipynb 10
async def get_file_content_as_markdown(service, file_id) -> str:
    """
    Get the file content as markdown string.

    Returns:
        File content as markdown string
    """
    try:
        content = (
            service.files()
            .export(fileId=file_id, mimeType=GDrive_MimeType_Enum.markdown.value)
            .execute()
        )
        return content.decode("utf-8") if isinstance(content, bytes) else content

    except HttpError as e:
        raise GoogleAPI_RouteError(doc_id=file_id) from e

# %% ../../../nbs/google/routes/drive.ipynb 12
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
    query: str = None,  # defaults "'{folder_id}' in parents"
    is_just_return_folders: bool = False,
    return_raw: bool = False,  # if return_raw will not execute recursively
    debug_api: bool = False,  # if True will print the query
):

    fields = fields or DEFAULT_GFOLDER_FIELDS

    res = None
    file_ls = []
    page_token = None

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

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

    if debug_api:
        print(query)

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

        except HttpError as e:
            raise GoogleAPI_RouteError("unable to retrieve folder contents", e=e) from e

        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/routes/drive.ipynb 14
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 and organizes into a flat file_ls

    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, is_just_return_folders=just_folders
    )
    file_ls.extend(new_files)

    if new_files and is_recursive:
        await dece.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/routes/drive.ipynb 16
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, is_just_return_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/routes/drive.ipynb 19
async def create_gfile(
    service: ga.Resource,
    media: MediaFileUpload,
    file_name: str,
    mime_type_enum: GDrive_MimeType_Enum = None,
    debug_prn: bool = False,
    parent_folder: str = None,
    debug_api: bool = False,
    return_fields: List[str] = None,
):

    if mime_type_enum is None:
        mime_type_enum = GDrive_MimeType_Enum.gdoc

    file_metadata = {
        "name": file_name,
        "mimeType": mime_type_enum.value,
    }

    if parent_folder:
        file_metadata.update(
            {
                "parents": (
                    parent_folder
                    if isinstance(parent_folder, list)
                    else [parent_folder]
                )
            }
        )

    return_fields = return_fields or DEFAULT_GDRIVE_FIELDS

    if debug_api:
        print({"file_metadata": file_metadata})

    try:
        file = (
            service.files()
            .create(
                body=file_metadata, media_body=media, fields=",".join(return_fields)
            )
            .execute()
        )
    except HttpError as e:
        raise GoogleAPI_RouteError("unable to create file", e=e) from e

    if debug_prn:
        print(f"Uploaded to Google Docs: {file['webViewLink']}")

    return file

# %% ../../../nbs/google/routes/drive.ipynb 20
async def replace_gfile(
    service: ga.Resource,
    media: MediaFileUpload,
    file_id: str,
    return_fields: List[str] = None,
):
    return_fields = return_fields or DEFAULT_GDRIVE_FIELDS

    try:
        file = (
            service.files()
            .update(fileId=file_id, media_body=media, fields=",".join(return_fields))
            .execute()
        )
    except HttpError as e:
        raise GoogleAPI_RouteError("unable to replace file", e=e) from e

    return file

# %% ../../../nbs/google/routes/drive.ipynb 21
async def delete_gfile(
    service: ga.Resource,
    file_id: str,
):
    try:
        file = service.files().delete(fileId=file_id).execute()

    except HttpError as e:
        raise GoogleAPI_RouteError("unable to delete file", e=e) from e

    return True

# %% ../../../nbs/google/routes/drive.ipynb 23
@auto_retry_async(
    max_retry=3,
    exceptions_to_retry=[HttpError],
    sleep_beween_retry=2,
)
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,
):

    content = (
        service.files()
        .export(fileId=file_id, mimeType=export_mime_type.value)
        .execute()
    )

    if not content:
        raise GoogleAPI_RouteError(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):
    from domolibrary_extensions.google.GDrive import GDrive_MimeType_Enum

    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_markdown(service: ga.Resource, file_id: str, output_path: str):
    from domolibrary_extensions.google.GDrive import GDrive_MimeType_Enum

    return await export_file_as(
        service=service,
        file_id=file_id,
        export_mime_type=GDrive_MimeType_Enum.markdown,
        export_fn=download_txt,
        output_path=output_path,
    )


async def export_as_zip(service: ga.Resource, file_id: str, output_path: str):
    from domolibrary_extensions.google.GDrive import GDrive_MimeType_Enum

    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,
    )

# %% ../../../nbs/google/routes/drive.ipynb 24
async def export_metadata_to_file(
    metadata_list: List[Dict[str, str]], output_path: str, format: str = "csv"
) -> str:
    """
    Export metadata to CSV or JSON file.

    Args:
        metadata_list: List of metadata dictionaries from extract_metadata_from_all_files
        output_path: Path where to save the file
        format: 'csv' or 'json'

    Returns:
        Path to the created file
    """
    output_path = Path(output_path)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    if format.lower() == "csv":
        if not str(output_path).endswith(".csv"):
            output_path = output_path.with_suffix(".csv")

        with open(output_path, "w", newline="", encoding="utf-8") as csvfile:
            if metadata_list:
                fieldnames = metadata_list[0].keys()
                writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
                writer.writeheader()
                writer.writerows(metadata_list)

    elif format.lower() == "json":
        if not str(output_path).endswith(".json"):
            output_path = output_path.with_suffix(".json")

        with open(output_path, "w", encoding="utf-8") as jsonfile:
            json.dump(metadata_list, jsonfile, indent=2, ensure_ascii=False)

    else:
        raise ValueError("Format must be 'csv' or 'json'")

    return str(output_path)
