# -*- coding: utf-8 -*-
"""Quick utility for validating PostgreSQL connectivity (ASCII-only output)."""
from __future__ import annotations

import argparse
import os
import sys

PROJECT_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
if PROJECT_ROOT not in sys.path:
    sys.path.insert(0, PROJECT_ROOT)

from database.connection import DatabaseConnection


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="PostgreSQL connectivity smoke test")
    parser.add_argument("--dsn", help="Override TEST_DB_DSN / env value")
    parser.add_argument(
        "--query",
        default="SELECT 1 AS ok",
        help="Custom SQL to run after connection (default: SELECT 1 AS ok)",
    )
    parser.add_argument(
        "--timeout",
        type=int,
        default=10,
        help="connect_timeout seconds passed to psycopg2 (capped at 20, default: 10)",
    )
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    dsn = args.dsn or os.environ.get("TEST_DB_DSN")
    if not dsn:
        print("Missing DSN. Use --dsn or TEST_DB_DSN.", file=sys.stderr)
        return 2

    print(f"Trying connection: {dsn}")
    try:
        timeout = max(1, min(args.timeout, 20))
        conn = DatabaseConnection(dsn, connect_timeout=timeout)
    except Exception as exc:  # pragma: no cover - diagnostic output
        print("Connection failed:", exc, file=sys.stderr)
        return 1

    try:
        result = conn.query(args.query)
        print("Connection OK, query result:")
        for row in result:
            print(row)
        conn.close()
        return 0
    except Exception as exc:  # pragma: no cover - diagnostic output
        print("Connection succeeded but query failed:", exc, file=sys.stderr)
        try:
            conn.close()
        finally:
            return 3


if __name__ == "__main__":
    raise SystemExit(main())