"""
KRX 히스토리컬 데이터 수집 스크립트

KRX OpenAPI (data-dbg.krx.co.kr)를 사용하여 KOSPI/KOSDAQ 종목의
일봉 OHLCV 데이터를 DB에 저장합니다.

- KIS API 의존성 완전 제거
- 날짜 1건 호출 → 해당일 전 종목 OHLCV 일괄 수신 (효율적)

사용법:
    python collect_krx_data.py                            # 전체 수집 (2020-01-01~)
    python collect_krx_data.py --limit 5                  # 최근 5거래일 (테스트)
    python collect_krx_data.py --start-date 20240101      # 시작일 지정
    python collect_krx_data.py --update                   # 증분 업데이트
    python collect_krx_data.py --market kospi             # KOSPI만
"""

import argparse
import asyncio
import logging
import os
import sys
import time
from datetime import date, datetime, timedelta
from pathlib import Path

import aiohttp
import pandas as pd

# 프로젝트 루트를 Python 경로에 추가
sys.path.insert(0, str(Path(__file__).parent))

from sqlalchemy import text
from sqlalchemy.dialects.postgresql import insert as pg_insert

from app.core.database import AsyncSessionLocal, engine, Base
from app.models.models import StockDailyPrice

# ─── KRX OpenAPI 설정 ─────────────────────────────────────────────
# 실제 데이터 엔드포인트 (openapi.krx.co.kr 는 포털, 데이터는 data-dbg)
KRX_DATA_BASE = "https://data-dbg.krx.co.kr/svc/apis"

try:
    from app.core.config import settings as _cfg
    KRX_API_KEY: str = _cfg.KRX_API_KEY
except Exception:
    from dotenv import load_dotenv
    load_dotenv(Path(__file__).parent.parent / ".env")
    KRX_API_KEY = os.environ.get("KRX_API_KEY", "")

# ─── ETF/특수 종목 제외 키워드 ────────────────────────────────────
SKIP_KEYWORDS = [
    "ETF", "레버리지", "인버스", "KODEX", "TIGER", "ARIRANG",
    "KBSTAR", "HANARO", "SOL", "ACE", "KOSEF", "KINDEX",
    "선물", "스팩", "리츠", "SPAC",
]

# 날짜별 동시 요청 수 (KRX API 쓰로틀링 고려)
COLLECT_CONCURRENCY = 3
REQUEST_INTERVAL = 0.5   # 요청 간 최소 대기(초)

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
)
logger = logging.getLogger(__name__)


# ─── 유틸 ─────────────────────────────────────────────────────────

def generate_trading_dates(start: str, end: str) -> list[str]:
    """start ~ end 사이 평일(월~금) 날짜 문자열 목록 생성 (YYYYMMDD)."""
    d = datetime.strptime(start, "%Y%m%d").date()
    e = datetime.strptime(end, "%Y%m%d").date()
    dates: list[str] = []
    while d <= e:
        if d.weekday() < 5:   # 0=Mon … 4=Fri
            dates.append(d.strftime("%Y%m%d"))
        d += timedelta(days=1)
    return dates


def _int(val) -> int:
    """문자열/숫자를 int로 안전 변환 (쉼표 포함 처리)."""
    try:
        return int(str(val).replace(",", "").strip() or "0")
    except (ValueError, TypeError):
        return 0


def _float(val) -> float:
    """문자열/숫자를 float로 안전 변환."""
    try:
        return float(str(val).replace(",", "").strip() or "0")
    except (ValueError, TypeError):
        return 0.0


# ─── KRX API 호출 ─────────────────────────────────────────────────

async def _fetch_krx(
    session: aiohttp.ClientSession,
    endpoint: str,
    bas_dd: str,
) -> list:
    """KRX OpenAPI 단건 호출 → OutBlock_1 리스트 반환.

    endpoint 예: "sto/stk_bydd_trd"
    """
    url = f"{KRX_DATA_BASE}/{endpoint}"
    try:
        async with session.get(
            url,
            params={"basDd": bas_dd},
            headers={"AUTH_KEY": KRX_API_KEY},
            timeout=aiohttp.ClientTimeout(total=30),
        ) as resp:
            if resp.status != 200:
                logger.warning("KRX %s %s → HTTP %s", endpoint, bas_dd, resp.status)
                return []
            data = await resp.json(content_type=None)
            return data.get("OutBlock_1") or []
    except Exception as exc:
        logger.error("KRX API 오류 [%s %s]: %s", endpoint, bas_dd, exc)
        return []


