import pandas as pd
import requests
from bs4 import BeautifulSoup
import json
import pymysql
from pymysql.cursors import DictCursor
import os
import time

def load_csv(file_path):
    """
    加载CSV文件并返回DataFrame。
    
    Args:
        file_path (str): CSV文件的路径
    
    Returns:
        pd.DataFrame: 从CSV文件加载的数据
    """
    return pd.read_csv(file_path)

def load_excel(file_path, sheet_name=0):
    """
    加载Excel文件并返回DataFrame。
    
    Args:
        file_path (str): Excel文件的路径
        sheet_name (int, str, optional): 要加载的工作表名称或索引，默认为0
    
    Returns:
        pd.DataFrame: 从Excel文件加载的数据
    """
    return pd.read_excel(file_path, sheet_name=sheet_name)

def load_json(file_path):
    """
    加载JSON文件并返回DataFrame。
    
    Args:
        file_path (str): JSON文件的路径
    
    Returns:
        pd.DataFrame: 从JSON文件加载的数据
    """
    with open(file_path, 'r', encoding='utf-8') as f:
        data = json.load(f)
    return pd.DataFrame(data)

def load_from_api(url, params=None):
    """
    从API获取数据并返回DataFrame。
    
    Args:
        url (str): API的URL地址
        params (dict, optional): 发送到API的参数，默认为None
    
    Returns:
        pd.DataFrame: 从API获取的数据
    
    Raises:
        requests.exceptions.HTTPError: 如果API请求失败
    """
    response = requests.get(url, params=params)
    response.raise_for_status()
    return pd.DataFrame(response.json())

def load_from_web(url, table_selector):
    """
    从网页表格获取数据并返回DataFrame。
    
    Args:
        url (str): 网页的URL地址
        table_selector (str): 用于选择表格的CSS选择器
    
    Returns:
        pd.DataFrame: 从网页表格获取的数据
    
    Raises:
        requests.exceptions.HTTPError: 如果网页请求失败
    """
    response = requests.get(url)
    response.raise_for_status()
    soup = BeautifulSoup(response.content, 'html.parser')
    table = soup.select_one(table_selector)
    return pd.read_html(str(table))[0]

def get_db_config():
    """
    获取数据库配置信息。
    
    从环境变量 DB_INFO 指定的配置文件中读取数据库连接信息。
    
    Returns:
        dict: 包含数据库配置信息的字典
    
    Raises:
        ValueError: 如果环境变量 DB_INFO 未设置
        FileNotFoundError: 如果配置文件不存在
        PermissionError: 如果没有权限读取配置文件
        json.JSONDecodeError: 如果配置文件不是有效的JSON格式
    """
    # 从环境变量获取配置文件路径，不设置默认值
    CONFIG_PATH = os.getenv('DB_INFO')
    if not CONFIG_PATH:
        # 环境变量未设置时的提示
        error_msg = "环境变量 DB_INFO 未设置。请按照以下方式配置：\n"
        error_msg += "1. 环境变量配置：\n"
        error_msg += "   - Windows (图形界面): 设置环境变量 DB_INFO 为配置文件路径\n"
        error_msg += "   - Windows (PowerShell): 执行 $env:DB_INFO='配置文件路径'\n"
        error_msg += "   - Linux/Mac: 执行 export DB_INFO=配置文件路径\n"
        raise ValueError(error_msg)
    
    try:
        # 尝试读取配置文件
        with open(CONFIG_PATH) as f:
            return json.load(f)
    except FileNotFoundError:
        # 文件不存在时的提示
        error_msg = f"配置文件 {CONFIG_PATH} 不存在。请确保文件路径正确。\n"
        raise FileNotFoundError(error_msg)
    except PermissionError:
        # 权限不足时的提示
        raise PermissionError(f"没有权限读取配置文件 {CONFIG_PATH}")

DB_CONFIG = get_db_config()

