"""Tabular data media type for litemind.

Provides the ``Table`` class which wraps tabular data backed by a file
(CSV, Excel, etc.). Tables can be created from URIs, pandas DataFrames,
numpy arrays, or CSV strings, and rendered as Markdown for LLM consumption.
"""

from typing import Optional, Union

from pandas import DataFrame

from litemind.media.media_uri import MediaURI
from litemind.media.types.media_text import Text
from litemind.utils.load_table import load_table_from_uri


class Table(MediaURI):
    """Media that stores tabular data backed by a file (CSV, Excel, etc.).

    The underlying data is lazily loaded into a pandas DataFrame on first
    access. Tables can be created from URIs, DataFrames, numpy arrays,
    or CSV strings.
    """

    def __init__(self, uri: str, extension: Optional[str] = None, **kwargs):
        """Create a new table media from a file URI.

        Parameters
        ----------
        uri : str
            The URI of the table file (e.g. a CSV or Excel file).
        extension : str, optional
            The expected file extension. If provided and the URI does not
            end with this extension, a ValueError is raised.
        **kwargs
            Additional keyword arguments forwarded to ``MediaURI``.

        Raises
        ------
        ValueError
            If *extension* is provided but does not match the URI.
        """

        # Check that the file extension is valid:
        if extension is not None:
            if not uri.endswith(extension):
                raise ValueError(
                    f"Invalid table URI: '{uri}' (must have a valid table file extension)"
                )

        super().__init__(uri=uri, extension=extension, **kwargs)

        self.dataframe: Optional[DataFrame] = None

    @classmethod
    def from_table(
        cls,
        table: Union["ndarray", "DataFrame", list, tuple],
        filepath: Optional[str] = None,
    ):
        """Create a Table media from tabular data.

        Parameters
        ----------
        table : numpy.ndarray, pandas.DataFrame, list, or tuple
            The tabular data. Arrays and sequences are converted to a
            DataFrame automatically.
        filepath : str, optional
            Destination CSV file path. A temporary file is created if None.

        Returns
        -------
        Table
            A new Table media referencing the saved CSV file.

        Raises
        ------
        ValueError
            If *table* is not a supported type.
        """

        if filepath is None:
            # Create temporary file:
            import tempfile

            from litemind.utils.temp_file_manager import register_temp_file

            filepath = register_temp_file(
                tempfile.NamedTemporaryFile(delete=False).name
            )

        # If dataframe is a numpy array, or array-like, convert to DataFrame:
        if isinstance(table, (list, tuple)) or hasattr(table, "__array__"):
            import pandas as pd

            table = pd.DataFrame(table)

        # At this point `table` should be a DataFrame:
        elif not isinstance(table, DataFrame):
            raise ValueError(
                f"Parameter `table` must be a pandas DataFrame or an array-like object, not {type(table)}"
            )

        # Save dataframe to file:
        table.to_csv(
            filepath,
            index=False,  # avoid the synthetic index column
            na_rep="",  # match original null representation
            float_format=None,  # let 'round_trip' preserve original
        )

        # URI:
        table_uri = "file://" + filepath

        # Create Table from URI:
        return Table(uri=table_uri)

    @classmethod
    def from_csv_string(cls, csv_table_str: str):
        """Create a Table media from a CSV-formatted string.

        Parameters
        ----------
        csv_table_str : str
            The CSV string representing the table data.

        Returns
        -------
        Table
            A new Table media backed by a temporary CSV file.

        Raises
        ------
        ValueError
            If *csv_table_str* is not a string.
        """

        from io import StringIO

        import pandas as pd

        # Check that it is really a string:
        if not isinstance(csv_table_str, str):
            raise ValueError(
                f"Parameter `csv_table_str` must be a string, not {type(csv_table_str)}"
            )

        dataframe = pd.read_csv(
            StringIO(csv_table_str),
            dtype=str,  # keep all columns textual
            keep_default_na=False,  # don't turn 'NA' into NaN
            float_precision="round_trip",  # exact float repr
        )

        # Create a table from the DataFrame:
        return cls.from_table(dataframe)

    def load_from_uri(self):
        """Load the table data from the URI into ``self.dataframe``.

        Raises
        ------
        ValueError
            If no URI is set.
        """

        # Check if the URI is provided
        if self.uri is None:
            raise ValueError("No URI provided to load the table data")

        # Load the table data from the URI
        self.dataframe = load_table_from_uri(self.uri)

    def to_dataframe(self):
        """Get the table data as a pandas DataFrame.

        Loads the data from the URI if not already loaded.

        Returns
        -------
        pandas.DataFrame
            The table data.
        """

        # If the dataframe is not loaded, load it from the URI:
        if self.dataframe is None:
            self.load_from_uri()

        return self.dataframe

    def to_markdown(self):
        """Render the table as a Markdown-formatted string.

        Uses the ``tabulate`` library with GitHub-flavoured table format.

        Returns
        -------
        str
            The table rendered as a Markdown fenced code block.
        """

        #
        if self.dataframe is None:
            self.load_from_uri()

        # Convert DataFrame to markdown table string:
        from tabulate import tabulate

        table_markdown = tabulate(
            self.dataframe, headers="keys", tablefmt="github", showindex=True
        )

        # Convert table to markdown:
        markdown = f"```dataframe\n{table_markdown}\n```"

        return markdown

    def to_markdown_text_media(self):
        """Convert the table to a Text media with Markdown formatting.

        Returns
        -------
        Text
            A Text media containing the Markdown table.
        """

        # Get Markdown:
        markdown = self.to_markdown()

        return Text(markdown)
