from __future__ import annotations

import base64
import csv
import io
import json
import zlib
from copy import deepcopy
from datetime import datetime
from typing import Any, Iterable, Optional, Sequence, cast

import ndjson

from unstructured.documents.coordinates import PixelSpace
from unstructured.documents.elements import (
    TYPE_TO_TEXT_ELEMENT_MAP,
    CheckBox,
    Element,
    ElementMetadata,
)
from unstructured.partition.common.common import exactly_one
from unstructured.utils import Point, dependency_exists, requires_dependencies

if dependency_exists("pandas"):
    import pandas as pd

# ================================================================================================
# SERIALIZATION/DESERIALIZATION (SERDE) RELATED FUNCTIONS
# ================================================================================================
# These serde functions will likely relocate to `unstructured.documents.elements` since they are
# so closely related to elements and this staging "brick" is deprecated.
# ================================================================================================

# == DESERIALIZERS ===============================


def elements_from_base64_gzipped_json(b64_encoded_elements: str) -> list[Element]:
    """Restore Base64-encoded gzipped JSON elements to element objects.

    This is used to when deserializing `ElementMetadata.orig_elements` from its compressed form in
    JSON and dict forms and perhaps for other purposes.
    """
    # -- Base64 str -> gzip-encoded (JSON) bytes --
    decoded_b64_bytes = base64.b64decode(b64_encoded_elements)
    # -- undo gzip compression --
    elements_json_bytes = zlib.decompress(decoded_b64_bytes)
    # -- JSON (bytes) to JSON (str) --
    elements_json_str = elements_json_bytes.decode("utf-8")
    # -- JSON (str) -> dicts --
    element_dicts = json.loads(elements_json_str)
    # -- dicts -> elements --
    return elements_from_dicts(element_dicts)


def elements_from_dicts(element_dicts: Iterable[dict[str, Any]]) -> list[Element]:
    """Convert a list of element-dicts to a list of elements."""
    elements: list[Element] = []

    for item in element_dicts:
        element_id: str = item.get("element_id", None)
        metadata = (
            ElementMetadata()
            if item.get("metadata") is None
            else ElementMetadata.from_dict(item["metadata"])
        )

        if item.get("type") in TYPE_TO_TEXT_ELEMENT_MAP:
            ElementCls = TYPE_TO_TEXT_ELEMENT_MAP[item["type"]]
            elements.append(ElementCls(text=item["text"], element_id=element_id, metadata=metadata))
        elif item.get("type") == "CheckBox":
            elements.append(
                CheckBox(checked=item["checked"], element_id=element_id, metadata=metadata)
            )

    return elements


# -- legacy aliases for elements_from_dicts() --
isd_to_elements = elements_from_dicts
dict_to_elements = elements_from_dicts


def elements_from_json(
    filename: str = "", text: str = "", encoding: str = "utf-8"
) -> list[Element]:
    """Loads a list of elements from a JSON file or a string."""
    exactly_one(filename=filename, text=text)

    if filename:
        with open(filename, encoding=encoding) as f:
            element_dicts = json.load(f)
    else:
        element_dicts = json.loads(text)

    return elements_from_dicts(element_dicts)


# == SERIALIZERS =================================


def elements_to_base64_gzipped_json(elements: Iterable[Element]) -> str:
    """Convert `elements` to Base64-encoded gzipped JSON.

    This is used to when serializing `ElementMetadata.orig_elements` to make it as compact as
    possible when transported as JSON, for example in an HTTP response. This compressed form is also
    present when elements are in dict form ("element_dicts"). This function is not coupled to that
    purpose however and could have other uses.
    """
    # -- adjust floating-point precision of coordinates down for a more compact str value --
    precision_adjusted_elements = _fix_metadata_field_precision(elements)
    # -- serialize elements as dicts --
    element_dicts = elements_to_dicts(precision_adjusted_elements)
    # -- serialize the dicts to JSON (bytes) --
    json_bytes = json.dumps(element_dicts, sort_keys=True).encode("utf-8")
    # -- compress the JSON bytes with gzip compression --
    deflated_bytes = zlib.compress(json_bytes)
    # -- base64-encode those bytes so they can be serialized as a JSON string value --
    b64_deflated_bytes = base64.b64encode(deflated_bytes)
    # -- convert to a string suitable for serializing in JSON --
    return b64_deflated_bytes.decode("utf-8")


def elements_to_dicts(elements: Iterable[Element]) -> list[dict[str, Any]]:
    """Convert document elements to element-dicts."""
    return [e.to_dict() for e in elements]


# -- legacy aliases for elements_to_dicts() --
convert_to_isd = elements_to_dicts
convert_to_dict = elements_to_dicts


