#!/usr/bin/env python3
"""
测试Python版本的DuckDB Delta扫描查询功能
对比Go版本的实现
"""

import os
import sys
import json
import logging
from datetime import datetime

# 添加src目录到Python路径
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'src'))

from slowquery_proxy.db import DuckDBStore
from slowquery_proxy.auth import get_s3_creds_by_github_token

# 设置日志
logging.basicConfig(level=logging.INFO,filename="../logs/slowquery_proxy.log",format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

def test_delta_scan_query():
    """测试Delta扫描查询，使用您提供的SQL"""
    
    # 设置环境变量
    os.environ["APP_ENV"] = "prod"
    os.environ["LOG_CRED_DEBUG"] = "true"
    
    store = DuckDBStore()
    
    try:
        # 设置S3认证（使用us-east-1区域）
        region = "us-east-1"
        store.set_aws_s3(region)
        logger.info(f"S3认证设置完成，区域: {region}")
        
        # 您提供的测试SQL
        test_sql = """
        SELECT  query 
        FROM delta_scan( 
            's3://o11y-prod-shared-us-east-1/deltalake/1372813089209061633/ed869d81-91c5-4a87-a28d-20795c738d3f/slowlogs' 
        ) 
        Where "time" >= 1749826800 AND "time" < 1749827800 
        LIMIT 1;
        """
        
        logger.info(f"执行SQL查询: {test_sql.strip()}")
        
        # 执行查询
        start_time = datetime.now()
        rows, meta = store.execute_sql(test_sql)
        end_time = datetime.now()
        
        # 输出结果
        logger.info(f"查询完成，耗时: {(end_time - start_time).total_seconds():.2f}秒")
        logger.info(f"返回行数: {meta['row_count']}")
        logger.info(f"列名: {meta['columns']}")
        logger.info(f"执行时间: {meta['elapsed_ms']}毫秒")
        
        if rows:
            logger.info("查询结果:")
            for i, row in enumerate(rows):
                logger.info(f"行 {i+1}: {json.dumps(row, ensure_ascii=False, indent=2)}")
        else:
            logger.info("没有返回任何数据")
            
        return True
        
    except Exception as e:
        logger.error(f"查询失败: {str(e)}")
        import traceback
        traceback.print_exc()
        return False
        
    finally:
        store.close()

def test_s3_credentials():
    """测试S3认证功能"""
    try:
        creds = get_s3_creds_by_github_token()
        logger.info(f"S3认证成功:")
        logger.info(f"  Access Key ID: {creds.access_key_id[:4]}****")
        logger.info(f"  有Session Token: {'是' if creds.session_token else '否'}")
        logger.info(f"  过期时间: {creds.expiration}")
        return True
    except Exception as e:
        logger.error(f"S3认证失败: {str(e)}")
        return False

if __name__ == "__main__":
    logger.info("开始测试Python DuckDB Delta扫描查询...")
    
    # 首先测试S3认证
    if not test_s3_credentials():
        logger.error("S3认证测试失败，无法继续")
        sys.exit(1)
    
    # 然后测试Delta扫描查询
    success = test_delta_scan_query()
    
    if success:
        logger.info("测试通过！Python版本可以正常执行Delta扫描查询")
    else:
        logger.error("测试失败！Python版本无法正常执行Delta扫描查询")
        sys.exit(1)