"""通达信「沪深京日线数据完整包」全量导入 candles 表。

用法（在 backend 目录下）：
    uv run python scripts/import_tdx_day.py C:/Users/cirry/Downloads/hsjday [symbol ...]
    # symbol 为可选的 6 位代码过滤（如 000001 002671），只重导这些标的
    uv run python scripts/import_tdx_day.py <目录> --no-clear
    # --no-clear：不清空任何行，纯 upsert（用于给已导入的底座回补 amount 成交额）

- 解析 vipdoc 的 .day 二进制文件（每条 32 字节）：
    日期(YYYYMMDD) 开 高 低 收（×100） 成交额(元, float32) 成交量(股) 保留
- 只导入 stock_basic 里登记的股票（自动排除指数/基金/可转债/回购）；
  sh000001(上证指数) 与 sz000001(平安银行) 这类代码冲突也由此化解。
- 价格为**不复权**：全量模式导入前清空已有的非 DEMO 行情；指定 symbol 过滤时
  只清空这些标的（用于修复被复权口径污染的个别股票），其余不动。
- amount 为 TDX 原生 float32（元），精度 ~6 位有效数字，展示用途足够；
  ON CONFLICT 时仅更新 amount 列，不动 OHLCV/turnover（避免与换手率回补互相干扰）。
- 写入用 asyncpg execute_many + ON CONFLICT DO UPDATE，可重复执行（幂等）。
"""
from __future__ import annotations

import argparse
import asyncio
import struct
import sys
import time
from pathlib import Path

import asyncpg

# .env 里的 DATABASE_URL 是 SQLAlchemy 格式，asyncpg 需要 libpq 格式
DEFAULT_URL = "postgresql://postgres:postgres@localhost:5432/stock"
BATCH = 20_000  # 每批 upsert 行数


def load_db_url() -> str:
    env = Path(__file__).resolve().parent.parent / ".env"
    if env.exists():
        for line in env.read_text(encoding="utf-8").splitlines():
            line = line.strip()
            if line.startswith("DATABASE_URL=postgresql+asyncpg://"):
                return "postgresql://" + line.split("://", 1)[1]
    return DEFAULT_URL


