import datetime
import logging
import math
from typing import List, TYPE_CHECKING, Dict, Any, Optional

from dagshub.common.helpers import log_message
from dagshub.common.util import lazy_load, wrap_bytes
from dagshub.data_engine import dtypes
from dagshub.data_engine.client.models import autogenerated_columns
from dagshub.data_engine.dtypes import MetadataFieldType
from dagshub.data_engine.model.metadata.dtypes import DatasourceFieldInfo, DatapointMetadataUpdateEntry
from dagshub.data_engine.model.metadata.util import _get_datetime_utc_offset
from dagshub.data_engine.model.metadata.validation import UploadingMetadataInfo, MAX_STRING_FIELD_LENGTH
from dagshub.data_engine.model.schema_util import metadata_type_lookup, special_metadata_handlers

if TYPE_CHECKING:
    from dagshub.data_engine.model.datasource import Datasource
    import pandas
else:
    pandas = lazy_load("pandas")

logger = logging.getLogger(__name__)


def _transform_strings_to_documents(ds: "Datasource", metadata: List[DatapointMetadataUpdateEntry], field_name: str):
    log_message(f"Uploading field {field_name} as a document field due to long string values")
    ds.metadata_field(field_name).set_type(dtypes.Document).apply()
    for m in metadata:
        if m.key != field_name:
            continue
        m.valueType = MetadataFieldType.BLOB
        m.value = wrap_bytes(m.value.encode("utf-8"))


def run_preupload_transforms(
    ds: "Datasource",
    metadata: List[DatapointMetadataUpdateEntry],
    precalculated_info: Dict[str, UploadingMetadataInfo],
):
    """
    Run transformations on metadata that should happen before uploading.

    Current transformations:
        - New string fields with large values get converted to a Document type
    """

    for info in precalculated_info.values():
        if (
            info.field_type == MetadataFieldType.STRING
            and info.existing_metadata_in_ds is None
            and len(info.longest_value) > MAX_STRING_FIELD_LENGTH
        ):
            _transform_strings_to_documents(ds, metadata, info.field_name)


def _add_metadata(
    field_info: DatasourceFieldInfo,
    buf: List[DatapointMetadataUpdateEntry],
    dp_path: str,
    key: str,
    value: Any,
    is_pandas: bool,
):
    """
    Adds the new metadata to the buf buffer.
    Have to pass in the buffer in order to graduate the already added metadata to multivalue, if that happens.
    """
    if key in autogenerated_columns:
        return
    if value is None:
        return
    if is_pandas:
        # Pandas specific: since pandas doesn't distinguish between None and NaN, don't upload it
        if isinstance(value, float) and math.isnan(value):
            return
        # Pandas specific: Cast pandas Timestamp to datetime
        if isinstance(value, list) or isinstance(value, pandas._libs.tslibs.timestamps.Timestamp):
            if isinstance(value, list):
                if isinstance(value[0], pandas._libs.tslibs.timestamps.Timestamp):
                    value = [val.to_pydatetime() for val in value]
            else:
                value = value.to_pydatetime()

    def _generate_metadata_update_entry(val: Any, is_multivalue: bool) -> Optional[DatapointMetadataUpdateEntry]:
        field_type = field_info.field_value_types.get(key)
        value_type = type(val)

        # Run special handling first
        if value_type in special_metadata_handlers:
            val, converted_field_type = special_metadata_handlers[value_type](val)
            if val is None:
                return None
            if field_type is not None and converted_field_type != field_type:
                raise ValueError(f"Can't put value of type {value_type} into field of type {field_type}")
            return DatapointMetadataUpdateEntry(
                url=dp_path,
                key=key,
                value=val,
                valueType=converted_field_type,
                allowMultiple=is_multivalue,
                timeZone=None,
            )

        tz = None
        if field_type is None:
            field_type = metadata_type_lookup[type(val)]
            field_info.field_value_types[key] = field_type
        # Don't override bytes if they're not bytes - probably just undownloaded values
        if field_type == MetadataFieldType.BLOB and not isinstance(val, bytes):
            if key not in field_info.document_fields:
                return None
        # Pandas quirk - integers are floats on the backend
        if is_pandas:
            if field_type == MetadataFieldType.INTEGER:
                val = int(val)
        if isinstance(val, str) and key in field_info.document_fields:
            val = val.encode("utf-8")
        if isinstance(val, bytes):
            val = wrap_bytes(val)
        if isinstance(val, datetime.datetime):
            tz = _get_datetime_utc_offset(val)
            val = int(val.timestamp() * 1000)
        return DatapointMetadataUpdateEntry(
            url=dp_path,
            key=key,
            value=str(val),
            valueType=field_type,
            allowMultiple=is_multivalue,
            timeZone=tz,
        )

    if isinstance(value, list):
        if key not in field_info.multivalue_fields:
            field_info.multivalue_fields.add(key)
            # Promote all the existing uploading metadata to multivalue
            for update_entry in buf:
                if update_entry.key == key:
                    update_entry.allowMultiple = True
        for sub_val in value:
            metadata_entry = _generate_metadata_update_entry(sub_val, True)
            if metadata_entry is not None:
                buf.append(metadata_entry)
    else:
        metadata_entry = _generate_metadata_update_entry(value, key in field_info.multivalue_fields)
        if metadata_entry is not None:
            buf.append(metadata_entry)
