"""Table and Series upload logic.

Serialises tabular data to Parquet and uploads as an artifact.
Accepts DataFrames (pandas/polars), list of dicts, or dict of lists —
no hard dependency on any DataFrame library.
"""

from __future__ import annotations

import tempfile
from pathlib import Path
from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
    from ._types import Run

# Input types for log_table
TableData = Any  # DataFrame | list[dict] | dict[str, list]


def _to_parquet_bytes(*, data: TableData) -> bytes:
    """Convert table data to Parquet bytes.

    Detection is duck-typed — checks for methods, not isinstance:
    - write_parquet() → polars DataFrame
    - to_parquet() → pandas DataFrame
    - dict of lists / list of dicts → plain Python data
    """
    # Has write_parquet? (polars DataFrame)
    if hasattr(data, "write_parquet"):
        return _write_parquet_via_write(data)

    # Has to_parquet but not write_parquet? (pandas DataFrame)
    # Convert through polars when available to avoid pyarrow dependency
    if hasattr(data, "to_parquet") and hasattr(data, "columns"):
        return _write_parquet_from_pandas_like(data)

    # dict of lists: {"col": [1, 2, 3]}
    if isinstance(data, dict) and all(isinstance(v, list) for v in data.values()):
        return _write_parquet_from_dict(data)

    # list of dicts: [{"col": 1}, {"col": 2}]
    if isinstance(data, list) and len(data) > 0 and isinstance(data[0], dict):
        cols: dict[str, list[Any]] = {}
        for row in data:
            for key, val in row.items():
                cols.setdefault(key, []).append(val)
        return _write_parquet_from_dict(cols)

    raise TypeError(f"Expected a DataFrame, list[dict], or dict[str, list], " f"got {type(data).__name__}")


def _write_parquet_via_write(df: Any) -> bytes:
    """Write via write_parquet() (polars-style)."""
    with tempfile.TemporaryDirectory() as tmp:
        path = Path(tmp) / "data.parquet"
        df.write_parquet(path)
        return path.read_bytes()


def _write_parquet_via_to(df: Any) -> bytes:
    """Write via to_parquet() (pandas-style)."""
    with tempfile.TemporaryDirectory() as tmp:
        path = Path(tmp) / "data.parquet"
        df.to_parquet(path)
        return path.read_bytes()


def _write_parquet_from_pandas_like(df: Any) -> bytes:
    """Convert a pandas-like DataFrame to polars first, fall back to to_parquet()."""
    try:
        import polars as pl

        return _write_parquet_via_write(pl.from_pandas(df))
    except (ImportError, TypeError):
        return _write_parquet_via_to(df)


def _write_parquet_from_dict(data: dict[str, list[Any]]) -> bytes:
    """Write a dict-of-lists to Parquet.

    Tries polars first, then pandas. At least one must be installed.
    """
    try:
        import polars as pl

        return _write_parquet_via_write(pl.DataFrame(data))
    except ImportError:
        pass

    try:
        import pandas as pd  # type: ignore[import-untyped]

        return _write_parquet_via_to(pd.DataFrame(data))
    except ImportError:
        pass

    raise ImportError(
        "Either polars or pandas is required to serialize dict/list data to Parquet. "
        "Install one: pip install polars  or  pip install pandas pyarrow"
    )


def _validate_series(*, series: object) -> None:
    """Check that the object looks like a Series."""
    if not (hasattr(series, "name") and hasattr(series, "shape") and hasattr(series, "to_frame")):
        raise TypeError(
            f"Expected a Series-like object with name, shape, and to_frame(), " f"got {type(series).__name__}"
        )


def upload_table(
    *,
    run: "Run",
    name: str,
    data: TableData,
    metadata: dict[str, object] | None = None,
) -> str:
    """Upload tabular data as a Parquet artifact."""
    parquet_bytes = _to_parquet_bytes(data=data)

    version_id = run._client.create_artifact_version(
        run.id,
        name,
        "table",
        filter_data=metadata,
    )

    try:
        run._client.upload_artifact_file(
            version_id,
            "data.parquet",
            parquet_bytes,
            "application/octet-stream",
        )
        run._client.complete_artifact_version(version_id)
    except Exception:
        run._client.fail_artifact_version(version_id)
        raise

    return version_id


def upload_series(
    *,
    run: "Run",
    name: str,
    series: object,
    metadata: dict[str, object] | None = None,
) -> str:
    """Upload a Series by converting to a single-column DataFrame."""
    _validate_series(series=series)

    df = series.to_frame()  # type: ignore[attr-defined]
    return upload_table(
        run=run,
        name=name,
        data=df,
        metadata=metadata,
    )
