#!/usr/bin/env python3
"""
kis_trader/utils/export_ticks_parquet.py
DB의 ws_ticks 데이터를 Parquet 파일로 변환하여 디스크에 저장합니다.
주기적인 백그라운드 작업이나 On-demand 호출을 위해 설계되었습니다.
"""
import os
import argparse
import sys
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")))
from datetime import datetime, timedelta
import polars as pl

from kis_trader.utils.logger import get_logger
logger = get_logger("export_ticks_parquet")
from database import TradeDB, _DB_HOST, _DB_USER, _DB_PASS, _DB_NAME, _DB_PORT

# 저장 폴더 경로 설정
TICKS_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "data", "ticks"))
os.makedirs(os.path.join(TICKS_DIR, "KR"), exist_ok=True)
os.makedirs(os.path.join(TICKS_DIR, "US"), exist_ok=True)

# Polars 용 URI 생성 (connectorx)
DB_URI = f"mysql://{_DB_USER}:{_DB_PASS}@{_DB_HOST}:{_DB_PORT}/{_DB_NAME}"

def export_day(market: str, day: str) -> bool:
    """특정 일자의 데이터를 Parquet으로 변환합니다 (day: YYYYMMDD)"""
    market = market.strip().upper()
    table = "ws_ticks_us" if market == "US" else "ws_ticks"
    
    start_str = f"{day}000000"
    end_str = f"{day}235959"
    
    # 쿼리 작성 (MariaDB에서 직접 Polars로 로드)
    query = f"""
        SELECT code, tick_time, price, volume, source, recv_ts, tick_time_raw
        FROM {table}
        WHERE market = '{market}'
          AND tick_time >= '{start_str}'
          AND tick_time <= '{end_str}'
    """
    
    out_file = os.path.join(TICKS_DIR, market, f"{day}.parquet")
    
    try:
        logger.info(f"[{market}] {day} 틱 데이터 로딩 중... ({table})")
        # polars의 read_database 사용 (connectorx 백엔드 사용 시 매우 빠름, 기본 adbc)
        # pip install adbc-driver-mysql 를 권장하지만 일단 pandas 폴백 혹은 pyarrow 활용
        df = pl.read_database_uri(query, uri=DB_URI)
        
        if df.height == 0:
            logger.warning(f"[{market}] {day} 데이터가 없습니다.")
            return False
            
        logger.info(f"[{market}] {day} 로딩 완료 ({df.height} rows). Parquet 저장 시작...")
        df.write_parquet(out_file)
        logger.info(f"[{market}] {day} Parquet 저장 완료: {out_file}")
        return True
    except Exception as e:
        logger.error(f"[{market}] {day} Parquet 저장 실패: {e}")
        return False

def main():
    parser = argparse.ArgumentParser(description="틱 데이터 Parquet 익스포터")
    parser.add_argument("--market", type=str, default="KR", help="마켓 (KR/US)")
    parser.add_argument("--start", type=str, help="시작일 (YYYY-MM-DD)")
    parser.add_argument("--end", type=str, help="종료일 (YYYY-MM-DD)")
    
    args = parser.parse_args()
    
    if not args.start or not args.end:
        logger.error("start 와 end 인자가 필요합니다.")
        return

    try:
        start_dt = datetime.strptime(args.start, "%Y-%m-%d")
        end_dt = datetime.strptime(args.end, "%Y-%m-%d")
    except Exception as e:
        logger.error(f"날짜 형식 오류: {e}")
        return
        
    cur_dt = start_dt
    while cur_dt <= end_dt:
        day_str = cur_dt.strftime("%Y%m%d")
        export_day(args.market, day_str)
        cur_dt += timedelta(days=1)

if __name__ == "__main__":
    main()