# coding:utf-8
# 更新港股列表的代码

import requests
import pandas as pd
from sqlalchemy import create_engine, text
import sys
import os

# 将项目根目录添加到Python路径，以便导入config
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from src.scripts.config import XUEQIU_HEADERS

def collect_hk_stock_codes(db_url):
    """
    采集雪球港股列表数据，并存储到数据库。
    """
    engine = create_engine(db_url)
    headers = XUEQIU_HEADERS
    base_url = "https://stock.xueqiu.com/v5/stock/screener/quote/list.json"
    page = 1
    page_size = 90
    all_data = []

    print("--- Starting to collect Hong Kong stock codes ---")
    
    # 采集前先清空表
    try:
        with engine.begin() as conn:
            conn.execute(text("TRUNCATE TABLE gp_code_hk"))
        print("Table `gp_code_hk` has been truncated.")
    except Exception as e:
        print(f"Error truncating table `gp_code_hk`: {e}")
        return

    while True:
        params = {
            'page': page,
            'size': page_size,
            'order': 'desc',
            'order_by': 'market_capital',
            'market': 'HK',
            'type': 'hk',
            'is_delay': 'true'
        }
        
        print(f"Fetching page {page}...")
        try:
            response = requests.get(base_url, headers=headers, params=params, timeout=20)
            if response.status_code != 200:
                print(f"Request failed with status code {response.status_code}")
                break
            
            data = response.json()
            if data.get('error_code') != 0:
                print(f"API error: {data.get('error_description')}")
                break

            stock_list = data.get('data', {}).get('list', [])
            if not stock_list:
                print("No more data found. Collection finished.")
                break

            all_data.extend(stock_list)
            
            # 如果获取到的数据少于每页数量，说明是最后一页
            if len(stock_list) < page_size:
                print("Reached the last page. Collection finished.")
                break
            
            page += 1

        except requests.exceptions.RequestException as e:
            print(f"Request exception on page {page}: {e}")
            break
    
    if all_data:
        print(f"--- Collected a total of {len(all_data)} stocks. Preparing to save to database. ---")
        df = pd.DataFrame(all_data)
        
        # 数据映射和转换
        df_to_save = pd.DataFrame()
        df_to_save['gp_name'] = df['name']
        df_to_save['gp_code'] = df['symbol']
        df_to_save['gp_code_two'] = 'HK' + df['symbol'].astype(str)
        df_to_save['market_cap'] = df['market_capital']
        
        try:
            df_to_save.to_sql('gp_code_hk', engine, if_exists='append', index=False)
            print("--- Successfully saved all data to `gp_code_hk`. ---")
        except Exception as e:
            print(f"Error saving data to database: {e}")
    else:
        print("--- No data collected. ---")
    
    engine.dispose()

if __name__ == "__main__":
    db_url = 'mysql+pymysql://root:Chlry#$.8@192.168.18.199:3306/db_gp_cj'
    collect_hk_stock_codes(db_url) 