def parse_day_file(path: Path) -> list[tuple[int, float, float, float, float, float, float]]:
    """解析单个 .day 文件 -> [(date, open, high, low, close, volume(股), amount(元)), ...]"""
    raw = path.read_bytes()
    unpack = struct.Struct("<IIIIIfII").unpack_from
    out = []
    for i in range(len(raw) // 32):
        date, o, h, l, c, amount, vol, _reserved = unpack(raw, i * 32)
        out.append((date, o / 100.0, h / 100.0, l / 100.0, c / 100.0, float(vol), float(amount)))
    return out


async def main(root: Path, symbols: list[str] | None = None, no_clear: bool = False) -> None:
    if not root.exists():
        sys.exit(f"目录不存在: {root}")

    conn = await asyncpg.connect(load_db_url())
    try:
        # 股票清单：ts_code 形如 000001.SZ，用于过滤指数/基金/转债
        rows = await conn.fetch("SELECT ts_code, symbol FROM stock_basic WHERE list_status = 'L'")
        by_exchange: dict[str, set[str]] = {"sh": set(), "sz": set(), "bj": set()}
        for r in rows:
            suffix = r["ts_code"].split(".")[-1].lower()  # SH/SZ/BJ -> sh/sz/bj
            if suffix in by_exchange:
                by_exchange[suffix].add(r["symbol"])
        print(f"stock_basic 在市股票: " + ", ".join(f"{k}={len(v)}" for k, v in by_exchange.items()))

        files = sorted(root.glob("*/lday/*.day"))
        print(f"发现 .day 文件: {len(files)} 个")

        if no_clear:
            print("--no-clear：不清空任何行，纯 upsert 回补 amount")
        elif symbols:
            # 清空旧行情（保留 DEMO 合成数据），避免 qfq/不复权混用；
            # 带 symbol 过滤时只清空目标标的（修复个别被污染的股票，不动其余底座）
            deleted = await conn.execute(
                "DELETE FROM candles WHERE symbol = ANY($1)", symbols
            )
            print(f"清空目标标的 {symbols}: {deleted}")
            keep = set(symbols)
            files = [p for p in files if p.name[2:8] in keep]
            print(f"过滤后待导入 .day 文件: {len(files)} 个")
        else:
            deleted = await conn.execute("DELETE FROM candles WHERE symbol <> 'DEMO'")
            print(f"清空旧行情: {deleted}")

        if no_clear:
            # 回填模式：只写 amount，不动 OHLCV/turnover（底座已就位，避免全表重写）
            upsert_sql = """
                INSERT INTO candles (symbol, timeframe, ts, open, high, low, close, volume, amount)
                VALUES ($1, '1d', to_timestamp($2::text, 'YYYYMMDD')::timestamp, $3, $4, $5, $6, $7, $8)
                ON CONFLICT (symbol, timeframe, ts) DO UPDATE
                SET amount = EXCLUDED.amount
            """
        else:
            upsert_sql = """
                INSERT INTO candles (symbol, timeframe, ts, open, high, low, close, volume, amount)
                VALUES ($1, '1d', to_timestamp($2::text, 'YYYYMMDD')::timestamp, $3, $4, $5, $6, $7, $8)
                ON CONFLICT (symbol, timeframe, ts) DO UPDATE
                SET open = EXCLUDED.open, high = EXCLUDED.high, low = EXCLUDED.low,
                    close = EXCLUDED.close, volume = EXCLUDED.volume, amount = EXCLUDED.amount
            """

        t0 = time.time()
        total_stocks = 0
        skipped = 0
        batch: list[tuple] = []

        rows_done = 0

        async def flush() -> None:
            nonlocal batch, rows_done
            if batch:
                await conn.executemany(upsert_sql, batch)
                rows_done += len(batch)
                batch = []

        for n, path in enumerate(files, 1):
            market = path.name[:2].lower()  # sh / sz / bj
            code = path.name[2:8]
            if code not in by_exchange.get(market, set()):
                skipped += 1
                continue
            for date, o, h, l, c, v, amount in parse_day_file(path):
                batch.append((code, str(date), o, h, l, c, v, amount))
            total_stocks += 1
            if len(batch) >= BATCH:
                await flush()
            if n % 500 == 0:
                elapsed = time.time() - t0
                print(f"  进度 {n}/{len(files)} 文件, 已入库 {total_stocks} 只股票, "
                      f"{rows_done + len(batch):,} 行, {elapsed:.0f}s")
        await flush()

        cnt = await conn.fetchval("SELECT count(*) FROM candles WHERE symbol <> 'DEMO'")
        span = await conn.fetchrow(
            "SELECT min(ts) AS lo, max(ts) AS hi FROM candles WHERE symbol <> 'DEMO'"
        )
        with_amt = await conn.fetchval(
            "SELECT count(*) FROM candles WHERE symbol <> 'DEMO' AND amount IS NOT NULL"
        )
        print(f"\n完成: {total_stocks} 只股票, {cnt:,} 行日线, "
              f"范围 {span['lo']:%Y-%m-%d} ~ {span['hi']:%Y-%m-%d}, "
              f"含成交额 {with_amt:,} 行, "
              f"跳过非股票文件 {skipped} 个, 耗时 {time.time() - t0:.0f}s")
    finally:
        await conn.close()


if __name__ == "__main__":
    ap = argparse.ArgumentParser(description="TDX 沪深京日线全量导入 candles")
    ap.add_argument("root", help="hsjday 目录（其下 */lday/*.day）")
    ap.add_argument("symbols", nargs="*", help="可选的 6 位代码过滤")
    ap.add_argument("--no-clear", action="store_true",
                    help="不清空任何行，纯 upsert（amount 回补模式）")
    a = ap.parse_args()
    asyncio.run(main(Path(a.root), a.symbols or None, a.no_clear))