def elements_to_json(
    elements: Iterable[Element],
    filename: Optional[str] = None,
    indent: int = 4,
    encoding: str = "utf-8",
) -> str:
    """Serialize `elements` to a JSON array.

    Also writes the JSON to `filename` if it is provided, encoded using `encoding`.

    The JSON is returned as a string.
    """
    # -- serialize `elements` as a JSON array (str) --
    precision_adjusted_elements = _fix_metadata_field_precision(elements)
    element_dicts = elements_to_dicts(precision_adjusted_elements)
    json_str = json.dumps(element_dicts, indent=indent, sort_keys=True)

    if filename is not None:
        with open(filename, "w", encoding=encoding) as f:
            f.write(json_str)

    return json_str


def elements_to_ndjson(
    elements: Iterable[Element],
    filename: Optional[str] = None,
    encoding: str = "utf-8",
) -> str:
    """Serialize `elements` to a JSON array.

    Also writes the JSON to `filename` if it is provided, encoded using `encoding`.

    The JSON is returned as a string.
    """
    # -- serialize `elements` as a JSON array (str) --
    precision_adjusted_elements = _fix_metadata_field_precision(elements)
    element_dicts = elements_to_dicts(precision_adjusted_elements)
    ndjson_str = ndjson.dumps(element_dicts, sort_keys=True)

    if filename is not None:
        with open(filename, "w", encoding=encoding) as f:
            f.write(ndjson_str)

    return ndjson_str


def _fix_metadata_field_precision(elements: Iterable[Element]) -> list[Element]:
    out_elements: list[Element] = []
    for element in elements:
        el = deepcopy(element)
        if el.metadata.coordinates:
            precision = 1 if isinstance(el.metadata.coordinates.system, PixelSpace) else 2
            points = el.metadata.coordinates.points
            assert points is not None
            rounded_points: list[Point] = []
            for point in points:
                x, y = point
                rounded_point = (round(x, precision), round(y, precision))
                rounded_points.append(rounded_point)
            el.metadata.coordinates.points = tuple(rounded_points)

        if el.metadata.detection_class_prob:
            el.metadata.detection_class_prob = round(el.metadata.detection_class_prob, 5)

        out_elements.append(el)

    return out_elements


# ================================================================================================


def _get_metadata_table_fieldnames() -> list[str]:
    metadata_fields = list(ElementMetadata.__annotations__.keys())
    metadata_fields.remove("coordinates")
    metadata_fields.extend(
        [
            "sender",
            "coordinates_points",
            "coordinates_system",
            "coordinates_layout_width",
            "coordinates_layout_height",
        ],
    )
    return metadata_fields


TABLE_FIELDNAMES: list[str] = [
    "type",
    "text",
    "element_id",
] + _get_metadata_table_fieldnames()


def convert_to_text(elements: Iterable[Element]) -> str:
    """Convert elements into clean, concatenated text."""
    return "\n".join([e.text for e in elements if hasattr(e, "text") and e.text])


def elements_to_text(
    elements: Iterable[Element], filename: Optional[str] = None, encoding: str = "utf-8"
) -> Optional[str]:
    """Convert text from each of `elements` into clean, concatenated text.

    Saves to a txt file if filename is specified. Otherwise, return the text of the elements as a
    string.
    """
    element_cct = convert_to_text(elements)
    if filename is not None:
        with open(filename, "w", encoding=encoding) as f:
            f.write(element_cct)
            return None
    else:
        return element_cct


def flatten_dict(
    dictionary: dict[str, Any],
    parent_key: str = "",
    separator: str = "_",
    flatten_lists: bool = False,
    remove_none: bool = False,
    keys_to_omit: Optional[Sequence[str]] = None,
) -> dict[str, Any]:
    """Flattens a nested dictionary into a single level dictionary.

    keys_to_omit is a list of keys that don't get flattened. If omitting a nested key, format as
    {parent_key}{separator}{key}. If flatten_lists is True, then lists and tuples are flattened as
    well. If remove_none is True, then None keys/values are removed from the flattened
    dictionary.
    """
    keys_to_omit = keys_to_omit if keys_to_omit else []
    flattened_dict: dict[str, Any] = {}
    for key, value in dictionary.items():
        new_key = f"{parent_key}{separator}{key}" if parent_key else key
        if new_key in keys_to_omit:
            flattened_dict[new_key] = value
        elif value is None and remove_none:
            continue
        elif isinstance(value, dict):
            value = cast("dict[str, Any]", value)
            flattened_dict.update(
                flatten_dict(
                    value, new_key, separator, flatten_lists, remove_none, keys_to_omit=keys_to_omit
                ),
            )
        elif isinstance(value, (list, tuple)) and flatten_lists:
            value = cast("list[Any] | tuple[Any]", value)
            for index, item in enumerate(value):
                flattened_dict.update(
                    flatten_dict(
                        {f"{new_key}{separator}{index}": item},
                        "",
                        separator,
                        flatten_lists,
                        remove_none,
                        keys_to_omit=keys_to_omit,
                    )
                )
        else:
            flattened_dict[new_key] = value

    return flattened_dict


