from __future__ import annotations

from contextlib import contextmanager
from dataclasses import dataclass, field
from pathlib import Path
from time import time
from typing import TYPE_CHECKING, Any, Generator, Optional, Union

from dateutil import parser
from pydantic import Field, Secret

from unstructured_ingest.data_types.file_data import FileDataSourceMetadata
from unstructured_ingest.errors_v2 import ProviderError, UserError
from unstructured_ingest.logger import logger
from unstructured_ingest.processes.connector_registry import (
    DestinationRegistryEntry,
    SourceRegistryEntry,
)
from unstructured_ingest.processes.connectors.fsspec.fsspec import (
    FsspecAccessConfig,
    FsspecConnectionConfig,
    FsspecDownloader,
    FsspecDownloaderConfig,
    FsspecIndexer,
    FsspecIndexerConfig,
    FsspecUploader,
    FsspecUploaderConfig,
)
from unstructured_ingest.processes.utils.blob_storage import (
    BlobStoreUploadStager,
    BlobStoreUploadStagerConfig,
)
from unstructured_ingest.utils.dep_check import requires_dependencies
from unstructured_ingest.utils.string_and_date_utils import json_to_dict

if TYPE_CHECKING:
    from gcsfs import GCSFileSystem

CONNECTOR_TYPE = "gcs"


class GcsIndexerConfig(FsspecIndexerConfig):
    pass


service_account_key_description = """
  Options:
  - ``None``, GCSFS will attempt to guess your credentials in the
following order: gcloud CLI default, gcsfs cached token, google compute
metadata service, anonymous.
  - ``'google_default'``, your default gcloud credentials will be used,
    which are typically established by doing ``gcloud login`` in a terminal.
  - ``'cache'``, credentials from previously successful gcsfs
    authentication will be used (use this after "browser" auth succeeded)
  - ``'anon'``, no authentication is performed, and you can only
    access data which is accessible to allUsers (in this case, the project and
    access level parameters are meaningless)
  - ``'browser'``, you get an access code with which you can
    authenticate via a specially provided URL
  - if ``'cloud'``, we assume we are running within google compute
    or google container engine, and query the internal metadata directly for
    a token.
  - you may supply a token generated by the
    [gcloud](https://cloud.google.com/sdk/docs/)
    utility; this is either a python dictionary or the name of a file
    containing the JSON returned by logging in with the gcloud CLI tool.
  """


class GcsAccessConfig(FsspecAccessConfig):
    service_account_key: Optional[str] = Field(
        default=None, description=service_account_key_description
    )
    token: Union[str, dict, None] = Field(init=False, default=None)

    def model_post_init(self, __context: Any) -> None:
        ALLOWED_AUTH_VALUES = "google_default", "cache", "anon", "browser", "cloud"

        # Case: null value
        if not self.service_account_key:
            return

        # Case: one of auth constants
        if self.service_account_key in ALLOWED_AUTH_VALUES:
            self.token = self.service_account_key
            return

        # Case: token as json
        if isinstance(json_to_dict(self.service_account_key), dict):
            self.token = json_to_dict(self.service_account_key)
            return

        # Case: path to token
        if Path(self.service_account_key).is_file():
            self.token = self.service_account_key
            return

        raise ValueError("Invalid auth token value")


class GcsConnectionConfig(FsspecConnectionConfig):
    supported_protocols: list[str] = field(default_factory=lambda: ["gs", "gcs"], init=False)
    access_config: Secret[GcsAccessConfig] = Field(default=GcsAccessConfig(), validate_default=True)
    connector_type: str = Field(default=CONNECTOR_TYPE, init=False)

    @requires_dependencies(["gcsfs", "fsspec"], extras="gcs")
    @contextmanager
    def get_client(self, protocol: str) -> Generator["GCSFileSystem", None, None]:
        with super().get_client(protocol=protocol) as client:
            yield client

    def wrap_error(self, e: Exception) -> Exception:
        # https://github.com/fsspec/gcsfs/blob/main/gcsfs/retry.py#L79
        from gcsfs.retry import HttpError

        if isinstance(e, FileNotFoundError):
            raise UserError(f"File not found: {e}")
        if isinstance(e, OSError) and "Forbidden" in str(e):
            raise UserError(e)
        if isinstance(e, ValueError) and "Bad Request" in str(e):
            raise UserError(e)
        if isinstance(e, HttpError) and (http_error_code := e.code):
            message = e.message or e
            if 400 <= http_error_code < 500:
                raise UserError(message)
            if http_error_code >= 500:
                raise ProviderError(message)
        logger.error(f"unhandled exception from gcs ({type(e)}): {e}", exc_info=True)
        return e


@dataclass
class GcsIndexer(FsspecIndexer):
    connection_config: GcsConnectionConfig
    index_config: GcsIndexerConfig
    connector_type: str = CONNECTOR_TYPE

    def get_metadata(self, file_info: dict) -> FileDataSourceMetadata:
        path = file_info["name"]
        date_created = None
        date_modified = None
        if modified_at_str := file_info.get("updated"):
            date_modified = str(parser.parse(modified_at_str).timestamp())
        if created_at_str := file_info.get("timeCreated"):
            date_created = str(parser.parse(created_at_str).timestamp())

        file_size = file_info.get("size") if "size" in file_info else None

        version = file_info.get("etag")
        record_locator = {
            "protocol": self.index_config.protocol,
            "remote_file_path": self.index_config.remote_url,
            "file_id": file_info.get("id"),
        }
        return FileDataSourceMetadata(
            date_created=date_created,
            date_modified=date_modified,
            date_processed=str(time()),
            version=version,
            url=f"{self.index_config.protocol}://{path}",
            record_locator=record_locator,
            filesize_bytes=file_size,
        )


class GcsDownloaderConfig(FsspecDownloaderConfig):
    pass


@dataclass
class GcsDownloader(FsspecDownloader):
    protocol: str = "gcs"
    connection_config: GcsConnectionConfig
    connector_type: str = CONNECTOR_TYPE
    download_config: Optional[GcsDownloaderConfig] = field(default_factory=GcsDownloaderConfig)


class GcsUploaderConfig(FsspecUploaderConfig):
    pass


@dataclass
class GcsUploader(FsspecUploader):
    connector_type: str = CONNECTOR_TYPE
    connection_config: GcsConnectionConfig
    upload_config: GcsUploaderConfig = field(default=None)


gcs_source_entry = SourceRegistryEntry(
    indexer=GcsIndexer,
    indexer_config=GcsIndexerConfig,
    downloader=GcsDownloader,
    downloader_config=GcsDownloaderConfig,
    connection_config=GcsConnectionConfig,
)

gcs_destination_entry = DestinationRegistryEntry(
    uploader=GcsUploader,
    uploader_config=GcsUploaderConfig,
    connection_config=GcsConnectionConfig,
    upload_stager_config=BlobStoreUploadStagerConfig,
    upload_stager=BlobStoreUploadStager,
)
