#!/usr/bin/env python3
"""import_full_stocks.py — 导入全量A股+港股列表到stocks表

数据来源：深交所/上交所公开列表（通过akshare或腾讯API）
运行：python3 import_full_stocks.py
"""
import sys, json, time, urllib.request
from pathlib import Path

DB_PATH = Path("/home/hmo/MoFin/data/mofin.db")

def fetch_tencent_batch(codes):
    """腾讯批量查询股票名称"""
    url = f"http://qt.gtimg.cn/q={','.join(codes)}"
    try:
        req = urllib.request.Request(url, headers={"User-Agent": "Mozilla/5.0"})
        proxy = urllib.request.ProxyHandler({})
        opener = urllib.request.build_opener(proxy)
        with opener.open(req, timeout=15) as r:
            text = r.read().decode("gbk")
        results = {}
        for line in text.strip().split("\n"):
            if "~" not in line:
                continue
            parts = line.split("~")
            name_part = parts[0] if parts else ""
            code = ""
            m = __import__('re').search(r'_(sh|sz|hk)(\d+)', name_part)
            if m:
                code = m.group(2)
            name = parts[1] if len(parts) > 1 else ""
            market = parts[2] if len(parts) > 2 else ""
            if code and name:
                results[code] = (name, market)
        return results
    except Exception as e:
        print(f"  腾讯API错误: {e}", file=sys.stderr)
        return {}

def main():
    import sqlite3
    conn = sqlite3.connect(str(DB_PATH))
    
    # 获取已有代码
    existing = set(r[0] for r in conn.execute("SELECT code FROM stocks").fetchall())
    print(f"当前stocks表已有: {len(existing)}只")
    
    # 生成待查询的A股代码范围（深市000/001/002/003/300/301，沪市600/601/603/605/688/689）
    prefixes = {
        "深市A": [f"{i:03d}" for i in range(0, 10)],  # 000-009
        "深市中小": [f"{i:03d}" for i in range(10, 50)],  # 010-049→实际用001/002
        "深市创业": [f"{i:03d}" for i in range(300, 302)],  # 300-301→实际用300
        "沪市A": [f"{i:03d}" for i in range(600, 606)],  # 600-605
        "沪市科创": [f"{i:03d}" for i in range(688, 690)],  # 688-689
    }
    
    # 实际代码规则：深市000/001/002/003/300/301，沪市600/601/603/605/688
    code_ranges = []
    for prefix in ["000", "001", "002", "003", "300", "301"]:
        for suffix in range(1, 1000):
            code_ranges.append(f"{prefix}{suffix:03d}")
    for prefix in ["600", "601", "603", "605", "688"]:
        for suffix in range(1, 1000):
            code_ranges.append(f"{prefix}{suffix:03d}")
    
    print(f"待查代码总量: {len(code_ranges)}")
    
    # 分批查询（每批30个）
    batch_size = 30
    new_count = 0
    for i in range(0, len(code_ranges), batch_size):
        batch = code_ranges[i:i+batch_size]
        # 过滤已存在的
        batch = [c for c in batch if c not in existing]
        if not batch:
            continue
        
        symbols = []
        for c in batch:
            if c.startswith(("5", "6", "9")):
                symbols.append(f"sh{c}")
            else:
                symbols.append(f"sz{c}")
        
        results = fetch_tencent_batch(symbols)
        for code, (name, market) in results.items():
            if code not in existing:
                try:
                    conn.execute(
                        "INSERT OR IGNORE INTO stocks (code, name) VALUES (?, ?)",
                        (code, name)
                    )
                    new_count += 1
                    existing.add(code)
                except Exception:
                    pass
        
        if (i // batch_size) % 50 == 0:
            print(f"  进度: {i}/{len(code_ranges)}, 新增{new_count}")
    
    conn.commit()
    total = conn.execute("SELECT COUNT(*) FROM stocks").fetchone()[0]
    print(f"\n完成: 新增{new_count}, 总{total}只")
    conn.close()

if __name__ == "__main__":
    main()