"""BigQuery client and query operations"""

from typing import List, Dict, Any, Optional, Iterator
from datetime import datetime
from google.cloud import bigquery
from google.cloud.exceptions import NotFound
import time


class BigQueryClient:
    """BigQuery客户端封装"""
    
    def __init__(self, project_id: Optional[str] = None):
        """初始化BigQuery客户端
        
        Args:
            project_id: GCP项目ID，如果为None则使用默认项目
        """
        if project_id:
            self.client = bigquery.Client(project=project_id)
        else:
            self.client = bigquery.Client()
        
        self.job_config = bigquery.QueryJobConfig()
        self.job_config.use_query_cache = True
        self.job_config.use_legacy_sql = False
    
    def execute_query(self, query: str, 
                     parameters: Optional[List[bigquery.ScalarQueryParameter]] = None,
                     timeout: int = 300) -> Iterator[Dict]:
        """执行查询并返回结果迭代器
        
        Args:
            query: SQL查询语句
            parameters: 查询参数列表
            timeout: 查询超时时间（秒）
            
        Yields:
            查询结果行
        """
        job_config = self.job_config
        if parameters:
            job_config.query_parameters = parameters
        
        try:
            query_job = self.client.query(query, job_config=job_config)
            
            # 等待查询完成
            results = query_job.result(timeout=timeout)
            
            # 逐行返回结果
            for row in results:
                yield dict(row)
                
        except Exception as e:
            print(f"✗ 查询执行失败: {e}")
            raise
    
    def execute_query_to_list(self, query: str,
                             parameters: Optional[List[bigquery.ScalarQueryParameter]] = None,
                             timeout: int = 300,
                             max_results: int = 100000) -> List[Dict]:
        """执行查询并返回结果列表
        
        Args:
            query: SQL查询语句
            parameters: 查询参数列表
            timeout: 查询超时时间（秒）
            max_results: 最大结果数量
            
        Returns:
            查询结果列表
        """
        results = []
        count = 0
        
        for row in self.execute_query(query, parameters, timeout):
            results.append(row)
            count += 1
            if count >= max_results:
                print(f"⚠️ 达到最大结果数限制: {max_results}")
                break
        
        return results
    
    def table_exists(self, dataset_id: str, table_id: str) -> bool:
        """检查表是否存在
        
        Args:
            dataset_id: 数据集ID (格式: project.dataset)
            table_id: 表ID
            
        Returns:
            表是否存在
        """
        try:
            # 分离项目和数据集
            if '.' in dataset_id:
                project, dataset_name = dataset_id.split('.', 1)
            else:
                project = self.client.project
                dataset_name = dataset_id
            
            dataset_ref = bigquery.DatasetReference(project, dataset_name)
            table_ref = dataset_ref.table(table_id)
            
            self.client.get_table(table_ref)
            return True
        except NotFound:
            return False
        except Exception as e:
            print(f"检查表存在性失败 {dataset_id}.{table_id}: {e}")
            return False
    
    def get_table_schema(self, dataset_id: str, table_id: str) -> List[Dict]:
        """获取表结构
        
        Args:
            dataset_id: 数据集ID
            table_id: 表ID
            
        Returns:
            表字段列表
        """
        try:
            if '.' in dataset_id:
                project, dataset_name = dataset_id.split('.', 1)
            else:
                project = self.client.project
                dataset_name = dataset_id
            
            dataset_ref = bigquery.DatasetReference(project, dataset_name)
            table_ref = dataset_ref.table(table_id)
            
            table = self.client.get_table(table_ref)
            
            return [
                {
                    'name': field.name,
                    'type': field.field_type,
                    'mode': field.mode,
                    'description': field.description
                }
                for field in table.schema
            ]
        except Exception as e:
            print(f"获取表结构失败 {dataset_id}.{table_id}: {e}")
            return []
    
    def list_tables(self, dataset_id: str, prefix: Optional[str] = None) -> List[str]:
        """列出数据集中的表
        
        Args:
            dataset_id: 数据集ID
            prefix: 表名前缀过滤
            
        Returns:
            表名列表
        """
        try:
            if '.' in dataset_id:
                project, dataset_name = dataset_id.split('.', 1)
            else:
                project = self.client.project
                dataset_name = dataset_id
            
            dataset_ref = bigquery.DatasetReference(project, dataset_name)
            
            tables = []
            for table in self.client.list_tables(dataset_ref):
                if prefix is None or table.table_id.startswith(prefix):
                    tables.append(table.table_id)
            
            return tables
        except Exception as e:
            print(f"列出表失败 {dataset_id}: {e}")
            return []
    
    def get_table_row_count(self, dataset_id: str, table_id: str) -> int:
        """获取表行数
        
        Args:
            dataset_id: 数据集ID
            table_id: 表ID
            
        Returns:
            表行数
        """
        query = f"""
        SELECT COUNT(*) as row_count
        FROM `{dataset_id}.{table_id}`
        """
        
        try:
            results = list(self.execute_query(query))
            if results:
                return results[0].get('row_count', 0)
            return 0
        except Exception:
            return 0
    
    def build_date_filter(self, date_field: str, 
                         start_date: datetime,
                         end_date: datetime) -> str:
        """构建日期过滤条件
        
        Args:
            date_field: 日期字段名
            start_date: 开始日期
            end_date: 结束日期
            
        Returns:
            SQL WHERE子句
        """
        start_str = start_date.strftime('%Y-%m-%d')
        end_str = end_date.strftime('%Y-%m-%d')
        
        return f"""
        DATE({date_field}) >= '{start_str}'
        AND DATE({date_field}) <= '{end_str}'
        """
    
    def estimate_query_cost(self, query: str) -> Dict[str, Any]:
        """估算查询成本
        
        Args:
            query: SQL查询
            
        Returns:
            包含预估字节数和成本的字典
        """
        job_config = bigquery.QueryJobConfig(dry_run=True, use_query_cache=False)
        
        try:
            query_job = self.client.query(query, job_config=job_config)
            
            bytes_processed = query_job.total_bytes_processed
            gb_processed = bytes_processed / (1024**3)
            
            # BigQuery定价: $5 per TB
            estimated_cost = (bytes_processed / (1024**4)) * 5
            
            return {
                'bytes_processed': bytes_processed,
                'gb_processed': round(gb_processed, 2),
                'estimated_cost_usd': round(estimated_cost, 4)
            }
        except Exception as e:
            print(f"估算查询成本失败: {e}")
            return {
                'bytes_processed': 0,
                'gb_processed': 0,
                'estimated_cost_usd': 0
            }