async def fetch_daily_prices(
    session: aiohttp.ClientSession,
    market: str,
    bas_dd: str,
) -> pd.DataFrame:
    """특정 날짜의 전 종목 OHLCV 조회.

    KRX 응답 필드:
        ISU_SRT_CD  : 단축 종목코드 (e.g. "005930")
        ISU_ABBRV   : 종목 약칭
        TDD_OPNPRC  : 시가
        TDD_HGPRC   : 고가
        TDD_LWPRC   : 저가
        TDD_CLSPRC  : 종가
        ACC_TRDVOL  : 누적 거래량
        FLUC_RT     : 등락률 (%)
    """
    endpoint = "sto/stk_bydd_trd" if market == "kospi" else "sto/ksq_bydd_trd"
    items = await _fetch_krx(session, endpoint, bas_dd)
    if not items:
        return pd.DataFrame()

    trade_date = datetime.strptime(bas_dd, "%Y%m%d").date()
    rows: list[dict] = []

    for item in items:
        name = (item.get("ISU_ABBRV") or item.get("ISU_NM") or "").strip()
        if any(kw in name for kw in SKIP_KEYWORDS):
            continue

        code = item.get("ISU_SRT_CD", "").strip()
        if not code:
            continue

        volume = _int(item.get("ACC_TRDVOL"))
        if volume == 0:
            continue   # 거래정지일 제외

        close = _int(item.get("TDD_CLSPRC"))
        if close == 0:
            continue

        rows.append({
            "stock_code":  code,
            "stock_name":  name,
            "date":        trade_date,
            "open":        _int(item.get("TDD_OPNPRC")),
            "high":        _int(item.get("TDD_HGPRC")),
            "low":         _int(item.get("TDD_LWPRC")),
            "close":       close,
            "volume":      volume,
            "change_rate": _float(item.get("FLUC_RT")) / 100.0,
        })

    return pd.DataFrame(rows) if rows else pd.DataFrame()


async def get_stock_codes(
    session: aiohttp.ClientSession,
    market: str,
    bas_dd: str,
) -> set[str]:
    """특정 날짜 기준 상장 종목 코드 집합 조회 (선택적 필터용)."""
    endpoint = "sto/stk_isu_base_info" if market == "kospi" else "sto/ksq_isu_base_info"
    items = await _fetch_krx(session, endpoint, bas_dd)
    codes: set[str] = set()
    for item in items:
        code = item.get("ISU_SRT_CD", "").strip()
        if code:
            codes.add(code)
    return codes


# ─── DB 헬퍼 ──────────────────────────────────────────────────────

async def create_tables():
    async with engine.begin() as conn:
        await conn.run_sync(Base.metadata.create_all)
    logger.info("DB 테이블 확인/생성 완료")


async def get_last_trade_date(session) -> date | None:
    """DB 내 마지막 거래일 조회 (전 종목 공통)."""
    result = await session.execute(
        text("SELECT MAX(trade_date) FROM stock_daily_prices")
    )
    return result.scalar()


async def upsert_daily_prices(session, df: pd.DataFrame) -> int:
    """DB에 일봉 데이터 UPSERT (배치 500건)."""
    if df.empty:
        return 0

    records = []
    for _, row in df.iterrows():
        trade_date = row["date"]
        if isinstance(trade_date, pd.Timestamp):
            trade_date = trade_date.date()
        records.append({
            "stock_code":  row["stock_code"],
            "stock_name":  row["stock_name"],
            "trade_date":  trade_date,
            "open_price":  int(row["open"]),
            "high_price":  int(row["high"]),
            "low_price":   int(row["low"]),
            "close_price": int(row["close"]),
            "volume":      int(row["volume"]),
            "change_rate": float(row["change_rate"]) if pd.notna(row["change_rate"]) else None,
            "market_cap":  None,
        })

    batch_size = 500
    inserted = 0
    for i in range(0, len(records), batch_size):
        batch = records[i : i + batch_size]
        stmt = pg_insert(StockDailyPrice).values(batch)
        stmt = stmt.on_conflict_do_update(
            constraint="uq_stock_daily",
            set_={
                "stock_name":  stmt.excluded.stock_name,
                "open_price":  stmt.excluded.open_price,
                "high_price":  stmt.excluded.high_price,
                "low_price":   stmt.excluded.low_price,
                "close_price": stmt.excluded.close_price,
                "volume":      stmt.excluded.volume,
                "change_rate": stmt.excluded.change_rate,
                "market_cap":  stmt.excluded.market_cap,
            },
        )
        await session.execute(stmt)
        inserted += len(batch)

    await session.commit()
    return inserted


# ─── 날짜별 수집 ──────────────────────────────────────────────────

