import csv
import os
import sqlite3
import sys
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")
DB_PATH = os.path.join(os.path.dirname(__file__), "db", "futures.db")

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 is_trading_day(exchange="CZCE"):
    """检查指定交易所今天是否开市"""
    today = datetime.now().strftime("%Y%m%d")
    rows = tushare_query(
        "trade_cal",
        params={"exchange": exchange, "start_date": today, "end_date": today},
        fields="cal_date,is_open",
    )
    if rows:
        return rows[0]["is_open"] == "1"
    # 查不到（未来日期未更新）时保守放行
    return True


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 is_trading_day(exchange="CZCE"):
    """检查指定交易所今天是否开市"""
    today = datetime.now().strftime("%Y%m%d")
    rows = tushare_query(
        "trade_cal",
        params={"exchange": exchange, "start_date": today, "end_date": today},
        fields="cal_date,is_open",
    )
    if rows:
        return rows[0]["is_open"] == "1"
    return True  # 查不到时保守放行


# ── 合约列表 ──────────────────────────────────

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 get_all_contracts(exchange, fut_code):
    return tushare_query(
        "fut_basic",
        params={"exchange": exchange, "fut_code": fut_code, "fut_type": "1"},
        fields="ts_code,symbol,name,list_date,delist_date",
    )


# ── 增量更新 CSV ──────────────────────────────

def update_contract_csv(fut_code, contract):
    """
    增量更新单份合约 CSV。
    从已有数据的最后交易日重新拉取（含该日），实现补漏 + 修正。
    返回新增/修改的行，或 None（无变化）。
    """
    ts_code = contract["ts_code"]
    csv_name = ts_code.split(".")[0] + ".csv"
    out_dir = os.path.join(DATA_DIR, fut_code)
    csv_path = os.path.join(out_dir, csv_name)
    os.makedirs(out_dir, exist_ok=True)

    # 读取已有数据，获取最后交易日
    existing = []
    if os.path.exists(csv_path):
        with open(csv_path) as f:
            reader = csv.DictReader(f)
            existing = list(reader)

    last_date = existing[-1]["trade_date"] if existing else None

    # 从最后交易日开始拉取（包含该日，以便修正）
    params = {"ts_code": ts_code}
    if last_date:
        params["start_date"] = last_date

    rows = tushare_query("fut_daily", params=params, fields=",".join(FUT_DAILY_FIELDS))
    if not rows:
        return None

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

    if not existing:
        # 全新合约
        with open(csv_path, "w", newline="") as f:
            w = csv.DictWriter(f, fieldnames=FUT_DAILY_FIELDS)
            w.writeheader()
            w.writerows(rows)
        return rows

    if rows[0]["trade_date"] == last_date:
        # 最后交易日可能有修正 → 替换最后一行再追加
        existing = existing[:-1]
        merged = existing + rows
    elif rows[0]["trade_date"] > last_date:
        # 纯新增数据
        merged = existing + rows
    else:
        merged = existing + rows  # 安全兜底

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

    return rows


def update_all_contracts(fut_code, contracts):
    """批量增量更新，返回所有有变化的行"""
    all_new = []
    for c in contracts:
        code = c["ts_code"].split(".")[0]
        new_rows = update_contract_csv(fut_code, c)
        if new_rows is None:
            print(f"  {code:12s}  无新数据")
        else:
            print(f"  {code:12s}  +{len(new_rows)} 条 ({new_rows[0]['trade_date']} ~ {new_rows[-1]['trade_date']})")
            all_new.extend(new_rows)
        time.sleep(0.3)
    return all_new


# ── 全量拉取（初始化用） ─────────────────────

def fetch_all_csv(fut_code, contracts):
    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"]
        csv_name = ts_code.split(".")[0] + ".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)


# ── 数据库 ────────────────────────────────────

def init_db():
    conn = sqlite3.connect(DB_PATH)
    conn.execute("""
        CREATE TABLE IF NOT EXISTS daily (
            ts_code     TEXT NOT NULL,
            trade_date  TEXT NOT NULL,
            pre_close   REAL, pre_settle REAL,
            open REAL, high REAL, low REAL, close REAL, settle REAL,
            change1 REAL, change2 REAL,
            vol REAL, amount REAL, oi REAL, oi_chg REAL,
            delv_settle REAL,
            PRIMARY KEY (ts_code, trade_date)
        )
    """)
    conn.execute("""
        CREATE TABLE IF NOT EXISTS contracts (
            ts_code TEXT PRIMARY KEY,
            symbol TEXT, name TEXT, exchange TEXT,
            fut_code TEXT, list_date TEXT, delist_date TEXT
        )
    """)
    conn.commit()
    conn.close()


