"""
etf_backtest.py — ETF 액티브 매매 백테스트 (RSI 기반 분할 매수 & 슈팅 익절)
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
역할: 테마성 ETF (원자력, 전력망 등) 의 눌림목 분할 매수와 슈팅 익절 전략 검증
      일봉 데이터로 백테스트 수행, 거래 내역은 즉시 저장 (Atomic Save)

전략 개요:
  - 매수: RSI 35/30/25 이하에서 3 분할 매수 (30%/30%/40%)
  - 매도: 수익률 +4% 이상 또는 RSI 70 이상에서 전량 익절
  - 손절: 평단가 대비 -10% 에서 전량 손절 (리스크 관리)
  - 유니버스: 사용자가 지정한 ETF 종목 (예: URA, ICLN, QCLN 등)

KIS API 준수:
  - SafeRequest: API 429 에러 대비 재시도 로직
  - 메신저 알림: 텔레그램 (HTML), 매터모스트 (Markdown) 분리
  - Atomic Save: 매매 발생 시 즉시 JSON 저장 (재시작 안전성)
"""

import yfinance as yf
import pandas as pd
import numpy as np
import json
import os
import time
import random
import requests
from datetime import datetime
from typing import Dict, List, Optional

# 로깅 설정
import logging
logging.basicConfig(
    format='[%(asctime)s] %(message)s',
    datefmt='%H:%M:%S',
    level=logging.INFO,
)
logger = logging.getLogger("ETFBacktest")

# ==============================================================================
# [메신저 알림 기능] 
# 텔레그램 (HTML) 및 매터모스트 (Markdown) 분리 구현
# ==============================================================================
def msg_tg(token: str, chat_id: str, trade_info: Dict):
    """
    텔레그램 메시지 전송 (HTML 방식)
    이미지 태그 (URL 링크) 방식을 우선하여 속도 저하를 방지합니다.
    """
    url = f"https://api.telegram.org/bot{token}/sendMessage"
    
    # 텔레그램 메시지 상세 구성
    text = (
        f"<b>🔔 [ETF 액티브 봇] {trade_info['type']}</b>\n\n"
        f"▪️ <b>종목명:</b> {trade_info['ticker']}\n"
        f"▪️ <b>체결일자:</b> {trade_info['date']}\n"
        f"▪️ <b>체결단가:</b> {trade_info['price']:,.2f}원\n"
        f"▪️ <b>체결수량:</b> {trade_info['qty']:,}주\n"
        f"▪️ <b>현재 RSI:</b> {trade_info['rsi']:.1f}\n"
    )
    
    if "profit_loss" in trade_info:
        text += f"▪️ <b>실현손익:</b> {trade_info['profit_loss']:,.0f}원\n"
        
    text += f"▪️ <b>현재자산 (현금):</b> {trade_info['capital']:,.0f}원\n"
    text += f"\n📊 <a href='https://finance.yahoo.com/quote/{trade_info['ticker']}/chart'>[야후 파이낸스 차트 확인]</a>"

    payload = {
        "chat_id": chat_id,
        "text": text,
        "parse_mode": "HTML",
        "disable_web_page_preview": False
    }
    try:
        requests.post(url, data=payload, timeout=5)
    except Exception as e:
        logger.debug(f"[텔레그램 전송 에러] {e}")


def msg_mm(webhook_url: str, trade_info: Dict):
    """
    매터모스트 메시지 전송 (Markdown 방식)
    """
    text = (
        f"### 🔔 [ETF 액티브 봇] {trade_info['type']}\n"
        f"- **종목명:** {trade_info['ticker']}\n"
        f"- **체결일자:** {trade_info['date']}\n"
        f"- **체결단가:** {trade_info['price']:,.2f}원\n"
        f"- **체결수량:** {trade_info['qty']:,}주\n"
        f"- **현재 RSI:** {trade_info['rsi']:.1f}\n"
    )
    
    if "profit_loss" in trade_info:
        text += f"- **실현손익:** {trade_info['profit_loss']:,.0f}원\n"
        
    text += f"- **현재자산 (현금):** {trade_info['capital']:,.0f}원\n"
    text += f"\n📊 [야후 파이낸스 차트 확인](https://finance.yahoo.com/quote/{trade_info['ticker']}/chart)"

    payload = {
        "text": text
    }
    try:
        requests.post(webhook_url, json=payload, timeout=5)
    except Exception as e:
        logger.debug(f"[매터모스트 전송 에러] {e}")


