#!/usr/bin/env python3
"""
Script để test Cloudflare Tunnel URL và kiểm tra dữ liệu PostgreSQL
"""

import sys
import psycopg2
from pathlib import Path

# Colors
GREEN = '\033[0;32m'
YELLOW = '\033[1;33m'
RED = '\033[0;31m'
BLUE = '\033[0;34m'
NC = '\033[0m'  # No Color

def log(msg):
    print(f"{BLUE}[TEST]{NC} {msg}")

def success(msg):
    print(f"{GREEN}✅{NC} {msg}")

def warn(msg):
    print(f"{YELLOW}⚠️{NC} {msg}")

def error(msg):
    print(f"{RED}❌{NC} {msg}")

def test_connection(host, port, database, user, password):
    """Test kết nối PostgreSQL"""
    log(f"Đang test kết nối: {host}:{port}/{database}")
    
    try:
        conn = psycopg2.connect(
            host=host,
            port=port,
            database=database,
            user=user,
            password=password,
            connect_timeout=10
        )
        success("Kết nối thành công!")
        
        cursor = conn.cursor()
        
        # Test 1: PostgreSQL version
        cursor.execute("SELECT version();")
        version = cursor.fetchone()[0]
        log(f"PostgreSQL version: {version[:50]}...")
        
        # Test 2: Kiểm tra database
        cursor.execute("SELECT current_database();")
        db_name = cursor.fetchone()[0]
        success(f"Đang kết nối database: {db_name}")
        
        # Test 3: Kiểm tra số lượng tables
        cursor.execute("""
            SELECT COUNT(*) 
            FROM information_schema.tables 
            WHERE table_schema = 'public'
        """)
        table_count = cursor.fetchone()[0]
        success(f"Số lượng tables: {table_count}")
        
        # Test 4: Kiểm tra dữ liệu trong các bảng chính
        tables_to_check = [
            'core_legaldocument',
            'core_fine',
            'core_procedure',
            'core_office',
            'core_advisory',
        ]
        
        print(f"\n{BLUE}📊 Kiểm tra dữ liệu:{NC}")
        print("=" * 60)
        
        total_docs = 0
        for table in tables_to_check:
            try:
                cursor.execute(f"SELECT COUNT(*) FROM {table};")
                count = cursor.fetchone()[0]
                total_docs += count
                if count > 0:
                    success(f"{table}: {count:,} records")
                else:
                    warn(f"{table}: 0 records (trống)")
            except Exception as e:
                error(f"{table}: Lỗi - {e}")
        
        print("=" * 60)
        success(f"Tổng số documents: {total_docs:,}")
        
        # Test 5: Kiểm tra một vài records mẫu
        if total_docs > 0:
            print(f"\n{BLUE}📄 Sample records:{NC}")
            try:
                cursor.execute("""
                    SELECT id, title, created_at 
                    FROM core_legaldocument 
                    ORDER BY created_at DESC 
                    LIMIT 5
                """)
                records = cursor.fetchall()
                for idx, (doc_id, title, created_at) in enumerate(records, 1):
                    title_short = (title[:50] + "...") if title and len(title) > 50 else (title or "N/A")
                    print(f"  {idx}. [{doc_id}] {title_short}")
                    print(f"     Created: {created_at}")
            except Exception as e:
                warn(f"Không thể lấy sample records: {e}")
        
        cursor.close()
        conn.close()
        
        return True, total_docs
        
    except psycopg2.OperationalError as e:
        error(f"Lỗi kết nối: {e}")
        return False, 0
    except Exception as e:
        error(f"Lỗi: {e}")
        return False, 0

def main():
    print(f"{BLUE}{'='*60}{NC}")
    print(f"{BLUE}Test Cloudflare Tunnel và kiểm tra dữ liệu PostgreSQL{NC}")
    print(f"{BLUE}{'='*60}{NC}\n")
    
    # Lấy thông tin từ user
    print("Nhập thông tin Cloudflare Tunnel:")
    tunnel_host = input("  Host (ví dụ: xyz.trycloudflare.com hoặc xyz.cfargotunnel.com): ").strip()
    tunnel_port = input("  Port (ví dụ: 5432): ").strip() or "5432"
    
    # Database credentials
    db_user = input("  Database user (mặc định: hue_remote): ").strip() or "hue_remote"
    db_password = input("  Database password (mặc định: huepass123): ").strip() or "huepass123"
    db_name = input("  Database name (mặc định: hue_portal): ").strip() or "hue_portal"
    
    print(f"\n{BLUE}🔍 Đang test kết nối...{NC}\n")
    
    # Test connection
    success_conn, doc_count = test_connection(
        tunnel_host,
        int(tunnel_port),
        db_name,
        db_user,
        db_password
    )
    
    if success_conn:
        print(f"\n{BLUE}{'='*60}{NC}")
        success("KẾT NỐI THÀNH CÔNG!")
        if doc_count > 0:
            success(f"Có {doc_count:,} documents trong database")
            print(f"\n{BLUE}📝 DATABASE_URL để dùng:{NC}")
            print(f"postgres://{db_user}:{db_password}@{tunnel_host}:{tunnel_port}/{db_name}")
        else:
            warn("Database trống, chưa có dữ liệu")
        print(f"{BLUE}{'='*60}{NC}")
    else:
        print(f"\n{BLUE}{'='*60}{NC}")
        error("KHÔNG KẾT NỐI ĐƯỢC!")
        print(f"\n{BLUE}💡 Kiểm tra:{NC}")
        print("  1. Cloudflare Tunnel URL có đúng không?")
        print("  2. Tunnel có đang chạy trên server không?")
        print("  3. PostgreSQL có đang chạy không?")
        print("  4. Firewall có block port không?")
        print(f"{BLUE}{'='*60}{NC}")
        sys.exit(1)

if __name__ == "__main__":
    main()