#!/usr/bin/env python3
"""
直接在容器内导入数据，从容器外复制parquet内容通过base64
"""

import sys
import base64
import io
import pandas as pd
from vnpy.trader.object import BarData
from vnpy.trader.constant import Exchange, Interval
from vnpy.trader.database import get_database

def main():
    print("🚀 [IMPORT] 开始导入 510300.SSE 日线数据")
    
    # 读取base64编码的parquet数据
    if len(sys.argv) > 1:
        b64_data = sys.argv[1]
        binary_data = base64.b64decode(b64_data)
        buffer = io.BytesIO(binary_data)
        df = pd.read_parquet(buffer)
    else:
        print("❌ 请提供base64编码的parquet数据")
        return
    
    print(f"✅ [IMPORT] 读取完成，共 {len(df)} 行")
    print(f"   时间范围: {df['trade_date'].min()} ~ {df['trade_date'].max()}")
    
    # 转换为BarData
    print(f"🔧 [IMPORT] 转换为 BarData...")
    bars = []
    for idx, row in df.iterrows():
        bar = BarData(
            symbol="510300",
            exchange=Exchange.SSE,
            interval=Interval.DAILY,
            datetime=row["trade_date"],
            open_price=row["open"],
            high_price=row["high"],
            low_price=row["low"],
            close_price=row["close"],
            volume=row["volume"],
            turnover=row["amount"],
            gateway_name="DATA"
        )
        bars.append(bar)
    
    print(f"✅ [IMPORT] 转换完成，共 {len(bars)} 个BarData")
    
    # 写入数据库
    print(f"💾 [IMPORT] 写入数据库...")
    db = get_database()
    
    # 保存数据
    db.save_bar_data(bars)
    print(f"✅ [IMPORT] 保存完成，共写入 {len(bars)} 条")
    
    # 验证
    from datetime import datetime
    start_target = datetime(2021, 1, 1)
    end_target = datetime(2026, 3, 1)
    target_bars = db.load_bar_data("510300", Exchange.SSE, Interval.DAILY, start_target, end_target)
    print(f"✅ [IMPORT] 验证目标区间 2021-01-01 ~ 2026-03-01: 找到 {len(target_bars)} 条数据")
    
    if len(target_bars) > 0:
        print(f"   第一条: {target_bars[0].datetime}")
        print(f"   最后一条: {target_bars[-1].datetime}")
        print(f"   ✅ 满足需求！导入成功！")
    
    print("\n🎉 [IMPORT] 导入完成！")

if __name__ == '__main__':
    main()