#!/usr/bin/env python3
"""
Database Connection Test Script

Tests connectivity from the local laptop to the PostgreSQL database
running in Docker on the Synology NAS at 20.20.20.20.

Tests performed:
  1. TCP connectivity to the NAS port
  2. PostgreSQL authentication
  3. Schema inspection (list tables)
  4. Data verification (row counts, latest records)
  5. Integration with the project's db.py module

Usage:
  python testing/test_db_connection.py
  python testing/test_db_connection.py --host 20.20.20.20 --port 5433
  python testing/test_db_connection.py --verbose

Environment variables (or .env file in this directory):
  PG_HOST     - NAS IP address (default: 20.20.20.20)
  PG_PORT     - PostgreSQL port on NAS (default: 5433)
  PG_DB       - Database name (default: hyper)
  PG_USER     - Database user (default: hyper)
  PG_PASSWORD - Database password (default: kaqpaaoi0)
  PG_TIMEOUT  - Connection timeout in seconds (default: 10)
"""

import os
import sys
import socket
import argparse
import logging

try:
    from dotenv import load_dotenv
    _TESTING_DIR = os.path.dirname(os.path.abspath(__file__))
    load_dotenv(os.path.join(_TESTING_DIR, ".env"))
except ImportError:
    pass

_PROJECT_ROOT = os.path.dirname(_TESTING_DIR)
sys.path.insert(0, _PROJECT_ROOT)

PG_HOST = os.environ.get("PG_HOST", "20.20.20.20")
PG_PORT = int(os.environ.get("PG_PORT", "5433"))
PG_DB = os.environ.get("PG_DB", "hyper")
PG_USER = os.environ.get("PG_USER", "hyper")
PG_PASSWORD = os.environ.get("PG_PASSWORD", "kaqpaaoi0")
PG_TIMEOUT = int(os.environ.get("PG_TIMEOUT", "10"))


class TestResult:
    """Track pass/fail results across all test stages."""

    def __init__(self):
        self.passed = 0
        self.failed = 0
        self.results = []

    def add(self, name, success, message):
        self.results.append((name, success, message))
        if success:
            self.passed += 1
        else:
            self.failed += 1

    def summary(self):
        total = self.passed + self.failed
        return (
            f"\n{'=' * 60}\n"
            f"Results: {self.passed}/{total} passed, {self.failed} failed\n"
            f"{'=' * 60}"
        )


def test_tcp_connectivity(host, port, timeout):
    """Stage 1: Verify TCP reachability to the NAS port."""
    try:
        sock = socket.create_connection((host, port), timeout=timeout)
        sock.close()
        return True, f"TCP connection to {host}:{port} succeeded"
    except socket.timeout:
        return False, f"TCP connection to {host}:{port} timed out after {timeout}s"
    except ConnectionRefusedError:
        return False, (
            f"Connection refused to {host}:{port}. "
            "Port may be closed or Synology DSM firewall is blocking it."
        )
    except socket.gaierror:
        return False, f"Could not resolve hostname: {host}"
    except Exception as e:
        return False, f"TCP connection error: {e}"


def test_pg_auth(conn_str, timeout):
    """Stage 2: Verify PostgreSQL authentication with credentials."""
    try:
        import psycopg2
    except ImportError:
        return False, (
            "psycopg2 is not installed. "
            "Install it with: pip install psycopg2-binary"
        )

    try:
        conn = psycopg2.connect(conn_str, connect_timeout=timeout)
        cur = conn.cursor()
        cur.execute("SELECT 1")
        cur.fetchone()
        cur.close()
        conn.close()
        masked = conn_str.split("@")[1] if "@" in conn_str else conn_str
        return True, f"PostgreSQL authentication successful ({masked})"
    except psycopg2.OperationalError as e:
        return False, f"PostgreSQL connection failed: {e}"
    except Exception as e:
        return False, f"PostgreSQL error: {e}"


def test_schema_inspection(conn_str):
    """Stage 3: List tables in the public schema."""
    try:
        import psycopg2
        conn = psycopg2.connect(conn_str)
        cur = conn.cursor()
        cur.execute(
            """
            SELECT table_name
            FROM information_schema.tables
            WHERE table_schema = 'public'
              AND table_type = 'BASE TABLE'
            ORDER BY table_name
            """
        )
        tables = cur.fetchall()
        cur.close()
        conn.close()
        table_names = [t[0] for t in tables]
        sample = table_names[:5] if table_names else "none"
        return True, f"Found {len(table_names)} tables. Sample: {sample}"
    except Exception as e:
        return False, f"Schema inspection failed: {e}"