def sync_contracts_to_db(fut_code, exchange, contracts):
    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()
    for c in contracts:
        cursor.execute(
            "INSERT OR REPLACE INTO contracts VALUES (?, ?, ?, ?, ?, ?, ?)",
            (c["ts_code"], c["symbol"], c["name"], exchange, fut_code,
             c["list_date"], c.get("delist_date")),
        )
    conn.commit()
    conn.close()
    print(f"  contracts 表同步 {len(contracts)} 条")


def sync_rows_to_db(rows):
    """只同步有变化的行到 daily 表（增量更新用）"""
    if not rows:
        return
    conn = sqlite3.connect(DB_PATH)
    placeholders = ",".join("?" for _ in FUT_DAILY_FIELDS)
    cols = ",".join(FUT_DAILY_FIELDS)
    sql = f"INSERT OR REPLACE INTO daily ({cols}) VALUES ({placeholders})"
    vals = [[r.get(c) for c in FUT_DAILY_FIELDS] for r in rows]
    conn.executemany(sql, vals)
    conn.commit()
    conn.close()
    print(f"  daily 表更新 {len(rows)} 条")


def sync_all_csv_to_db():
    """全量同步所有 CSV 到 daily 表（初始化用）"""
    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()
    for root, _, files in os.walk(DATA_DIR):
        for fname in sorted(files):
            if not fname.endswith(".csv"):
                continue
            fpath = os.path.join(root, fname)
            with open(fpath) as f:
                reader = csv.DictReader(f)
                rows = list(reader)
            if not rows:
                continue
            placeholders = ",".join("?" for _ in FUT_DAILY_FIELDS)
            cols = ",".join(FUT_DAILY_FIELDS)
            sql = f"INSERT OR REPLACE INTO daily ({cols}) VALUES ({placeholders})"
            vals = [[r.get(c) for c in FUT_DAILY_FIELDS] for r in rows]
            cursor.executemany(sql, vals)
            print(f"  {fname:12s}  {len(rows)} 条")
    conn.commit()
    conn.close()


# ── 主入口 ────────────────────────────────────

FG_EXTRA_CODES = {f"FG{suffix}.ZCE" for suffix in
                  [f"14{i:02d}" for i in range(1, 13)] +
                  [f"15{i:02d}" for i in range(1, 13)] +
                  [f"16{i:02d}" for i in range(1, 13)] +
                  [f"17{i:02d}" for i in range(1, 13)] +
                  [f"18{i:02d}" for i in range(1, 13)] +
                  [f"19{i:02d}" for i in range(1, 13)] +
                  [f"20{i:02d}" for i in range(1, 13)] +
                  [f"21{i:02d}" for i in range(1, 13)] +
                  [f"22{i:02d}" for i in range(1, 13)] +
                  [f"23{i:02d}" for i in range(1, 13)] +
                  [f"24{i:02d}" for i in range(1, 13)] +
                  [f"25{i:02d}" for i in range(1, 13)] +
                  [f"26{i:02d}" for i in range(1, 8)]}


if __name__ == "__main__":
    full_init = "--init" in sys.argv

    init_db()
    active = get_active_contracts("CZCE", "FG")

    if full_init:
        all_fg = get_all_contracts("CZCE", "FG")
        active_codes = {c["ts_code"] for c in active}
        extra = [c for c in all_fg if c["ts_code"] in FG_EXTRA_CODES and c["ts_code"] not in active_codes]
        contracts = active + extra
        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','-')}")
        fetch_all_csv("FG", contracts)
        sync_contracts_to_db("FG", "CZCE", contracts)
        sync_all_csv_to_db()
    else:
        if not is_trading_day():
            print("今天非交易日，无需更新")
            print("完成")
            sys.exit(0)

        print(f"增量更新: {len(active)} 个活跃合约")
        new_rows = update_all_contracts("FG", active)
        if new_rows:
            sync_contracts_to_db("FG", "CZCE", active)
            sync_rows_to_db(new_rows)
        else:
            print("  无新数据")

    print("完成")