"""Result and Record classes for query output."""

from __future__ import annotations

from collections.abc import Iterator
from typing import Any

from piggie.types import Edge, Vertex


class Record:
    """A single row from a query result, with keyed and positional access."""

    __slots__ = ("_data", "_keys")

    def __init__(self, keys: list[str], values: tuple[Any, ...]) -> None:
        self._keys = keys
        self._data = dict(zip(keys, values))

    def __getitem__(self, key: str | int) -> Any:
        if isinstance(key, int):
            return list(self._data.values())[key]
        return self._data[key]

    def get(self, key: str, default: Any = None) -> Any:
        return self._data.get(key, default)

    def keys(self) -> list[str]:
        return list(self._keys)

    def values(self) -> list[Any]:
        return list(self._data.values())

    def items(self) -> list[tuple[str, Any]]:
        return list(self._data.items())

    def __repr__(self) -> str:
        return f"Record({self._data})"

    def __len__(self) -> int:
        return len(self._data)


class Result:
    """Container for query results with DataFrame and dict output.

    The primary output path is ``to_df()`` for pandas DataFrames.
    Also supports ``to_dicts()`` for plain Python dicts, ``to_arrow()``
    for Apache Arrow tables, and ``records`` for Record iteration.
    """

    def __init__(self, rows: list[tuple[Any, ...]], columns: list[str]) -> None:
        self._rows = rows
        self._columns = columns

    @property
    def columns(self) -> list[str]:
        return list(self._columns)

    @property
    def records(self) -> list[Record]:
        return [Record(self._columns, row) for row in self._rows]

    def single(self) -> Record:
        """Return the single record, raising if count != 1."""
        if len(self._rows) != 1:
            raise ValueError(
                f"Expected exactly 1 record, got {len(self._rows)}"
            )
        return Record(self._columns, self._rows[0])

    def __len__(self) -> int:
        return len(self._rows)

    def __iter__(self) -> Iterator[Record]:
        return iter(self.records)

    def __bool__(self) -> bool:
        return len(self._rows) > 0

    def to_dicts(self, *, expand: bool = False) -> list[dict[str, Any]]:
        """Convert results to a list of dicts.

        Args:
            expand: If True, expand Vertex/Edge properties into flat columns.
        """
        result: list[dict[str, Any]] = []
        for row in self._rows:
            d = dict(zip(self._columns, row))
            if expand:
                d = _expand_row(d)
            result.append(d)
        return result

    def to_df(self, *, expand: bool = False) -> Any:
        """Convert results to a pandas DataFrame.

        Args:
            expand: If True, expand Vertex/Edge properties into columns.

        Returns:
            pandas.DataFrame

        Raises:
            ImportError: If pandas is not installed.
        """
        try:
            import pandas as pd
        except ImportError:
            raise ImportError(
                "pandas is required for to_df(). "
                "Install it with: pip install piggie[pandas]"
            ) from None

        data = self.to_dicts(expand=expand)
        if not data:
            return pd.DataFrame(columns=self._columns)
        return pd.DataFrame(data)

    def to_arrow(self, *, expand: bool = False) -> Any:
        """Convert results to a PyArrow Table.

        Args:
            expand: If True, expand Vertex/Edge properties into columns.

        Returns:
            pyarrow.Table

        Raises:
            ImportError: If pyarrow is not installed.
        """
        try:
            import pyarrow as pa
        except ImportError:
            raise ImportError(
                "pyarrow is required for to_arrow(). "
                "Install it with: pip install piggie[arrow]"
            ) from None

        data = self.to_dicts(expand=expand)
        if not data:
            return pa.table({col: [] for col in self._columns})
        return pa.Table.from_pylist(data)

    def __repr__(self) -> str:
        return f"Result(rows={len(self._rows)}, columns={self._columns})"


def _expand_row(row: dict[str, Any]) -> dict[str, Any]:
    """Expand Vertex/Edge objects in a row into flat property columns."""
    expanded: dict[str, Any] = {}
    for key, val in row.items():
        if isinstance(val, Vertex):
            expanded[f"{key}._id"] = val.id
            expanded[f"{key}._label"] = val.label
            for pk, pv in val.properties.items():
                expanded[f"{key}.{pk}"] = pv
        elif isinstance(val, Edge):
            expanded[f"{key}._id"] = val.id
            expanded[f"{key}._label"] = val.label
            expanded[f"{key}._start_id"] = val.start_id
            expanded[f"{key}._end_id"] = val.end_id
            for pk, pv in val.properties.items():
                expanded[f"{key}.{pk}"] = pv
        else:
            expanded[key] = val
    return expanded