def fetch_dataframe(sql, db_name='rpt'):
    """
    从数据库执行SQL查询并返回DataFrame。
    
    Args:
        sql (str): 要执行的SQL查询语句
        db_name (str, optional): 数据库标识，默认为'rpt'
    
    Returns:
        pd.DataFrame: 查询结果的DataFrame，如果没有数据则返回空DataFrame
    
    Raises:
        ValueError: 如果指定的数据库标识不存在
        Exception: 如果执行SQL查询时发生错误
    """
    if db_name not in DB_CONFIG:
        raise ValueError(f"未知数据库标识: {db_name}")
    # 使用配置文件中的连接信息
    conn_params = DB_CONFIG[db_name]
    conn = pymysql.connect(**conn_params)
    
    try:
        with conn.cursor(DictCursor) as cursor:
            cursor.execute(sql)
            res = cursor.fetchall()
            if res:
                df = pd.DataFrame(res)
            else:
                # 如果没有数据，返回空DataFrame
                df = pd.DataFrame()
            return df
    except Exception as e:
        raise e
    finally:
        conn.close()

def fetch_data_from_db(tbl_name, order_sql, db_name, batch_size=20000, batch_no=0):
    """
    从数据库批量获取数据并保存为Parquet文件。
    
    Args:
        tbl_name (str): 表名
        order_sql (str): 排序字段
        db_name (str): 数据库标识
        batch_size (int, optional): 每批获取的记录数，默认为20000
        batch_no (int, optional): 起始批次号，默认为0
    
    Returns:
        int: 实际执行的批次数
    
    Raises:
        ValueError: 如果指定的数据库标识不存在
        Exception: 如果执行SQL查询时发生错误
    """
    if db_name not in DB_CONFIG:
        raise ValueError(f"未知数据库标识: {db_name}")
    # 使用配置文件中的连接信息
    conn_params = DB_CONFIG[db_name]
    conn = pymysql.connect(**conn_params)
    cur = conn.cursor() # 生成游标对象
    offset = batch_no*batch_size #初始偏移量为0
    cur.execute(f'SELECT COUNT(*) FROM {tbl_name}')
    total_records = cur.fetchone()[0]
    spend = 0
    
    rounds = int(total_records/batch_size)+1
    print(f'start to fetch {total_records} rows data...')
    print(f'An estimated {rounds} iterations will be performed.')
    round_time = []
    st = time.time()
    while offset < total_records:
        sql=f"SELECT * FROM {tbl_name} ORDER BY {order_sql} LIMIT {batch_size} OFFSET {offset}" # SQL语句

        cur.execute(sql)
        res = cur.fetchall()
    
        df = pd.DataFrame(res, columns=[i[0] for i in cur.description])
        df.to_parquet(f'{tbl_name}_{batch_no}.pqt')
    
        batch_no += 1
        ed = time.time()
        round_time.append(ed-st)
        spend += ed-st
        remain_rounds = rounds - batch_no
        remain_time = remain_rounds*(sum(round_time)/len(round_time))
        print(f'[{batch_no}/{rounds}]: takes {ed-st:.1f} sec, estimated to be completed in {remain_time} seconds. ')
        st = ed
        offset += len(res)
    
    ed = time.time()
    
    print(f'takes total {spend + ed-st:.1f} sec to fetch all data')
    
    cur.close() # 关闭游标
    conn.close() # 关闭连接
    return rounds

def concat_tbl_data(tbl_name):
    """
    合并多个Parquet文件为一个文件。
    
    自动搜索当前目录中以 tbl_name_ 开头并以 .pqt 结尾的文件，
    将它们合并为一个名为 tbl_name.pqt 的文件。
    
    Args:
        tbl_name (str): 表名，用于搜索和命名文件
    
    Returns:
        None
    """
    import glob
    # 搜索当前目录中以 tbl_name_ 开头并以 .pqt 结尾的文件
    files = glob.glob(f'{tbl_name}_*.pqt')
    rounds = len(files)
    
    if rounds == 0:
        print(f'没有找到 {tbl_name}_*.pqt 文件')
        return
    
    st = time.time()
    df_list = []
    whole_tm = 0
    for i in range(rounds):
        df = pd.read_parquet(f'{tbl_name}_{i}.pqt')
        ed = time.time()
        print(f'[{i}/{rounds}]: {ed-st:.1f} sec to read')
        whole_tm += ed-st
        st = ed
        df_list.append(df)
    
    df_all = pd.concat(df_list,axis=0)
    print(df_all.shape)
    
    df_all.to_parquet(f'{tbl_name}.pqt')
    
    ed = time.time()
    whole_tm += ed-st
    print(f'takes {whole_tm:.1f} sec to do all thing!')