async def collect_date(
    i: int,
    total: int,
    bas_dd: str,
    market: str,
    semaphore: asyncio.Semaphore,
    http_session: aiohttp.ClientSession,
) -> tuple[bool, int]:
    """단일 날짜 수집 (세마포어 적용)."""
    async with semaphore:
        try:
            markets = []
            if market in ("kospi", "all"):
                markets.append("kospi")
            if market in ("kosdaq", "all"):
                markets.append("kosdaq")

            frames: list[pd.DataFrame] = []
            for mkt in markets:
                df = await fetch_daily_prices(http_session, mkt, bas_dd)
                if not df.empty:
                    frames.append(df)
                await asyncio.sleep(REQUEST_INTERVAL)

            if not frames:
                logger.info("[%d/%d] %s: 데이터 없음 (휴장일?)", i, total, bas_dd)
                return True, 0   # 휴장일은 성공으로 처리

            combined = pd.concat(frames, ignore_index=True)
            # 같은 날짜 내 코드 중복 제거 (KOSPI/KOSDAQ 겹침 방지)
            combined = combined.drop_duplicates(subset=["stock_code", "date"])

            async with AsyncSessionLocal() as db:
                count = await upsert_daily_prices(db, combined)

            logger.info("[%d/%d] %s: %d종목 저장", i, total, bas_dd, count)
            return True, count

        except Exception as exc:
            logger.error("[%d/%d] %s 실패: %s", i, total, bas_dd, exc)
            return False, 0


# ─── 메인 수집 로직 ───────────────────────────────────────────────

async def collect_all(
    market: str = "all",
    limit: int = 0,
    start_date: str = "20200101",
    update_mode: bool = False,
):
    """메인 수집 로직 (KRX OpenAPI 기반, 날짜별 전종목 일괄 수집)."""
    if not KRX_API_KEY:
        logger.error(
            "KRX_API_KEY가 설정되지 않았습니다.\n"
            "  1) https://openapi.krx.co.kr 에서 API 키 발급\n"
            "  2) .env 파일에 KRX_API_KEY=<발급받은키> 추가\n"
            "  3) 다시 실행하세요."
        )
        sys.exit(1)

    await create_tables()

    end_date = datetime.now().strftime("%Y%m%d")
    actual_start = start_date

    if update_mode:
        async with AsyncSessionLocal() as db:
            last_date = await get_last_trade_date(db)
        if last_date:
            next_day = last_date + timedelta(days=1)
            actual_start = next_day.strftime("%Y%m%d")
            if next_day >= date.today():
                logger.info("이미 최신 데이터입니다 (마지막: %s)", last_date)
                return
            logger.info("증분 업데이트: %s ~", actual_start)

    dates = generate_trading_dates(actual_start, end_date)
    if not dates:
        logger.warning("수집할 날짜가 없습니다.")
        return

    # --limit N → 최근 N거래일만 (테스트용)
    if limit > 0:
        dates = dates[-limit:]

    total = len(dates)
    logger.info("수집 날짜: %d일 (%s ~ %s)", total, dates[0], dates[-1])

    semaphore = asyncio.Semaphore(COLLECT_CONCURRENCY)
    conn = aiohttp.TCPConnector(limit=10)

    async with aiohttp.ClientSession(connector=conn) as http_session:
        results = await asyncio.gather(
            *[
                collect_date(i, total, d, market, semaphore, http_session)
                for i, d in enumerate(dates, 1)
            ],
            return_exceptions=True,
        )

    success  = sum(1 for r in results if isinstance(r, tuple) and r[0])
    fail     = total - success
    total_rows = sum(r[1] for r in results if isinstance(r, tuple))

    logger.info("=" * 60)
    logger.info("수집 완료!")
    logger.info("  성공: %d/%d일", success, total)
    logger.info("  실패: %d/%d일", fail, total)
    logger.info("  총 저장 행: %s", f"{total_rows:,}")
    logger.info("=" * 60)


def main():
    parser = argparse.ArgumentParser(description="KRX 히스토리컬 데이터 수집 (KRX OpenAPI)")
    parser.add_argument(
        "--market", choices=["kospi", "kosdaq", "all"], default="all",
        help="수집 시장 (기본: all)",
    )
    parser.add_argument(
        "--limit", type=int, default=0,
        help="최근 N거래일만 수집 (0=전체, 테스트용)",
    )
    parser.add_argument(
        "--start-date", default="20200101",
        help="수집 시작일 YYYYMMDD (기본: 20200101)",
    )
    parser.add_argument(
        "--update", action="store_true",
        help="증분 업데이트 모드 (마지막 저장일 이후만 수집)",
    )
    args = parser.parse_args()

    logger.info("KRX 데이터 수집 시작 (KRX OpenAPI)")
    logger.info("  시장: %s", args.market)
    logger.info("  날짜 제한: %s", "전체" if args.limit == 0 else f"최근 {args.limit}거래일")
    logger.info("  시작일: %s", args.start_date)
    logger.info("  증분 모드: %s", args.update)

    start_time = time.time()
    asyncio.run(collect_all(
        market=args.market,
        limit=args.limit,
        start_date=args.start_date,
        update_mode=args.update,
    ))
    elapsed = time.time() - start_time
    logger.info("총 소요시간: %.1f초 (%.1f분)", elapsed, elapsed / 60)


if __name__ == "__main__":
    main()