"""
Display operations mixin for DataFrame.

This mixin provides display and collection operations that can be mixed into
the DataFrame class to add display capabilities.
"""

from typing import Any, List, Union, TYPE_CHECKING

if TYPE_CHECKING:
    from ..dataframe import DataFrame

from ...spark_types import Row, StructType, StructField, StringType


class DisplayOperations:
    """Mixin providing display and collection operations for DataFrame."""

    def show(self: "DataFrame", n: int = 20, truncate: bool = True) -> None:
        """Display DataFrame content in a clean table format.

        Args:
            n: Number of rows to display (default: 20).
            truncate: Whether to truncate long values (default: True).

        Example:
            >>> df.show(5)
            DataFrame[3 rows, 3 columns]
            name    age  salary
            Alice   25   50000
            Bob     30   60000
            Charlie 35   70000
        """
        # Materialize lazy operations if needed
        if self._operations_queue:
            materialized = self._materialize_if_lazy()
            return materialized.show(n, truncate)

        print(f"DataFrame[{len(self.data)} rows, {len(self.schema.fields)} columns]")
        if not self.data:
            print("(empty)")
            return

        # Show first n rows
        display_data = self.data[:n]

        # Get column names
        columns = (
            list(display_data[0].keys()) if display_data else self.schema.fieldNames()
        )

        # Calculate column widths
        col_widths = {}
        for col in columns:
            # Start with column name width
            col_widths[col] = len(col)
            # Check data widths
            for row in display_data:
                value = str(row.get(col, "null"))
                if truncate and len(value) > 20:
                    value = value[:17] + "..."
                col_widths[col] = max(col_widths[col], len(value))

        # Print header (no extra padding) - add blank line for separation
        print()  # Add blank line between metadata and headers
        header_parts = []
        for col in columns:
            header_parts.append(col.ljust(col_widths[col]))
        print(" ".join(header_parts))

        # Print data rows (with padding for alignment)
        for row in display_data:
            row_parts = []
            for col in columns:
                value = str(row.get(col, "null"))
                if truncate and len(value) > 20:
                    value = value[:17] + "..."
                # Add padding to data but not headers
                padded_width = col_widths[col] + 2
                row_parts.append(value.ljust(padded_width))
            print(" ".join(row_parts))

        if len(self.data) > n:
            print(f"\n... ({len(self.data) - n} more rows)")

    def to_markdown(
        self: "DataFrame",
        n: int = 20,
        truncate: bool = True,
        underline_headers: bool = True,
    ) -> str:
        """
        Return DataFrame as a markdown table string.

        Args:
            n: Number of rows to show
            truncate: Whether to truncate long strings
            underline_headers: Whether to underline headers with = symbols

        Returns:
            String representation of DataFrame as markdown table
        """
        if not self.data:
            return f"DataFrame[{len(self.data)} rows, {len(self.schema.fields)} columns]\n\n(empty)"

        # Show first n rows
        display_data = self.data[:n]

        # Get column names
        columns = (
            list(display_data[0].keys()) if display_data else self.schema.fieldNames()
        )

        # Build markdown table
        lines = []
        lines.append(
            f"DataFrame[{len(self.data)} rows, {len(self.schema.fields)} columns]"
        )
        lines.append("")  # Blank line

        # Header row
        header_row = "| " + " | ".join(columns) + " |"
        lines.append(header_row)

        # Separator row - use underlines for better visual distinction
        if underline_headers:
            separator_row = (
                "| " + " | ".join(["=" * len(col) for col in columns]) + " |"
            )
        else:
            separator_row = "| " + " | ".join(["---" for _ in columns]) + " |"
        lines.append(separator_row)

        # Data rows
        for row in display_data:
            row_values = []
            for col in columns:
                value = str(row.get(col, "null"))
                if truncate and len(value) > 20:
                    value = value[:17] + "..."
                row_values.append(value)
            data_row = "| " + " | ".join(row_values) + " |"
            lines.append(data_row)

        if len(self.data) > n:
            lines.append(f"\n... ({len(self.data) - n} more rows)")

        return "\n".join(lines)

    def printSchema(self: "DataFrame") -> None:
        """Print DataFrame schema."""
        print("DataFrame Schema:")
        for field in self.schema.fields:
            nullable = "nullable" if field.nullable else "not nullable"
            print(
                f" |-- {field.name}: {field.dataType.__class__.__name__} ({nullable})"
            )

    def collect(self: "DataFrame") -> List[Row]:
        """Collect all data as list of Row objects."""
        if self._operations_queue:
            materialized = self._materialize_if_lazy()
            return self._get_collection_handler().collect(
                materialized.data, materialized.schema
            )
        return self._get_collection_handler().collect(self.data, self.schema)

    def take(self: "DataFrame", n: int) -> List[Row]:
        """Take first n rows as list of Row objects."""
        if self._operations_queue:
            materialized = self._materialize_if_lazy()
            return self._get_collection_handler().take(
                materialized.data, materialized.schema, n
            )
        return self._get_collection_handler().take(self.data, self.schema, n)

    def head(self: "DataFrame", n: int = 1) -> Union[Row, List[Row], None]:
        """Return first n rows."""
        if self._operations_queue:
            materialized = self._materialize_if_lazy()
            return self._get_collection_handler().head(
                materialized.data, materialized.schema, n
            )
        return self._get_collection_handler().head(self.data, self.schema, n)

    def tail(self: "DataFrame", n: int = 1) -> Union[Row, List[Row], None]:
        """Return last n rows."""
        if self._operations_queue:
            materialized = self._materialize_if_lazy()
            return self._get_collection_handler().tail(
                materialized.data, materialized.schema, n
            )
        return self._get_collection_handler().tail(self.data, self.schema, n)

    def toPandas(self: "DataFrame") -> Any:
        """Convert to pandas DataFrame (requires pandas as optional dependency)."""
        from ..export import DataFrameExporter

        return DataFrameExporter.to_pandas(self)

    def toJSON(self: "DataFrame") -> "DataFrame":
        """Return a single-column DataFrame of JSON strings."""
        import json

        json_rows = [{"value": json.dumps(row)} for row in self.data]

        schema = StructType([StructField("value", StringType())])
        from ..dataframe import DataFrame

        return DataFrame(json_rows, schema, self.storage)

    def count(self: "DataFrame") -> int:
        """Count number of rows."""
        # Materialize lazy operations if needed
        if self._operations_queue:
            materialized = self._materialize_if_lazy()
            # Don't call count() recursively - just return the length of materialized data
            return len(materialized.data)

        if self._cached_count is None:
            self._cached_count = len(self.data)
        return self._cached_count

    def isEmpty(self: "DataFrame") -> bool:
        """Check if DataFrame is empty (PySpark 3.3+).

        Returns:
            True if DataFrame has no rows
        """
        return len(self.data) == 0
