"""DataFrameRef: thin wrapper over DuckDB table or Polars LazyFrame."""
from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any, Literal

import duckdb
import polars as pl


@dataclass
class DataFrameRef:
    """Reference to a DataFrame backed by either DuckDB or Polars."""
    backend: Literal["duckdb", "polars"]
    name: str | None = None          # DuckDB table name
    lf: pl.LazyFrame | None = None   # Polars LazyFrame
    con: duckdb.DuckDBPyConnection | None = None  # DuckDB connection (for duckdb backend)
    source_path: str | None = None   # Original CSV path
    metadata: dict[str, Any] = field(default_factory=dict)

    def head(self, n: int = 50, offset: int = 0) -> pl.DataFrame:
        """Return first n rows as a Polars DataFrame."""
        if self.backend == "duckdb" and self.name and self.con:
            result = self.con.execute(
                f'SELECT * FROM "{self.name}" LIMIT {n} OFFSET {offset}'
            ).fetchall()
            col_names = [d[0] for d in self.con.execute(f'SELECT * FROM "{self.name}" LIMIT 0').description]
            return pl.DataFrame(result, schema=col_names, orient="row")
        elif self.backend == "polars" and self.lf is not None:
            return self.lf.slice(offset, n).collect()
        raise ValueError("Invalid DataFrameRef: no backend available")

    def schema(self) -> list[dict[str, str]]:
        """Return column names and types."""
        if self.backend == "duckdb" and self.name and self.con:
            rows = self.con.execute(
                f"SELECT column_name, data_type FROM information_schema.columns WHERE table_name = '{self.name}'"
            ).fetchall()
            return [{"name": r[0], "dtype": r[1]} for r in rows]
        elif self.backend == "polars" and self.lf is not None:
            schema = self.lf.schema
            return [{"name": name, "dtype": str(dtype)} for name, dtype in schema.items()]
        raise ValueError("Invalid DataFrameRef")

    def to_polars_df(self) -> pl.DataFrame:
        """Convert to a full Polars DataFrame."""
        if self.backend == "duckdb" and self.name and self.con:
            result = self.con.execute(f'SELECT * FROM "{self.name}"').fetchall()
            col_names = [d[0] for d in self.con.execute(f'SELECT * FROM "{self.name}" LIMIT 0').description]
            return pl.DataFrame(result, schema=col_names, orient="row")
        elif self.backend == "polars" and self.lf is not None:
            return self.lf.collect()
        raise ValueError("Invalid DataFrameRef")

    def to_lazy(self) -> pl.LazyFrame:
        """Get as Polars LazyFrame."""
        if self.backend == "polars" and self.lf is not None:
            return self.lf
        elif self.backend == "duckdb" and self.name and self.con:
            df = self.to_polars_df()
            return df.lazy()
        raise ValueError("Invalid DataFrameRef")

    def row_count(self) -> int:
        """Return total row count."""
        if self.backend == "duckdb" and self.name and self.con:
            return self.con.execute(f'SELECT COUNT(*) FROM "{self.name}"').fetchone()[0]
        elif self.backend == "polars" and self.lf is not None:
            return self.lf.select(pl.len()).collect().item()
        return 0

    def materialize_to_duckdb(self, table_name: str, con: duckdb.DuckDBPyConnection) -> DataFrameRef:
        """Materialize a Polars LazyFrame into a DuckDB table."""
        if self.backend == "polars" and self.lf is not None:
            df = self.lf.collect()
            # Use Apache Arrow for efficient transfer
            arrow_tbl = df.to_arrow()
            con.register("pl_df", arrow_tbl)
            con.execute(f'CREATE OR REPLACE TABLE "{table_name}" AS SELECT * FROM pl_df')
            con.unregister("pl_df")
            return DataFrameRef(backend="duckdb", name=table_name, con=con, source_path=self.source_path)
        return self

    def clone(self) -> DataFrameRef:
        """Create a shallow clone."""
        return DataFrameRef(
            backend=self.backend,
            name=self.name,
            lf=self.lf,
            con=self.con,
            source_path=self.source_path,
            metadata=dict(self.metadata),
        )