import logging
import psycopg2
import sqlite3




logger = logging.getLogger(__name__)

class PostgresCursorWrapper:
    def __init__(self, cursor):
        self.cursor = cursor

    def execute(self, query, args=None):
        # Convert SQLite ? placeholders to Postgres %s
        if query:
            # Replace ? with %s
            pg_query = query.replace('?', '%s')

            # Handle BLOB to BYTEA differences if necessary
            # (psycopg2 usually handles bytes objects correctly as BYTEA)
        else:
            pg_query = query

        try:
            return self.cursor.execute(pg_query, args)
        except Exception as e:
            logger.error(f"SQL Error: {e} | Query: {pg_query} | Args: {args}")
            raise e

    def fetchall(self):
        return self.cursor.fetchall()

    def fetchone(self):
        return self.cursor.fetchone()

    def __getattr__(self, name):
        return getattr(self.cursor, name)

class PostgresConnectionWrapper:
    def __init__(self, conn):
        self.conn = conn
        # Enable autocommit or manage it manually?
        # sqlite3 connection in app.py context manager does not seem to set autocommit,
        # but the context manager does commit() or rollback().

    def cursor(self):
        # Use RealDictCursor to mimic sqlite3.Row (dictionary access)
        return PostgresCursorWrapper(self.conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor))

    def commit(self):
        return self.conn.commit()

    def rollback(self):
        return self.conn.rollback()

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

    def __getattr__(self, name):
        return getattr(self.conn, name)

def get_db_connection(database_url):
    """
    Factory function to return a connection object.
    Supports 'sqlite://' (or just file path) and 'postgresql://'.
    """
    if database_url.startswith('postgres://') or database_url.startswith('postgresql://'):
        try:
            conn = psycopg2.connect(database_url)
            return PostgresConnectionWrapper(conn)
        except Exception as e:
            logger.error(f"Failed to connect to Postgres: {e}")
            raise e
    else:
        # Fallback to SQLite
        # If it starts with sqlite:///, strip it
        if database_url.startswith('sqlite:///'):
            db_path = database_url.replace('sqlite:///', '')
        else:
            db_path = database_url

        conn = sqlite3.connect(db_path)
        conn.row_factory = sqlite3.Row
        return conn