def _get_table_fieldnames(rows: list[dict[str, Any]]):
    return list(TABLE_FIELDNAMES)


def convert_to_csv(elements: Iterable[Element]) -> str:
    """Convert `elements` to CSV format."""
    rows: list[dict[str, Any]] = elements_to_dicts(elements)
    table_fieldnames = _get_table_fieldnames(rows)
    # NOTE(robinson) - flatten metadata and add it to the table
    for row in rows:
        metadata = row.pop("metadata")
        for key, value in flatten_dict(metadata).items():
            if key in table_fieldnames:
                row[key] = value

        if row.get("sent_from"):
            row["sender"] = row.get("sent_from")
            if isinstance(row["sender"], list):
                row["sender"] = row["sender"][0]

    with io.StringIO() as buffer:
        csv_writer = csv.DictWriter(buffer, fieldnames=table_fieldnames)
        csv_writer.writeheader()
        csv_writer.writerows(rows)
        return buffer.getvalue()


# -- legacy alias for convert_to_csv --
convert_to_isd_csv = convert_to_csv


@requires_dependencies(["pandas"])
def get_default_pandas_dtypes() -> dict[str, Any]:
    return {
        "text": pd.StringDtype(),  # type: ignore
        "type": pd.StringDtype(),  # type: ignore
        "element_id": pd.StringDtype(),  # type: ignore
        "filename": pd.StringDtype(),  # Optional[str]  # type: ignore
        "filetype": pd.StringDtype(),  # Optional[str]  # type: ignore
        "file_directory": pd.StringDtype(),  # Optional[str]  # type: ignore
        "last_modified": pd.StringDtype(),  # Optional[str]  # type: ignore
        "attached_to_filename": pd.StringDtype(),  # Optional[str]  # type: ignore
        "parent_id": pd.StringDtype(),  # Optional[str],  # type: ignore
        "category_depth": "Int64",  # Optional[int]
        "image_path": pd.StringDtype(),  # Optional[str]  # type: ignore
        "languages": object,  # Optional[list[str]]
        "page_number": "Int64",  # Optional[int]
        "page_name": pd.StringDtype(),  # Optional[str]  # type: ignore
        "url": pd.StringDtype(),  # Optional[str]  # type: ignore
        "link_urls": pd.StringDtype(),  # Optional[str]  # type: ignore
        "link_texts": object,  # Optional[list[str]]
        "links": object,
        "sent_from": object,  # Optional[list[str]],
        "sent_to": object,  # Optional[list[str]]
        "subject": pd.StringDtype(),  # Optional[str]  # type: ignore
        "section": pd.StringDtype(),  # Optional[str]  # type: ignore
        "header_footer_type": pd.StringDtype(),  # Optional[str]  # type: ignore
        "emphasized_text_contents": object,  # Optional[list[str]]
        "emphasized_text_tags": object,  # Optional[list[str]]
        "text_as_html": pd.StringDtype(),  # Optional[str]  # type: ignore
        "max_characters": "Int64",  # Optional[int]
        "is_continuation": "boolean",  # Optional[bool]
        "detection_class_prob": float,  # Optional[float],
        "sender": pd.StringDtype(),  # type: ignore
        "coordinates_points": object,
        "coordinates_system": pd.StringDtype(),  # type: ignore
        "coordinates_layout_width": float,
        "coordinates_layout_height": float,
        "data_source_url": pd.StringDtype(),  # Optional[str]  # type: ignore
        "data_source_version": pd.StringDtype(),  # Optional[str]  # type: ignore
        "data_source_record_locator": object,
        "data_source_date_created": pd.StringDtype(),  # Optional[str]  # type: ignore
        "data_source_date_modified": pd.StringDtype(),  # Optional[str]  # type: ignore
        "data_source_date_processed": pd.StringDtype(),  # Optional[str]  # type: ignore
        "data_source_permissions_data": object,
        "embeddings": object,
    }


