import json
from dataclasses import dataclass, field
from datetime import date, datetime
from pathlib import Path
from typing import TYPE_CHECKING, Any, Optional

from dateutil import parser
from pydantic import Field, Secret

from unstructured_ingest.error import DestinationConnectionError
from unstructured_ingest.utils.dep_check import requires_dependencies
from unstructured_ingest.v2.interfaces import (
    AccessConfig,
    ConnectionConfig,
    FileData,
    Uploader,
    UploaderConfig,
    UploadStager,
    UploadStagerConfig,
)
from unstructured_ingest.v2.logger import logger
from unstructured_ingest.v2.processes.connector_registry import (
    DestinationRegistryEntry,
)

if TYPE_CHECKING:
    from weaviate import Client

CONNECTOR_TYPE = "weaviate"


class WeaviateAccessConfig(AccessConfig):
    access_token: Optional[str] = Field(
        default=None, description="Used to create the bearer token."
    )
    api_key: Optional[str] = None
    client_secret: Optional[str] = None
    password: Optional[str] = None


class WeaviateConnectionConfig(ConnectionConfig):
    host_url: str = Field(description="Weaviate instance url")
    class_name: str = Field(
        description="Name of the class to push the records into, e.g: Pdf-elements"
    )
    access_config: Secret[WeaviateAccessConfig] = Field(
        default=WeaviateAccessConfig(), validate_default=True
    )
    username: Optional[str] = None
    anonymous: bool = Field(default=False, description="if set, all auth values will be ignored")
    scope: Optional[list[str]] = None
    refresh_token: Optional[str] = Field(
        default=None,
        description="Will tie this value to the bearer token. If not provided, "
        "the authentication will expire once the lifetime of the access token is up.",
    )
    connector_type: str = Field(default=CONNECTOR_TYPE, init=False)


class WeaviateUploadStagerConfig(UploadStagerConfig):
    pass


@dataclass
class WeaviateUploadStager(UploadStager):
    upload_stager_config: WeaviateUploadStagerConfig = field(
        default_factory=lambda: WeaviateUploadStagerConfig()
    )

    @staticmethod
    def parse_date_string(date_string: str) -> date:
        try:
            timestamp = float(date_string)
            return datetime.fromtimestamp(timestamp)
        except Exception as e:
            logger.debug(f"date {date_string} string not a timestamp: {e}")
        return parser.parse(date_string)

    @classmethod
    def conform_dict(cls, data: dict) -> None:
        """
        Updates the element dictionary to conform to the Weaviate schema
        """

        # Dict as string formatting
        if record_locator := data.get("metadata", {}).get("data_source", {}).get("record_locator"):
            # Explicit casting otherwise fails schema type checking
            data["metadata"]["data_source"]["record_locator"] = str(json.dumps(record_locator))

        # Array of items as string formatting
        if points := data.get("metadata", {}).get("coordinates", {}).get("points"):
            data["metadata"]["coordinates"]["points"] = str(json.dumps(points))

        if links := data.get("metadata", {}).get("links", {}):
            data["metadata"]["links"] = str(json.dumps(links))

        if permissions_data := (
            data.get("metadata", {}).get("data_source", {}).get("permissions_data")
        ):
            data["metadata"]["data_source"]["permissions_data"] = json.dumps(permissions_data)

        # Datetime formatting
        if date_created := data.get("metadata", {}).get("data_source", {}).get("date_created"):
            data["metadata"]["data_source"]["date_created"] = cls.parse_date_string(
                date_created
            ).strftime(
                "%Y-%m-%dT%H:%M:%S.%fZ",
            )

        if date_modified := data.get("metadata", {}).get("data_source", {}).get("date_modified"):
            data["metadata"]["data_source"]["date_modified"] = cls.parse_date_string(
                date_modified
            ).strftime(
                "%Y-%m-%dT%H:%M:%S.%fZ",
            )

        if date_processed := data.get("metadata", {}).get("data_source", {}).get("date_processed"):
            data["metadata"]["data_source"]["date_processed"] = cls.parse_date_string(
                date_processed
            ).strftime(
                "%Y-%m-%dT%H:%M:%S.%fZ",
            )

        if last_modified := data.get("metadata", {}).get("last_modified"):
            data["metadata"]["last_modified"] = cls.parse_date_string(last_modified).strftime(
                "%Y-%m-%dT%H:%M:%S.%fZ",
            )

        # String casting
        if version := data.get("metadata", {}).get("data_source", {}).get("version"):
            data["metadata"]["data_source"]["version"] = str(version)

        if page_number := data.get("metadata", {}).get("page_number"):
            data["metadata"]["page_number"] = str(page_number)

        if regex_metadata := data.get("metadata", {}).get("regex_metadata"):
            data["metadata"]["regex_metadata"] = str(json.dumps(regex_metadata))

    def run(
        self,
        elements_filepath: Path,
        file_data: FileData,
        output_dir: Path,
        output_filename: str,
        **kwargs: Any,
    ) -> Path:
        with open(elements_filepath) as elements_file:
            elements_contents = json.load(elements_file)
        for element in elements_contents:
            self.conform_dict(data=element)
        output_path = Path(output_dir) / Path(f"{output_filename}.json")
        with open(output_path, "w") as output_file:
            json.dump(elements_contents, output_file)
        return output_path


