from __future__ import annotations

import os
import sys
from typing import Any, Dict, Iterable, List, Optional, Tuple

import psycopg2
import psycopg2.extras


# --------------------------------------------------------------------------- #
# Type aliases
# --------------------------------------------------------------------------- #
Row = Dict[str, Any]
ColumnDef = Tuple[str, str, str]  # (name, type, constraints)


# --------------------------------------------------------------------------- #
# Helper class
# --------------------------------------------------------------------------- #
class PSQLDB:
    """
    A thin wrapper around psycopg2 that provides generic CRUD helpers.
    """

    def __init__(
        self,
        *,
        host: str = "localhost",
        port: int = 5432,
        database: str,
        user: str,
        password: str,
        sslmode: str = "prefer",
        autocommit: bool = True,
    ) -> None:
        """
        Open a connection to the PostgreSQL database.

        Parameters
        ----------
        host : str
            Database host.
        port : int
            Database port.
        database : str
            Database name.
        user : str
            Username.
        password : str
            Password.
        sslmode : str
            SSL mode (default: "prefer").
        autocommit : bool
            If True, each statement is committed automatically.
        """
        self.conn = psycopg2.connect(
            host=host,
            port=port,
            dbname=database,
            user=user,
            password=password,
            sslmode=sslmode,
        )
        self.conn.autocommit = autocommit
        self.cursor_factory = psycopg2.extras.RealDictCursor
        self.cursor = self.conn.cursor(cursor_factory=self.cursor_factory)

    # --------------------------------------------------------------------- #
    # Context manager support
    # --------------------------------------------------------------------- #
    def __enter__(self) -> "PSQLDB":
        return self

    def __exit__(self, exc_type, exc, tb) -> None:
        if exc_type:
            self.conn.rollback()
        else:
            self.conn.commit()
        self.cursor.close()
        self.conn.close()

    def close(self) -> None:
        """Explicitly close the underlying connection."""
        self.cursor.close()
        self.conn.close()

    # --------------------------------------------------------------------- #
    # Internal helpers
    # --------------------------------------------------------------------- #
    def _build_create_sql(self, table: str, columns: Iterable[ColumnDef], primary_key: Optional[str]) -> str:
        parts = []
        for name, col_type, constraints in columns:
            col_part = f"{name} {col_type}"
            if constraints:
                col_part += f" {constraints}"
            parts.append(col_part)

        pk_clause = f", PRIMARY KEY ({primary_key})" if primary_key else ""
        columns_sql = ", ".join(parts)
        return f"CREATE TABLE IF NOT EXISTS {table} ({columns_sql}{pk_clause});"

    def _build_set_clause(self, data: Dict[str, Any]) -> Tuple[str, List[Any]]:
        """
        Helper for UPDATE and INSERT. Returns a ``SET``‑style clause
        and the corresponding values list.

        Example:  data={'name':'bob', 'age':30}
                 → ("name = %s, age = %s", ['bob', 30])
        """
        if not data:
            raise ValueError("data dictionary must contain at least one key")

        clause = ", ".join(f"{k} = %s" for k in data)
        values = list(data.values())
        return clause, values

    # --------------------------------------------------------------------- #
    # Schema helpers
    # --------------------------------------------------------------------- #
    def create_table(
        self,
        table: str,
        columns: Iterable[ColumnDef],
        primary_key: Optional[str] = None,
        *,
        if_not_exists: bool = True,
        # You can add more constraints here – e.g. UNIQUE, CHECK, etc.
    ) -> None:
        """
        Create a new table with the supplied column set.

        Parameters
        ----------
        table : str
            Table name.
        columns : iterable of (name, type, constraints)
            Constraints are appended verbatim – e.g. "NOT NULL", "UNIQUE".
        primary_key : str | None
            Column name that should become the primary key. If None, no PK is added.
        if_not_exists : bool
            If True (default) the statement is `CREATE TABLE IF NOT EXISTS`.
        """
        sql = self._build_create_sql(table, columns, primary_key)
        if not if_not_exists:
            sql = sql.replace("IF NOT EXISTS", "")
        self.cursor.execute(sql)

    # --------------------------------------------------------------------- #
    # CRUD helpers
    # --------------------------------------------------------------------- #
    def insert(self, table: str, data: Dict[str, Any]) -> int:
        """
        Insert a single row and return the *generated* primary key
        (if the table uses a `SERIAL`/`IDENTITY` PK).
        If you supply the PK yourself, the returned value will just be that value.

        Parameters
        ----------
        table : str
            Table name.
        data : dict
            Column/value mapping.

        Returns
        -------
        int
            The value of the last inserted primary key.
        """
        cols = ", ".join(data.keys())
        placeholders = ", ".join(f"%s" for _ in data)
        values = list(data.values())

        sql = f"INSERT INTO {table} ({cols}) VALUES ({placeholders}) RETURNING *;"
        self.cursor.execute(sql, values)
        row = self.cursor.fetchone()
        # If you want *only* the PK you can change this to RETURNING id, etc.
        return row[next(iter(row))]  # first column value – usually the PK

    def insert_many(self, table: str, rows: Iterable[Dict[str, Any]], *, fetch_first: bool = False) -> List[int]:
        """
        Bulk insert – very fast thanks to execute_values.

        Parameters
        ----------
        table : str
            Table name.
        rows : iterable of dict
            Each dict represents one row.
        fetch_first : bool
            If True, the first row’s PK will be returned as a list of one element.
            Useful for a quick “INSERT … RETURNING …” when you only care about the first id.

        Returns
        -------
        list[int]
            List of primary‑key values of the inserted rows.
        """
        if not rows:
            return []

        keys = rows[0].keys()
        columns = ", ".join(keys)
        values = [tuple(row[k] for k in keys) for row in rows]
        placeholder = f"({', '.join(['%s'] * len(keys))})"

        sql = f"INSERT INTO {table} ({columns}) VALUES %s RETURNING *;"
        psycopg2.extras.execute_values(
            self.cursor,
            sql,
            values,
            template=placeholder,
            fetch=True,
        )
        # Return the PK column of every inserted row (first column by convention)
        return [r[next(iter(r))] for r in self.cursor.fetchall()]

    def get(self, table: str, pk_column: str, pk_value: Any, *, columns: Optional[List[str]] = None) -> Optional[Row]:
        """
        Retrieve a single row by its primary key.

        Parameters
        ----------
        table : str
            Table name.
        pk_column : str
            Primary‑key column name.
        pk_value : Any
            Value of the primary key.
        columns : list[str] | None
            If supplied, only those columns will be selected.

        Returns
        -------
        dict | None
            The row as a dictionary, or None if not found.
        """
        cols = ", ".join(columns) if columns else "*"
        sql = f"SELECT {cols} FROM {table} WHERE {pk_column} = %s LIMIT 1;"
        self.cursor.execute(sql, (pk_value,))
        return self.cursor.fetchone()

    def get_all(
        self,
        table: str,
        *,
        columns: Optional[List[str]] = None,
        where: Optional[str] = None,
        where_args: Optional[Iterable[Any]] = None,
        order_by: Optional[str] = None,
        limit: Optional[int] = None,
        offset: Optional[int] = None,
    ) -> List[Row]:
        """
        Retrieve many rows with optional filtering, sorting and pagination.

        Parameters
        ----------
        table : str
            Table name.
        columns : list[str] | None
            List of columns to fetch. If None → SELECT *.
        where : str | None
            Raw `WHERE` clause *without* the leading `WHERE`. Use `%s` for placeholders.
        where_args : iterable | None
            Values that match the placeholders in `where`.
        order_by : str | None
            Raw `ORDER BY` clause (no leading `ORDER BY`).
        limit : int | None
            Max rows to return.
        offset : int | None
            Skip the first N rows.

        Returns
        -------
        list[dict]
            List of rows as dictionaries.
        """
        cols = ", ".join(columns) if columns else "*"
        sql = f"SELECT {cols} FROM {table}"
        params: List[Any] = []

        if where:
            sql += f" WHERE {where}"
            if where_args:
                params.extend(where_args)

        if order_by:
            sql += f" ORDER BY {order_by}"

        if limit is not None:
            sql += " LIMIT %s"
            params.append(limit)

        if offset is not None:
            sql += " OFFSET %s"
            params.append(offset)

        self.cursor.execute(sql, params)
        return self.cursor.fetchall()

    def update(self, table: str, pk_column: str, pk_value: Any, data: Dict[str, Any]) -> bool:
        """
        Update one row by its primary key.

        Parameters
        ----------
        table : str
            Table name.
        pk_column : str
            Primary‑key column name.
        pk_value : Any
            Primary‑key value.
        data : dict
            Column/value pairs to update.

        Returns
        -------
        bool
            True if a row was modified, False otherwise.
        """
        if not data:
            return False

        set_clause = ", ".join(f"{col} = %s" for col in data)
        params = list(data.values()) + [pk_value]
        sql = f"UPDATE {table} SET {set_clause} WHERE {pk_column} = %s;"
        self.cursor.execute(sql, params)
        return self.cursor.rowcount > 0

    def delete(self, table: str, pk_column: str, pk_value: Any) -> bool:
        """
        Delete a row by its primary key.

        Returns
        -------
        bool
            True if a row was deleted.
        """
        sql = f"DELETE FROM {table} WHERE {pk_column} = %s;"
        self.cursor.execute(sql, (pk_value,))
        return self.cursor.rowcount > 0

