from __future__ import annotations
from pathlib import Path
import duckdb
import polars as pl


class DuckDBStore:
    """
    Lightweight DuckDB storage wrapper.
    """
    def __init__(self, path: str | Path):
        self.path = Path(path)
        self.path.parent.mkdir(parents=True, exist_ok=True)
        self.conn = duckdb.connect(self.path.as_posix())

    # Writes
    def write_table(
        self,
        name: str,
        df: pl.DataFrame,
        *,
        mode: str = "append",
    ) -> None:
        """
        Writes a Polars DataFrame to DuckDB.
        """
        if df.is_empty():
            return

        if mode not in {"append", "replace"}:
            raise ValueError("mode must be 'append' or 'replace'")

        self.conn.register("df_tmp", df.to_arrow())

        if mode == "replace":
            self.conn.execute(f"DROP TABLE IF EXISTS {name}")

        self.conn.execute(
            f"""
            CREATE TABLE IF NOT EXISTS {name} AS
            SELECT * FROM df_tmp
            """
        )

        if mode == "append":
            self.conn.execute(
                f"INSERT INTO {name} SELECT * FROM df_tmp"
            )

        self.conn.unregister("df_tmp")

    # Reads
    def read_table(self, name: str) -> pl.DataFrame:
        return pl.from_arrow(
            self.conn.execute(f"SELECT * FROM {name}").arrow()
        )

    def query(self, sql: str) -> pl.DataFrame:
        return pl.from_arrow(
            self.conn.execute(sql).arrow()
        )


    # Utilities
    def list_tables(self) -> list[str]:
        return [
            r[0]
            for r in self.conn.execute(
                "SHOW TABLES"
            ).fetchall()
        ]

    def close(self):
        self.conn.close()