def test_data_verification(conn_str):
    """Stage 4: Check row count and latest record for a table."""
    try:
        import psycopg2
        conn = psycopg2.connect(conn_str)
        cur = conn.cursor()
        cur.execute(
            """
            SELECT table_name
            FROM information_schema.tables
            WHERE table_schema = 'public'
              AND table_type = 'BASE TABLE'
            ORDER BY table_name
            LIMIT 1
            """
        )
        result = cur.fetchone()
        if not result:
            cur.close()
            conn.close()
            return False, "No tables found in database"

        table_name = result[0]
        cur.execute(f'SELECT COUNT(*) FROM "{table_name}"')
        count = cur.fetchone()[0]
        cur.execute(
            f'SELECT timestamp_ms FROM "{table_name}" '
            f'ORDER BY timestamp_ms DESC LIMIT 1'
        )
        latest = cur.fetchone()
        cur.close()
        conn.close()

        if count == 0:
            return True, f"Table '{table_name}' exists but is empty (0 rows)"
        latest_ts = latest[0] if latest else "N/A"
        return True, f"Table '{table_name}': {count} rows, latest timestamp_ms={latest_ts}"
    except Exception as e:
        return False, f"Data verification failed: {e}"


def test_db_py_integration(conn_str):
    """Stage 5: Verify the project's db.py module works with the remote connection."""
    try:
        import db
    except ImportError as e:
        return False, (
            f"Could not import db.py: {e}. "
            "Ensure the project root is accessible."
        )

    try:
        old_conn_str = db.PG_CONN_STR
        db.PG_CONN_STR = conn_str
        conn = db.get_connection()
        cur = conn.cursor()
        cur.execute("SELECT 1")
        cur.fetchone()
        cur.close()
        conn.close()
        db.PG_CONN_STR = old_conn_str
        return True, "db.py get_connection() works with remote connection string"
    except Exception as e:
        return False, f"db.py integration failed: {e}"


def main():
    parser = argparse.ArgumentParser(
        description="Test PostgreSQL database connection from laptop to NAS Docker container."
    )
    parser.add_argument("--host", default=PG_HOST, help=f"NAS IP address (default: {PG_HOST})")
    parser.add_argument("--port", type=int, default=PG_PORT, help=f"PostgreSQL port (default: {PG_PORT})")
    parser.add_argument("--db", default=PG_DB, help=f"Database name (default: {PG_DB})")
    parser.add_argument("--user", default=PG_USER, help=f"Database user (default: {PG_USER})")
    parser.add_argument("--password", default=PG_PASSWORD, help=f"Database password (default: {PG_PASSWORD})")
    parser.add_argument("--timeout", type=int, default=PG_TIMEOUT, help=f"Connection timeout in seconds (default: {PG_TIMEOUT})")
    parser.add_argument("--verbose", "-v", action="store_true", help="Verbose output")
    parser.add_argument("--skip-db-py", action="store_true", help="Skip db.py integration test")
    args = parser.parse_args()

    host = args.host
    port = args.port
    db_name = args.db
    user = args.user
    password = args.password
    timeout = args.timeout
    conn_str = f"postgresql://{user}:{password}@{host}:{port}/{db_name}"

    log_level = logging.DEBUG if args.verbose else logging.INFO
    logging.basicConfig(level=log_level, format="%(asctime)s - %(levelname)s - %(message)s")

    print(f"\n{'=' * 60}")
    print("Database Connection Test")
    print(f"{'=' * 60}")
    print(f"Target:     {host}:{port}")
    print(f"Database:   {db_name}")
    print(f"User:       {user}")
    print(f"Timeout:    {timeout}s")
    print(f"Connection: postgresql://{user}:***@{host}:{port}/{db_name}")
    print(f"{'=' * 60}\n")

    results = TestResult()

    print("[1/5] Testing TCP connectivity...")
    success, message = test_tcp_connectivity(host, port, timeout)
    print(f"  [{'PASS' if success else 'FAIL'}] {message}")
    results.add("TCP connectivity", success, message)

    print("\n[2/5] Testing PostgreSQL authentication...")
    success, message = test_pg_auth(conn_str, timeout)
    print(f"  [{'PASS' if success else 'FAIL'}] {message}")
    results.add("PostgreSQL authentication", success, message)

    if not success:
        print("\n  Authentication failed. Cannot proceed with further tests.")
        print("  Possible causes:")
        print("  - Wrong credentials")
        print("  - PostgreSQL not running or port not mapped (docker-compose up -d)")
        print("  - Synology DSM firewall blocking port 5433")
        print(f"\n{results.summary()}")
        sys.exit(1)

    print("\n[3/5] Testing schema inspection...")
    success, message = test_schema_inspection(conn_str)
    print(f"  [{'PASS' if success else 'FAIL'}] {message}")
    results.add("Schema inspection", success, message)

    print("\n[4/5] Testing data verification...")
    success, message = test_data_verification(conn_str)
    print(f"  [{'PASS' if success else 'FAIL'}] {message}")
    results.add("Data verification", success, message)

    if not args.skip_db_py:
        print("\n[5/5] Testing db.py integration...")
        success, message = test_db_py_integration(conn_str)
        print(f"  [{'PASS' if success else 'FAIL'}] {message}")
        results.add("db.py integration", success, message)

    print(f"\n{'=' * 60}")
    print("Summary")
    print(f"{'=' * 60}")
    for name, success, _ in results.results:
        print(f"  [{'PASS' if success else 'FAIL'}] {name}")
    print(f"\nTotal: {results.passed}/{results.passed + results.failed} passed")
    print(f"{'=' * 60}\n")

    sys.exit(0 if results.failed == 0 else 1)


if __name__ == "__main__":
    main()