from datetime import datetime, timezone
from flask import current_app
from databricks import sql


class DatabricksQuery:
    def __init__(self):
        """纯Python连接Databricks SQL，无JVM依赖"""
        try:
            # 读取配置（参数和原代码一致，无需修改）
            self.host = current_app.config["DATABRICKS_HOST"]
            self.token = current_app.config["DATABRICKS_TOKEN"]
            self.cluster_id = current_app.config["CLUSTER_ID"]

            # 校验参数
            if not all([self.host, self.token, self.cluster_id]):
                raise ValueError(
                    "必须配置DATABRICKS_HOST、DATABRICKS_TOKEN、CLUSTER_ID"
                )

            print("Databricks SQL 连接器初始化成功")
        except Exception as e:
            print(f"初始化失败: {e}")
            raise

    # 复用原有参数校验逻辑
    @staticmethod
    def _validate_table_name(table_name):
        if not table_name or not isinstance(table_name, str):
            raise ValueError(f"无效的表名: {table_name}")
        if ";" in table_name:
            raise ValueError(f"表名包含非法字符 ';': {table_name}")
        parts = table_name.split(".")
        if len(parts) != 3:
            raise ValueError(f"表名格式应为 catalog.schema.table: {table_name}")

    @staticmethod
    def _validate_vin(vin):
        if vin is None:
            return True
        if not isinstance(vin, str):
            raise ValueError(f"无效的 VIN: {vin} (必须是17位字母数字)")
        return True

    def query_table(
        self, table_name, vin=None, start_time=None, end_time=None,limit=None
    ):
        """纯Python查询表数据，无JVM依赖"""
        try:
            # 验证参数
            self._validate_table_name(table_name)
            self._validate_vin(vin)

            # 构建SQL查询语句（替代原PySpark DataFrame操作）
            select_fields = """
                Incident__Name AS incident_name,
                Incident__description AS incident_description,
                Meaning AS meaning,
                Signal__Name As signal_name,
                incident_time,
                oneid,
                speed,
                int_value,
                idc_tickcount_ms,
                latitude,
                longitude,
                mux_data,
                gps_heading,
                float_value
            """
            
            sql_query = f"SELECT {select_fields} FROM {table_name}"
            where_conditions = []

            # 拼接过滤条件
            if vin:
                where_conditions.append(f"oneid = '{vin}'")
            if start_time:
                start_dt = self._parse_iso_time(start_time).isoformat()
                where_conditions.append(f"incident_time >= '{start_dt}'")
            if end_time:
                end_dt = self._parse_iso_time(end_time).isoformat()
                where_conditions.append(f"incident_time <= '{end_dt}'")
            
            # 拼接WHERE子句
            if where_conditions:
                sql_query += " WHERE " + " AND ".join(where_conditions)
            # 拼接LIMIT
            if limit is not None:
                sql_query += f" LIMIT {limit}"
            
            # sql_query="select * from hive_metastore.taf_level_two_plus.accident_report_data WHERE oneid = 'LE4LG4GB6SL256937' limit 1"
            print(f"查询sql:{sql_query}")
            
            # 执行SQL（纯Python，无JVM）
            with sql.connect(
                server_hostname=self.host,
                http_path=f"sql/protocolv1/o/3035612432650494/{self.cluster_id}",  # 注意路径格式
                access_token=self.token,
            ) as connection:
                with connection.cursor() as cursor:
                    cursor.execute(sql_query)
                    # 获取列名和数据，转换为字典格式（和原代码返回格式一致）
                    columns = [desc[0] for desc in cursor.description]
                    results = []
                    for row in cursor.fetchall():
                        results.append(dict(zip(columns, row)))
                return results
            
        except Exception as e:
            print(f"查询失败: {str(e)}")
            raise

    # 复用原有时间解析逻辑
    def _parse_iso_time(self, time_str):
        try:
            if time_str.endswith("Z"):
                time_str = time_str[:-1] + "+00:00"
            dt = datetime.fromisoformat(time_str)
            if dt.tzinfo is None:
                dt = dt.replace(tzinfo=timezone.utc)
            return dt
        except Exception as e:
            raise ValueError(f"无效的时间格式 '{time_str}': {str(e)}")

    def stop(self):
        # 纯Python连接器无需关闭SparkSession，此处保留方法以兼容原有调用逻辑
        print("Databricks SQL 连接已释放")