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 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):
    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
    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.execute("""
        CREATE TABLE IF NOT EXISTS spread (
            trade_date TEXT PRIMARY KEY,
            spread REAL,
            main_ts_code TEXT,
            main_chg REAL
        )
    """)
    try:
        conn.execute("ALTER TABLE spread ADD COLUMN main_chg REAL")
    except sqlite3.OperationalError:
        pass
    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):
    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():
    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()


def sync_spread():
    conn = sqlite3.connect(DB_PATH)
    conn.row_factory = sqlite3.Row
    dates = conn.execute("SELECT DISTINCT trade_date FROM daily ORDER BY trade_date").fetchall()

    count = 0
    for d in dates:
        date_str = d["trade_date"]
        rows = conn.execute("""
            SELECT d.ts_code, d.close, d.open, d.vol
            FROM daily d
            JOIN contracts c ON d.ts_code = c.ts_code
            WHERE d.trade_date = ?
              AND c.delist_date > ?
              AND d.oi > 0
            ORDER BY c.delist_date
        """, (date_str, date_str)).fetchall()
        if len(rows) < 2:
            continue

        closes = []
        main_ts = None
        main_vol = -1
        main_chg = None
        for r in rows:
            try:
                closes.append(float(r["close"]))
            except (TypeError, ValueError):
                closes.append(None)
            if r["ts_code"][4:6] in ("01", "05", "09"):
                v = r["vol"]
                if v is not None and v > main_vol:
                    main_vol = v
                    main_ts = r["ts_code"]
                    try:
                        c = float(r["close"])
                        o = float(r["open"])
                        main_chg = c - o
                    except (TypeError, ValueError):
                        main_chg = None

        closes = [c for c in closes if c is not None]
        if len(closes) < 2:
            continue

        spread = closes[-1] - closes[0]
        conn.execute(
            "INSERT OR REPLACE INTO spread VALUES (?, ?, ?, ?)",
            (date_str, spread, main_ts, main_chg),
        )
        count += 1

    conn.commit()
    conn.close()
    print(f"  spread 表更新 {count} 个交易日")


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

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:
            print(f"  {c['ts_code'].split('.')[0]: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()
        sync_spread()
    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("  无新数据")

    sync_spread()
    print("完成")