@requires_dependencies(["pandas"])
def convert_to_dataframe(
    elements: Iterable[Element], drop_empty_cols: bool = True, set_dtypes: bool = False
) -> "pd.DataFrame":
    """Convert `elements` to a pandas DataFrame.

    The dataframe contains the following columns:
        text: the element text
        type: the text type (NarrativeText, Title, etc)

    Output is pd.DataFrame
    """
    element_dicts = elements_to_dicts(elements)
    for d in element_dicts:
        if metadata := d.pop("metadata", None):
            d.update(flatten_dict(metadata, keys_to_omit=["data_source_record_locator"]))
    df = pd.DataFrame.from_dict(element_dicts)  # type: ignore
    if set_dtypes:
        dt = {k: v for k, v in get_default_pandas_dtypes().items() if k in df.columns}
        df = df.astype(dt)  # type: ignore
    if drop_empty_cols:
        df.dropna(axis=1, how="all", inplace=True)  # type: ignore
    return df


def filter_element_types(
    elements: Iterable[Element],
    include_element_types: Optional[Sequence[type[Element]]] = None,
    exclude_element_types: Optional[Sequence[type[Element]]] = None,
) -> list[Element]:
    """Filters document elements by element type"""
    exactly_one(
        include_element_types=include_element_types,
        exclude_element_types=exclude_element_types,
    )

    filtered_elements: list[Element] = []
    if include_element_types:
        for element in elements:
            if type(element) in include_element_types:
                filtered_elements.append(element)

        return filtered_elements

    elif exclude_element_types:
        for element in elements:
            if type(element) not in exclude_element_types:
                filtered_elements.append(element)

        return filtered_elements

    return list(elements)


def convert_to_coco(
    elements: Iterable[Element],
    dataset_description: Optional[str] = None,
    dataset_version: str = "1.0",
    contributors: tuple[str] = ("Unstructured Developers",),
) -> dict[str, Any]:
    from unstructured.documents.elements import TYPE_TO_TEXT_ELEMENT_MAP

    coco_dataset: dict[str, Any] = {}
    # Handle Info
    coco_dataset["info"] = {
        "description": (
            dataset_description
            if dataset_description
            else f"Unstructured COCO Dataset {datetime.now().strftime('%Y-%m-%d')}"
        ),
        "version": dataset_version,
        "year": datetime.now().year,
        "contributors": ",".join(contributors),
        "date_created": datetime.now().date().isoformat(),
    }
    element_dicts = elements_to_dicts(elements)
    # Handle Images
    images = [
        {
            "width": (
                el["metadata"]["coordinates"]["layout_width"]
                if el["metadata"].get("coordinates")
                else None
            ),
            "height": (
                el["metadata"]["coordinates"]["layout_height"]
                if el["metadata"].get("coordinates")
                else None
            ),
            "file_directory": el["metadata"].get("file_directory", ""),
            "file_name": el["metadata"].get("filename", ""),
            "page_number": el["metadata"].get("page_number", ""),
        }
        for el in element_dicts
    ]
    images = list({tuple(sorted(d.items())): d for d in images}.values())
    for index, d in enumerate(images):
        d["id"] = index + 1
    coco_dataset["images"] = images
    # Handle Categories
    categories = sorted(set(TYPE_TO_TEXT_ELEMENT_MAP.keys()))
    categories = [{"id": i + 1, "name": cat} for i, cat in enumerate(categories)]
    coco_dataset["categories"] = categories
    # Handle Annotations
    annotations = [
        {
            "id": el["element_id"],
            "category_id": [x["id"] for x in categories if x["name"] == el["type"]][0],
            "bbox": (
                [
                    float(el["metadata"].get("coordinates")["points"][0][0]),
                    float(el["metadata"].get("coordinates")["points"][0][1]),
                    float(
                        abs(
                            el["metadata"].get("coordinates")["points"][0][0]
                            - el["metadata"].get("coordinates")["points"][2][0]
                        )
                    ),
                    float(
                        abs(
                            el["metadata"].get("coordinates")["points"][0][1]
                            - el["metadata"].get("coordinates")["points"][1][1]
                        )
                    ),
                ]
                if el["metadata"].get("coordinates")
                else []
            ),
            "area": (
                (
                    float(
                        abs(
                            el["metadata"].get("coordinates")["points"][0][0]
                            - el["metadata"].get("coordinates")["points"][2][0]
                        )
                    )
                    * float(
                        abs(
                            el["metadata"].get("coordinates")["points"][0][1]
                            - el["metadata"].get("coordinates")["points"][1][1]
                        )
                    )
                )
                if el["metadata"].get("coordinates")
                else None
            ),
        }
        for el in element_dicts
    ]
    coco_dataset["annotations"] = annotations
    return coco_dataset