class WeaviateUploaderConfig(UploaderConfig):
    batch_size: int = Field(default=100, description="Number of records per batch")


@dataclass
class WeaviateUploader(Uploader):
    upload_config: WeaviateUploaderConfig
    connection_config: WeaviateConnectionConfig
    connector_type: str = CONNECTOR_TYPE

    @requires_dependencies(["weaviate"], extras="weaviate")
    def get_client(self) -> "Client":
        from weaviate import Client

        auth = self._resolve_auth_method()
        return Client(url=self.connection_config.host_url, auth_client_secret=auth)

    def precheck(self) -> None:
        try:
            self.get_client()
        except Exception as e:
            logger.error(f"Failed to validate connection {e}", exc_info=True)
            raise DestinationConnectionError(f"failed to validate connection: {e}")

    @requires_dependencies(["weaviate"], extras="weaviate")
    def _resolve_auth_method(self):
        access_configs = self.connection_config.access_config.get_secret_value()
        connection_config = self.connection_config
        if connection_config.anonymous:
            return None

        if access_configs.access_token:
            from weaviate.auth import AuthBearerToken

            return AuthBearerToken(
                access_token=access_configs.access_token,
                refresh_token=connection_config.refresh_token,
            )
        elif access_configs.api_key:
            from weaviate.auth import AuthApiKey

            return AuthApiKey(api_key=access_configs.api_key)
        elif access_configs.client_secret:
            from weaviate.auth import AuthClientCredentials

            return AuthClientCredentials(
                client_secret=access_configs.client_secret, scope=connection_config.scope
            )
        elif connection_config.username and access_configs.password:
            from weaviate.auth import AuthClientPassword

            return AuthClientPassword(
                username=connection_config.username,
                password=access_configs.password,
                scope=connection_config.scope,
            )
        return None

    def run(self, path: Path, file_data: FileData, **kwargs: Any) -> None:
        with path.open("r") as file:
            elements_dict = json.load(file)
        logger.info(
            f"writing {len(elements_dict)} objects to destination "
            f"class {self.connection_config.class_name} "
            f"at {self.connection_config.host_url}",
        )

        client = self.get_client()
        client.batch.configure(batch_size=self.upload_config.batch_size)
        with client.batch as b:
            for e in elements_dict:
                vector = e.pop("embeddings", None)
                b.add_data_object(
                    e,
                    self.connection_config.class_name,
                    vector=vector,
                )


weaviate_destination_entry = DestinationRegistryEntry(
    connection_config=WeaviateConnectionConfig,
    uploader=WeaviateUploader,
    uploader_config=WeaviateUploaderConfig,
    upload_stager=WeaviateUploadStager,
    upload_stager_config=WeaviateUploadStagerConfig,
)
