from __future__ import annotations

from contextlib import contextmanager

import data


class PostgresServiceError(Exception):
    pass


def _load_driver():
    try:
        import psycopg
        return "psycopg", psycopg
    except Exception:
        pass
    try:
        import psycopg2
        return "psycopg2", psycopg2
    except Exception as e:
        raise PostgresServiceError("Nie je nainstalovany PostgreSQL driver psycopg ani psycopg2.") from e


def is_configured(conn: data.PostgresConnection) -> bool:
    return bool(
        conn.host.strip()
        and conn.database.strip()
        and conn.user.strip()
        and int(conn.port or 0) > 0
    )


def connection_kwargs(conn: data.PostgresConnection) -> dict:
    kwargs = {
        "host": conn.host,
        "port": int(conn.port or 5432),
        "dbname": conn.database,
        "user": conn.user,
        "password": conn.password,
        "connect_timeout": int(conn.connect_timeout or 5),
    }
    if conn.sslmode:
        kwargs["sslmode"] = conn.sslmode
    return kwargs


def test_connection(conn: data.PostgresConnection) -> None:
    if not conn.enabled:
        raise PostgresServiceError("PostgreSQL pripojenie nie je povolene pre instalaciu.")
    if not is_configured(conn):
        raise PostgresServiceError("PostgreSQL pripojenie nie je vyplnene.")
    driver_name, driver = _load_driver()
    try:
        pg = driver.connect(**connection_kwargs(conn))
        try:
            cur = pg.cursor()
            try:
                cur.execute("SELECT 1")
                cur.fetchone()
            finally:
                cur.close()
        finally:
            pg.close()
    except Exception as e:
        raise PostgresServiceError(f"PostgreSQL spojenie zlyhalo ({driver_name}): {e}") from e


def open_connection(conn: data.PostgresConnection):
    if not conn.enabled:
        raise PostgresServiceError("PostgreSQL pripojenie nie je povolene pre instalaciu.")
    if not is_configured(conn):
        raise PostgresServiceError("PostgreSQL pripojenie nie je vyplnene.")
    driver_name, driver = _load_driver()
    try:
        return driver.connect(**connection_kwargs(conn))
    except Exception as e:
        raise PostgresServiceError(f"PostgreSQL spojenie zlyhalo ({driver_name}): {e}") from e


@contextmanager
def connect(conn: data.PostgresConnection):
    if not conn.enabled:
        raise PostgresServiceError("PostgreSQL pripojenie nie je povolene pre instalaciu.")
    if not is_configured(conn):
        raise PostgresServiceError("PostgreSQL pripojenie nie je vyplnene.")
    driver_name, driver = _load_driver()
    try:
        pg = driver.connect(**connection_kwargs(conn))
    except Exception as e:
        raise PostgresServiceError(f"PostgreSQL spojenie zlyhalo ({driver_name}): {e}") from e
    try:
        yield pg
    finally:
        pg.close()