import requests
import sqlite3
import json
from typing import List


# 板块参数，根据 EastMoney API 中 cmd 字段区分
PLATE_MAP = {
    "沪深A股": "C._A",
    "上证A股": "C._A2",
    "深证A股": "C._A3"
}


DB_NAME = "eastmoney_stock.db"


def create_database():
    """初始化数据库"""
    conn = sqlite3.connect(DB_NAME)
    cursor = conn.cursor()
    cursor.execute('''
        CREATE TABLE IF NOT EXISTS stock_info (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            stock_code TEXT,
            stock_name TEXT,
            latest_price REAL,
            change_percent REAL,
            change_amount REAL,
            volume TEXT,
            turnover TEXT,
            amplitude REAL,
            high REAL,
            low REAL,
            open_price REAL,
            yesterday_close REAL,
            UNIQUE(stock_code)
        )
    ''')
    conn.commit()
    return conn


def fetch_stock_data(cmd: str, page: int = 1, page_size: int = 50) -> List[dict]:
    """ 拉取单页股票数据 """
    url = (
        "https://push2.eastmoney.com/api/qt/clist/get?"
        f"pn={page}&pz={page_size}&po=1&np=1&fltt=2&invt=2&fid=f3"
        f"&fs=b:MK0021&fields=f12,f14,f2,f3,f4,f5,f6,f7,f8,f9,f10,f11,f18,f15,f16,f17,f23"
        f"&cmd={cmd}"
    )

    headers = {
        "User-Agent": "Mozilla/5.0"
    }

    resp = requests.get(url, headers=headers)
    resp.encoding = "utf-8"

    data = json.loads(resp.text)

    if data.get("data") and data["data"].get("diff"):
        return data["data"]["diff"]
    return []


def save_to_database(stock_list: List[dict], conn):
    """ 保存数据进数据库，避免重复 """
    cursor = conn.cursor()
    for stock in stock_list:
        cursor.execute('''
            INSERT OR REPLACE INTO stock_info 
            (stock_code, stock_name, latest_price, change_percent, change_amount, volume, turnover, amplitude, high, low, open_price, yesterday_close)
            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
        ''', (
            stock.get("f12"), stock.get("f14"), stock.get("f2"), stock.get("f3"),
            stock.get("f4"), stock.get("f5"), stock.get("f6"), stock.get("f7"),
            stock.get("f15"), stock.get("f16"), stock.get("f17"), stock.get("f18")
        ))
    conn.commit()


def display(conn):
    cursor = conn.cursor()
    cursor.execute("SELECT stock_code, stock_name, latest_price, change_percent FROM stock_info LIMIT 50")
    rows = cursor.fetchall()
    print(f"{'股票代码':<10} {'股票名称':<10} {'最新价':<10} {'涨跌幅%':<10}")
    for row in rows:
        print(f"{row[0]:<10} {row[1]:<10} {row[2]:<10} {row[3]:<10}")


def main():
    conn = create_database()
    max_pages = 3  # ------- 可自行调整翻页数（每页50条） -------

    print("开始抓取数据 ...")

    for plate_name, cmd in PLATE_MAP.items():
        print(f"📌 抓取板块：{plate_name}")
        for page in range(1, max_pages + 1):
            stock_data = fetch_stock_data(cmd, page)
            if not stock_data:
                break
            save_to_database(stock_data, conn)
            print(f"✅ {plate_name} - 第 {page} 页数据完成！")

    print("\n✅ 数据抓取完毕！前50条展示如下：")
    display(conn)
    conn.close()


if __name__ == "__main__":
    main()