import csv
import os
import time
from datetime import datetime

import requests

TOKEN = "76efd8465f9f2591aa42a385268e06acf6b80b7a15be2267ad2281b7"
API_URL = "https://api.tushare.pro"
DATA_DIR = os.path.join(os.path.dirname(__file__), "data")

FUT_DAILY_FIELDS = [
    "ts_code", "trade_date", "pre_close", "pre_settle",
    "open", "high", "low", "close", "settle",
    "change1", "change2", "vol", "amount", "oi", "oi_chg",
    "delv_settle",
]


def tushare_query(api_name, params=None, fields=None):
    req = {"api_name": api_name, "token": TOKEN}
    if params:
        req["params"] = params
    if fields:
        req["fields"] = fields

    resp = requests.post(API_URL, json=req)
    data = resp.json()

    if data["code"] != 0:
        raise Exception(f"API error ({data['code']}): {data['msg']}")

    fields_list = data["data"]["fields"]
    items = data["data"]["items"]
    return [dict(zip(fields_list, item)) for item in items]


def get_active_contracts(exchange, fut_code):
    """获取指定交易所、品种仍在交易的合约列表"""
    contracts = tushare_query(
        "fut_basic",
        params={"exchange": exchange, "fut_code": fut_code, "fut_type": "1"},
        fields="ts_code,symbol,name,list_date,delist_date",
    )
    today = datetime.now().strftime("%Y%m%d")
    return [c for c in contracts if not c.get("delist_date") or c["delist_date"] > today]


def save_contract_csv(fut_code, contracts):
    """对每个活跃合约拉取日线行情并写入 data/{fut_code}/{code}.csv"""
    out_dir = os.path.join(DATA_DIR, fut_code)
    os.makedirs(out_dir, exist_ok=True)

    for c in contracts:
        ts_code = c["ts_code"]  # e.g. FG2609.ZCE
        csv_name = ts_code.split(".")[0] + ".csv"  # FG2609.csv
        csv_path = os.path.join(out_dir, csv_name)

        rows = tushare_query(
            "fut_daily",
            params={"ts_code": ts_code},
            fields=",".join(FUT_DAILY_FIELDS),
        )
        if not rows:
            print(f"  {csv_name:12s} 无数据，跳过")
            continue

        rows.sort(key=lambda r: r["trade_date"])

        with open(csv_path, "w", newline="") as f:
            w = csv.DictWriter(f, fieldnames=FUT_DAILY_FIELDS)
            w.writeheader()
            w.writerows(rows)

        print(f"  {csv_name:12s}  {len(rows)} 条记录 ({rows[0]['trade_date']} ~ {rows[-1]['trade_date']})")

        time.sleep(0.3)  # tushare 频率限制


if __name__ == "__main__":
    contracts = get_active_contracts("CZCE", "FG")
    print(f"活跃玻璃合约共 {len(contracts)} 个:")

    for c in contracts:
        code = c["ts_code"].split(".")[0]
        print(f"  {code:8s}  上市:{c['list_date']}  退市:{c.get('delist_date','-')}")

    print("\n拉取日线行情...")
    save_contract_csv("FG", contracts)
    print("完成")