"""Pandas integration for AiFastDb.

Provides helper functions to convert AiFastDb query results,
graph exports, and document store data to pandas DataFrames.

Usage:
    >>> from aifastdb.pandas_ext import to_dataframe, graph_to_dataframes
    >>> df = to_dataframe(store.all())                   # DocumentStore
    >>> df = to_dataframe(sg.list_persons())             # SocialGraphV2
    >>> entities_df, relations_df = graph_to_dataframes(sg.export_graph())
"""

from __future__ import annotations

from typing import Any, Dict, List, Optional, Sequence

__all__ = [
    "to_dataframe",
    "graph_to_dataframes",
    "query_result_to_dataframe",
    "embedding_matrix",
]


def _ensure_pandas():
    """Lazily import pandas and raise a helpful error if not installed."""
    try:
        import pandas as pd
        return pd
    except ImportError:
        raise ImportError(
            "pandas is required for DataFrame integration. "
            "Install it with: pip install aifastdb[pandas]  or  pip install pandas"
        )


def to_dataframe(
    records: Sequence[Dict[str, Any]],
    *,
    columns: Optional[List[str]] = None,
    drop_embedding: bool = True,
) -> "pd.DataFrame":
    """Convert a list of record dicts to a pandas DataFrame.

    Works with results from:
    - ``DocumentStore.all()``
    - ``DocumentStore.find_by_tag()``
    - ``SocialGraphV2.list_persons()``
    - ``SocialGraphV2.get_friends()``
    - ``Federation.query()["records"]``

    Args:
        records: Sequence of dicts (each becomes a row).
        columns: Optional column filter / ordering.
        drop_embedding: If True (default), drop the ``embedding`` column
            to keep the DataFrame compact.

    Returns:
        pandas.DataFrame
    """
    pd = _ensure_pandas()

    if not records:
        return pd.DataFrame(columns=columns or [])

    df = pd.DataFrame(records)

    if drop_embedding and "embedding" in df.columns:
        df = df.drop(columns=["embedding"])

    if columns:
        # Only keep columns that exist
        existing = [c for c in columns if c in df.columns]
        df = df[existing]

    return df


def graph_to_dataframes(
    export: Dict[str, Any],
) -> "tuple[pd.DataFrame, pd.DataFrame]":
    """Convert a SocialGraphV2 export_graph() result to two DataFrames.

    Args:
        export: The dict returned by ``sg.export_graph()``.
            Expected keys: ``"entities"`` (list[dict]) and ``"relations"`` (list[dict]).

    Returns:
        Tuple of ``(entities_df, relations_df)``.
    """
    pd = _ensure_pandas()

    entities = export.get("entities", [])
    relations = export.get("relations", [])

    entities_df = pd.DataFrame(entities) if entities else pd.DataFrame()
    relations_df = pd.DataFrame(relations) if relations else pd.DataFrame()

    return entities_df, relations_df


def query_result_to_dataframe(
    result: Dict[str, Any],
    *,
    drop_embedding: bool = True,
) -> "pd.DataFrame":
    """Convert a Federation query result to a DataFrame.

    The query result is expected to have a ``"records"`` key containing
    a list of record dicts.

    Args:
        result: The dict returned by ``fed.query()`` or ``fed.cross_query()``.
        drop_embedding: If True (default), drop embedding columns.

    Returns:
        pandas.DataFrame of the records.
    """
    records = result.get("records", [])
    return to_dataframe(records, drop_embedding=drop_embedding)


def embedding_matrix(
    records: Sequence[Dict[str, Any]],
    *,
    key: str = "embedding",
) -> "np.ndarray":
    """Extract embeddings from records as a 2D numpy array.

    Args:
        records: Sequence of dicts, each containing an embedding vector.
        key: The key in each dict that holds the embedding (default: ``"embedding"``).

    Returns:
        numpy.ndarray of shape ``(N, dim)`` with dtype ``float32``.
        Records without the embedding key are skipped.
    """
    import numpy as np

    vecs = [r[key] for r in records if key in r and r[key] is not None]

    if not vecs:
        return np.empty((0, 0), dtype=np.float32)

    return np.array(vecs, dtype=np.float32)
