import taosrest
import pandas as pd
import logging
from td_config import TDENGINE_CONFIG

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

def connect_to_tdengine():
    """
    连接到TDengine数据库
    """
    try:
        client = taosrest.RestClient(
            url=TDENGINE_CONFIG['url'],
            user=TDENGINE_CONFIG['user'],
            password=TDENGINE_CONFIG['password'],
            database=TDENGINE_CONFIG['database'],
            timeout=TDENGINE_CONFIG['timeout']
        )
        logger.info(f"成功连接到TDengine服务器 {TDENGINE_CONFIG['url']}")
        return client
    except Exception as e:
        logger.error(f"连接到TDengine时出错: {str(e)}")
        raise

def execute_query(conn, sql):
    """
    执行SQL查询并返回结果
    """
    try:
        logger.info(f"执行查询: {sql}")
        # 使用sql方法执行查询
        result = conn.sql(sql)
        
        # 将结果转换为DataFrame
        if result and isinstance(result, dict) and 'data' in result:
            data = result['data']
            if data:
                df = pd.DataFrame(data)
                logger.info(f"查询成功，返回{len(df)}行数据")
                return df
        
        logger.info("查询成功，但未返回数据")
        return pd.DataFrame()
            
    except Exception as e:
        logger.error(f"执行查询时出错: {str(e)}")
        raise

def get_database_info(conn):
    """
    获取数据库信息
    """
    try:
        super_tables_df = execute_query(conn, "SHOW STABLES")
        tables_df = execute_query(conn, "SHOW TABLES")
        db_info = execute_query(conn, f"SHOW {TDENGINE_CONFIG['database']}")
        
        return {
            "super_tables": super_tables_df,
            "tables": tables_df,
            "database_info": db_info
        }
    except Exception as e:
        logger.error(f"获取数据库信息时出错: {str(e)}")
        raise

def get_latest_timeseries_data(conn, table_name, limit=1):
    """
    获取指定表的最新时序数据
    """
    try:
        sql = f"SELECT * FROM {table_name} ORDER BY ts DESC LIMIT {limit}"
        return execute_query(conn, sql)
    except Exception as e:
        logger.error(f"获取最新时序数据时出错: {str(e)}")
        raise

def test_connection():
    """测试TDengine连接并执行基本查询"""
    try:
        conn = connect_to_tdengine()
        
        # 测试服务器版本
        version_df = execute_query(conn, "SELECT SERVER_VERSION()")
        version = version_df.iloc[0, 0] if not version_df.empty else "未知"
        logger.info(f"TDengine服务器版本: {version}")
        
        # # 测试数据库列表
        # status_df = execute_query(conn, "SHOW DATABASES")
        # logger.info(f"找到 {len(status_df)} 个数据库:")
        # for idx, row in status_df.iterrows():
        #     logger.info(f"  - {row['name']}")
        
        # # 测试当前数据库的表
        # if TDENGINE_CONFIG['database']:
        #     tables_df = execute_query(conn, "SHOW TABLES")
        #     logger.info(f"在数据库 '{TDENGINE_CONFIG['database']}' 中找到 {len(tables_df)} 个表")
        #     if not tables_df.empty:
        #         for idx, row in tables_df.iterrows():
        #             logger.info(f"  - {row['table_name']} (创建于 {row['created_time']})")
        
        # logger.info("连接测试完成")
        return True
    except Exception as e:
        logger.error(f"连接测试失败: {str(e)}")
        return False

if __name__ == "__main__":
    print("=" * 60)
    print(f"TDengine连接测试 - {TDENGINE_CONFIG['url']}")
    print("=" * 60)
    
    if test_connection():
        print("\n✓ 连接测试成功!")
    else:
        print("\n✗ 连接测试失败，请检查连接参数和错误日志。")
    
    print("\n提示: 要使用此模块，请先安装TDengine REST客户端:")
    print("pip install taosrest") 