# ==============================================================================
# [네트워크/API 안전성 관리] 
# 모든 API 요청은 SafeRequest 클래스를 거쳐 429 에러 발생 시 재시도합니다.
# ==============================================================================
class SafeRequest:
    def __init__(self, max_retries: int = 3, delay: float = 2.0):
        self.max_retries = max_retries
        self.delay = delay

    def fetch_data(self, ticker_symbol: str, start_date: str, end_date: str) -> Optional[pd.DataFrame]:
        """yfinance 데이터 다운로드 (재시도 로직 포함)"""
        for attempt in range(self.max_retries):
            try:
                logger.info(f"[{datetime.now().strftime('%H:%M:%S')}] {ticker_symbol} 데이터 요청 중... (시도: {attempt + 1}/{self.max_retries})")
                data = yf.download(ticker_symbol, start=start_date, end=end_date, progress=False)
                if not data.empty:
                    return data
            except Exception as e:
                logger.warning(f"[에러 발생] 데이터 요청 실패: {e}")
                time.sleep(self.delay)
        raise Exception("API 요청 한도 초과 또는 네트워크 오류가 지속됩니다.")


# ==============================================================================
# [핵심 매매 로직] RSI 기반 3 분할 매수 및 슈팅 익절 알고리즘
# ==============================================================================
class ETFActiveBacktester:
    """
    ETF 액티브 매매 백테스터
    
    사용법:
        tester = ETFActiveBacktester(
            ticker="URA",
            start_date="2023-01-01",
            end_date="2024-01-01",
            initial_capital=10000000
        )
        tester.run()
    """
    
    def __init__(self, ticker: str, start_date: str, end_date: str, initial_capital: float = 10000000):
        self.ticker = ticker
        self.start_date = start_date
        self.end_date = end_date
        
        # 기본 자본금 및 수수료/슬리피지 설정
        self.capital = initial_capital
        self.initial_capital = initial_capital
        self.fee_rate = 0.00015       # 수수료 0.015%
        self.slippage_rate = 0.0005   # 슬리피지 0.05%
        
        # 포트폴리오 상태
        self.position = 0          # 보유 수량
        self.total_invested = 0.0  # 총 투자 금액 (평단가 계산용)
        self.buy_step = 0          # 현재 몇 차 매수까지 진행되었는지 (0~3)
        
        # 로깅 설정 (Atomic Save)
        self.trade_log_file = f"{self.ticker}_etf_trade_history.json"
        if os.path.exists(self.trade_log_file):
            with open(self.trade_log_file, "r", encoding="utf-8") as f:
                try:
                    self.trade_history = json.load(f)
                except json.JSONDecodeError:
                    self.trade_history = []
        else:
            self.trade_history = []

    def save_trade_immediately(self, trade_record: Dict):
        """
        [Atomic Save] 데이터 파일은 프로그램 종료 시점이 아니라 이벤트 즉시 저장합니다.
        기존 데이터를 보존하며 추가합니다. (재시작 안전성)
        """
        self.trade_history.append(trade_record)
        with open(self.trade_log_file, "w", encoding="utf-8") as f:
            json.dump(self.trade_history, f, indent=4, ensure_ascii=False)
        logger.info(f"✅ 거래 로그 저장: {trade_record['date']} | {trade_record['type']}")

    def calculate_rsi(self, df: pd.DataFrame, period: int = 14) -> pd.DataFrame:
        """
        RSI(상대강도지수) 계산 함수. 초보자도 이해하기 쉬운 단순 이동평균 방식 사용.
        
        Args:
            df: 일봉 데이터 (Open, High, Low, Close 컬럼 필요)
            period: RSI 계산 기간 (기본 14 일)
        
        Returns:
            RSI 컬럼이 추가된 DataFrame
        """
        delta = df['Close'].diff()
        gain = (delta.where(delta > 0, 0)).ewm(alpha=1/period, adjust=False).mean()
        loss = (-delta.where(delta < 0, 0)).ewm(alpha=1/period, adjust=False).mean()
        rs = gain / loss
        df['RSI'] = 100 - (100 / (1 + rs))
        return df

    def run(self, enable_notifications: bool = False, tg_token: str = "", tg_chat_id: str = "", mm_webhook: str = ""):
        """
        백테스트 실행
        
        Args:
            enable_notifications: 메신저 알림 활성화 여부
            tg_token: 텔레그램 봇 토큰
            tg_chat_id: 텔레그램 채팅 ID
            mm_webhook: 매터모스트 웹훅 URL
        """
        api = SafeRequest()
        df = api.fetch_data(self.ticker, self.start_date, self.end_date)
        
        if df.empty:
            logger.error(f"❌ {self.ticker} 데이터를 가져올 수 없습니다.")
            return
        
        # DataFrame MultiIndex 구조 평탄화 (yfinance 최신버전 호환)
        if isinstance(df.columns, pd.MultiIndex):
            df.columns = df.columns.get_level_values(0)
            
        df = self.calculate_rsi(df)
        
        logger.info(f"\n=== [{self.ticker}] ETF 액티브 백테스트 시작 ===")
        logger.info(f"초기 자본금: {self.capital:,.0f}원")
        logger.info(f"테스트 기간: {self.start_date} ~ {self.end_date}\n")
        
        for index, row in df.iterrows():
            # 서버 부하 방지를 위해 일반 루프에 1~3 초 랜덤 딜레이 적용
            time.sleep(random.uniform(0.5, 1.5))
            
            if pd.isna(row.get('RSI')):
                continue  # RSI 계산을 위한 초기 기간 (14 일) 패스
                
            current_date = index.strftime('%Y-%m-%d')
            close_price = float(row['Close'])
            rsi = float(row['RSI'])
            
            # 평균 단가 계산
            avg_price = (self.total_invested / self.position) if self.position > 0 else 0.0
            
            # [1. 매도 로직] : 보유 수량이 있을 때 (익절 또는 손절)
            if self.position > 0:
                profit_rate = (close_price - avg_price) / avg_price if avg_price > 0 else 0
                
                # 익절 조건: 4% 이상 수익이 났거나, RSI 가 70(과열) 을 넘을 때 전량 매도
                if profit_rate >= 0.04 or rsi >= 70:
                    sell_price = close_price * (1 - self.fee_rate - self.slippage_rate)
                    revenue = self.position * sell_price
                    profit_loss = revenue - self.total_invested
                    self.capital += revenue
                    
                    record = {
                        "ticker": self.ticker,
                        "date": current_date,
                        "type": "SELL (슈팅 익절)",
                        "price": round(sell_price, 2),
                        "qty": self.position,
                        "rsi": round(rsi, 2),
                        "capital": round(self.capital, 2),
                        "profit_loss": round(profit_loss, 2)
                    }
                    self.save_trade_immediately(record)
                    
                    # 메신저 알림
                    if enable_notifications:
                        msg_tg(tg_token, tg_chat_id, record)
                        msg_mm(mm_webhook, record)
                    
                    # 상태 초기화
                    self.position = 0
                    self.total_invested = 0.0
                    self.buy_step = 0
                    continue

                # 손절 조건: 평단가 대비 -10% 이하로 떨어지면 리스크 관리 차원에서 손절
                elif profit_rate <= -0.10:
                    sell_price = close_price * (1 - self.fee_rate - self.slippage_rate)
                    revenue = self.position * sell_price
                    profit_loss = revenue - self.total_invested
                    self.capital += revenue
                    
                    record = {
                        "ticker": self.ticker,
                        "date": current_date,
                        "type": "SELL (안전 손절)",
                        "price": round(sell_price, 2),
                        "qty": self.position,
                        "rsi": round(rsi, 2),
                        "capital": round(self.capital, 2),
                        "profit_loss": round(profit_loss, 2)
                    }
                    self.save_trade_immediately(record)
                    
                    # 메신저 알림
                    if enable_notifications:
                        msg_tg(tg_token, tg_chat_id, record)
                        msg_mm(mm_webhook, record)
                    
                    # 상태 초기화
                    self.position = 0
                    self.total_invested = 0.0
                    self.buy_step = 0
                    continue

            # [2. 매수 로직] : RSI 에 따른 3 분할 매수
            # 자금 배분: 1 차 30%, 2 차 30%, 3 차 40%
            actual_buy_price = close_price * (1 + self.fee_rate + self.slippage_rate)
            buy_amount = 0
            buy_type = ""
            
            if self.buy_step == 0 and rsi < 35:
                # 1 차 매수 (RSI 35 미만)
                buy_amount = self.initial_capital * 0.3
                buy_type = "BUY (1 차 진입)"
                self.buy_step = 1
                
            elif self.buy_step == 1 and rsi < 30:
                # 2 차 매수 (RSI 30 미만)
                buy_amount = self.initial_capital * 0.3
                buy_type = "BUY (2 차 물타기)"
                self.buy_step = 2
                
            elif self.buy_step == 2 and rsi < 25:
                # 3 차 매수 (RSI 25 미만)
                buy_amount = self.initial_capital * 0.4  # 남은 현금 전부
                buy_type = "BUY (3 차 풀매수)"
                self.buy_step = 3

            # 매수 실행
            if buy_amount > 0 and self.capital >= buy_amount:
                quantity = int(buy_amount // actual_buy_price)
                if quantity > 0:
                    invest_cost = quantity * actual_buy_price
                    self.position += quantity
                    self.total_invested += invest_cost
                    self.capital -= invest_cost
                    
                    record = {
                        "ticker": self.ticker,
                        "date": current_date,
                        "type": buy_type,
                        "price": round(actual_buy_price, 2),
                        "qty": quantity,
                        "rsi": round(rsi, 2),
                        "capital": round(self.capital, 2)
                    }
                    self.save_trade_immediately(record)
                    
                    # 메신저 알림
                    if enable_notifications:
                        msg_tg(tg_token, tg_chat_id, record)
                        msg_mm(mm_webhook, record)

        # 백테스트 종료
        logger.info(f"\n=== 백테스트 종료 ===")
        avg_price = (self.total_invested / self.position) if self.position > 0 else 0.0
        final_assets = self.capital + (self.position * avg_price)
        
        # 성과 요약
        total_return = ((final_assets - self.initial_capital) / self.initial_capital) * 100
        
        logger.info(f"📊 [{self.ticker}] 백테스트 결과")
        logger.info(f"  초기 자본금: {self.initial_capital:,.0f}원")
        logger.info(f"  최종 자산: {final_assets:,.0f}원")
        logger.info(f"  총 수익률: {total_return:.2f}%")
        logger.info(f"  총 매매 횟수: {len([t for t in self.trade_history if 'SELL' in t['type']])}회")
        
        # 메신저 알림 (최종 결과)
        if enable_notifications:
            final_record = {
                "ticker": self.ticker,
                "date": datetime.now().strftime('%Y-%m-%d'),
                "type": "백테스트 완료",
                "price": 0,
                "qty": 0,
                "rsi": 0,
                "capital": round(final_assets, 2),
                "profit_loss": round(final_assets - self.initial_capital, 2)
            }
            msg_tg(tg_token, tg_chat_id, final_record)
            msg_mm(mm_webhook, final_record)


# ==============================================================================
# [실행부]
# ==============================================================================
if __name__ == "__main__":
    # 예시: URA (글로벌 X 우라늄 ETF) 로 테스트
    # 한국 ETF 는 '069500.KS' (KODEX 200) 형태
    tester = ETFActiveBacktester(
        ticker="URA",
        start_date="2023-01-01",
        end_date="2024-12-31",
        initial_capital=10000000  # 1 천만원
    )
    
    # 백테스트 실행 (메신저 알림은 주석 처리)
    tester.run(
        enable_notifications=False,  # True 로 설정 후 토큰 입력하면 알림 발송
        # tg_token="YOUR_BOT_TOKEN",
        # tg_chat_id="YOUR_CHAT_ID",
        # mm_webhook="YOUR_WEBHOOK_URL"
    )