# coding:utf-8

import requests
import pandas as pd
from sqlalchemy import create_engine
from config import XUEQIU_HEADERS

def fetch_and_store_stock_data(db_url, table_name, page_size=100):
    """
    Fetches stock data from the given API and stores it into a specified database table.
    Supports fetching data for multiple types: 'sha', 'sza', 'kcb'.
    """
    import requests
    import pandas as pd
    from sqlalchemy import create_engine

    base_url = 'https://stock.xueqiu.com/v5/stock/screener/quote/list.json'
    types = ['sha', 'sza', 'kcb']  # 数据类型
    headers = XUEQIU_HEADERS

    all_data = []

    for stock_type in types:
        params = {
            'page': 1,
            'size': page_size,
            'order': 'desc',
            'order_by': 'percent',
            'market': 'CN',
            'type': stock_type
        }

        # 初次请求以获取总页数
        response = requests.get(base_url, headers=headers, params=params)
        if response.status_code != 200:
            print(f"请求 {stock_type} 数据失败，状态码：{response.status_code}")
            continue

        data = response.json()
        total_count = data['data']['count']
        total_pages = (total_count // page_size) + 1

        for page in range(1, total_pages + 1):
            params['page'] = page
            response = requests.get(base_url, headers=headers, params=params)
            if response.status_code == 200:
                data = response.json()
                all_data.extend(data['data']['list'])
            else:
                print(f"请求 {stock_type} 数据第 {page} 页失败，状态码：{response.status_code}")

    # 转换为 DataFrame
    df = pd.DataFrame(all_data)

    if not df.empty:
        # 添加 id 列
        df['id'] = range(1, len(df) + 1)

        # 创建数据库连接
        engine = create_engine(db_url)

        # 将数据写入数据库表
        df.to_sql(table_name, con=engine, if_exists='replace', index=False)
        print(f"成功将数据写入数据库表 {table_name}。")
    else:
        print("未获取到任何数据。")


if __name__ == "__main__":
    # 示例调用
    db_url = 'mysql+pymysql://root:Chlry#$.8@192.168.18.199:3306/db_gp_cj'
    table_name = 'stock_changes'
    fetch_and_store_stock_data(db_url, table_name)