'''
📁 data_collection_engine/collect_all_cds.py

S&P500 등 여러 종목에 대해 CDS 데이터를 일괄 수집하고,
XGBoost 학습을 위한 통합 데이터셋도 자동 생성하는 스크립트
'''

import os
import time
import pandas as pd
from data_collection_engine.engine import DataCollectionEngine
from data_analysis_engine.dataset_builder import build_dataset

def collect_all_cds(ticker_list_path: str, start: str, end: str, save_dir: str = "data", delay_sec: int = 13):
    """
    여러 종목의 CDS 데이터를 일괄 수집하여 저장하고, 학습용 통합 데이터셋도 생성

    Parameters:
        ticker_list_path (str): 수집할 종목 코드가 들어 있는 CSV 파일 경로
        start (str): 시작일 (YYYY-MM-DD)
        end (str): 종료일 (YYYY-MM-DD)
        save_dir (str): 저장 디렉토리
        delay_sec (int): API 호출 간격 (기본 13초)
    """
    df = pd.read_csv(ticker_list_path)

    # 사용자에게 타입 선택 옵션 제공
    print("\n🔍 선택 가능한 종목 유형:")
    type_mapping = {
        "CS": "Common Stock - 일반 주식 (가장 보편적인 기업 지분)",
        "ETF": "Exchange Traded Fund - 지수 추종형 상장지수펀드",
        "ETN": "Exchange Traded Note - 채권 기반의 상장지수증권",
        "ADR": "American Depository Receipt - 해외 주식의 미국 상장 버전",
        "PREF": "Preferred Stock - 배당 우선주",
        "UNIT": "Unit - ETF 구성 단위 또는 묶음",
        "RIGHT": "Rights - 신주인수권, 일정 기간 내 주식 매입 권리",
        "FUND": "Mutual Fund - 공모펀드 또는 투자신탁",
        "SP": "Structured Product - 파생결합증권",
        "WARRANT": "Warrant - 일정 가격에 주식 구매 권한을 부여하는 증서",
        "INDEX": "Market Index - 종합 주가지수 등",
        "BOND": "Bond - 회사채 또는 정부채 등 고정수익 상품"
    }
    for k, v in type_mapping.items():
        print(f"{k:<7}: {v}")
    type_input = input("\n✅ 수집할 종목 유형을 입력하세요 (예: CS): ").strip().upper()

    if "type" not in df.columns:
        print("❗ 'type' 컬럼이 존재하지 않습니다. 전체 데이터 사용.")
        filtered_df = df.copy()
    else:
        filtered_df = df[df["type"] == type_input].copy()

    total_available = len(filtered_df)
    print(f"\n총 {total_available}개의 종목이 선택된 유형({type_input})에 해당합니다.")
    estimated_time = total_available * delay_sec / 60
    print(f"⏳ 예상 소요 시간: 약 {estimated_time:.1f}분 (요금제 기준 {delay_sec}초 간격)")
    count_limit = input(f"몇 개를 수집하시겠습니까? [기본: {total_available}]: ").strip()
    count = int(count_limit) if count_limit else total_available

    tickers = filtered_df.head(count)
    symbol_col = "symbol" if "symbol" in tickers.columns else "ticker"
    ticker_list = tickers[symbol_col].dropna().unique()

    # 종료 날짜 기준으로 저장 디렉토리 세분화 (예: data/CS/2024-12-31)
    dated_dir = os.path.join(save_dir, type_input, end)
    engine = DataCollectionEngine(data_dir=dated_dir)
    os.makedirs(dated_dir, exist_ok=True)

    failed = []
    X_all, y_all = [], []
    error_429_count = 0

    for i, symbol in enumerate(ticker_list):
        try:
            print(f"[{i+1}/{len(ticker_list)}] {symbol} 수집 중...")
            df = engine.collect(symbol.strip().upper(), start, end)
            X, y = build_dataset(df)
            X_all.append(X)
            y_all.append(y)
            error_429_count = 0
            time.sleep(delay_sec)
        except Exception as e:
            print(f"❌ {symbol} 수집 실패: {e}")
            failed.append(symbol)
            if "429" in str(e):
                error_429_count += 1
                if error_429_count >= 2:
                    print("⏸️ 연속된 429 오류 감지 → 60초 대기 중...")
                    time.sleep(60)

    if failed:
        print("\n⚠️ 일부 종목 수집 실패:", failed)
        failed_df = pd.DataFrame(failed, columns=["ticker"])
        failed_df.to_csv(os.path.join(dated_dir, "failed_tickers.csv"), index=False)
        print("📄 실패 종목 목록 저장 완료: failed_tickers.csv")
    else:
        print("\n✅ 전체 종목 수집 완료!")

    if X_all:
        X_total = pd.concat(X_all, ignore_index=True)
        y_total = pd.concat(y_all, ignore_index=True)
        X_total.to_csv(os.path.join(dated_dir, "X_total.csv"), index=False)
        y_total.to_csv(os.path.join(dated_dir, "y_total.csv"), index=False)
        print("✅ 통합 학습용 데이터 저장 완료: X_total.csv, y_total.csv")
    else:
        print("❗ 유효한 CDS 데이터가 없어 학습용 데이터셋을 생성하지 못했습니다.")

if __name__ == "__main__":
    print("CDS 일괄 수집을 시작합니다. 아래 정보를 입력해 주세요:")
    ticker_list_path = input("티커 리스트 CSV 경로 [기본: sp500_tickers.csv]: ").strip() or "sp500_tickers.csv"
    start = input("시작일 (YYYY-MM-DD) [기본: 2023-01-01]: ").strip() or "2023-01-01"
    end = input("종료일 (YYYY-MM-DD) [기본: 2024-12-31]: ").strip() or "2024-12-31"
    save_dir = input("저장 디렉토리 이름 [기본: data]: ").strip() or "data"
    delay = input("API 호출 간격 (초) [기본: 13]: ").strip()
    delay_sec = int(delay) if delay else 13

    collect_all_cds(
        ticker_list_path=ticker_list_path,
        start=start,
        end=end,
        save_dir=save_dir,
        delay_sec=delay_sec
    )