#
# # ----------------------------------------------------------------------- #
# # Demo usage
# # ----------------------------------------------------------------------- #
# def _demo():
#     db = os.getenv("POSTGRES_DB", "testdb")
#     user = os.getenv("POSTGRES_USER", "postgres")
#     pwd = os.getenv("POSTGRES_PASSWORD", "")
#     host = os.getenv("POSTGRES_HOST", "localhost")
#     port = int(os.getenv("POSTGRES_PORT", "5432"))
#
#     with PostgreSQLDB(
#         host=host, database=db, user=user, password=pwd, port=port
#     ) as db:
#
#         # 1. Create table
#         db.create_table(
#             "users",
#             [
#                 ("id", "SERIAL", ""),          # Identity / auto‑increment
#                 ("username", "TEXT", "NOT NULL"),
#                 ("email", "TEXT", "UNIQUE"),
#                 ("age", "INTEGER", "CHECK (age >= 0)"),
#             ],
#             primary_key="id",
#         )
#
#         # 2. Insert
#         user_id = db.insert(
#             "users",
#             {"username": "alice", "email": "alice@example.com", "age": 25},
#         )
#         print(f"Inserted user id={user_id}")
#
#         # 3. Bulk insert
#         bulk_ids = db.insert_many(
#             "users",
#             [
#                 {"username": "bob", "email": "bob@example.com", "age": 30},
#                 {"username": "carol", "email": "carol@example.com", "age": 22},
#             ],
#         )
#         print(f"Bulk inserted ids: {bulk_ids}")
#
#         # 4. Get single
#         user = db.get("users", "id", user_id)
#         print("Fetched user:", user)
#
#         # 5. Get many
#         all_users = db.get_all(
#             "users",
#             where="age >= %s",
#             where_args=[20],
#             order_by="age DESC",
#             limit=10,
#         )
#         print("All users:", all_users)
#
#         # 6. Update
#         updated = db.update("users", "id", user_id, {"age": 26, "email": "alice26@example.com"})
#         print("Update status:", updated)
#
#         # 7. Delete
#         deleted = db.delete("users", "id", user_id)
#         print("Deleted:", deleted)
#
# # ----------------------------------------------------------------------- #
# # Entry point for `python -m` usage
# # ----------------------------------------------------------------------- #
# if __name__ == "__main__":
#     _demo()