from datetime import datetime, timedelta
from ib_insync import Contract
from logger import LOGGER as log

def get_next_contract_month(ticker: str, offset=0, check_rollover=False) -> str:
    """Returns the next contract month in 'YYYYMM' format based on the ticker and specified conditions."""
    now = datetime.now()

    def next_quarter_month(date, offset=0):
        """Get the next quarter month (March, June, September, December) with an optional offset."""
        quarter_months = [3, 6, 9, 12]
        month = date.month
        year = date.year

        # Find the next quarter month
        next_month = min((m for m in quarter_months if m > month), default=quarter_months[0])
        if next_month <= month:
            year += 1
        next_month_index = (quarter_months.index(next_month) + offset) % len(quarter_months)
        next_month = quarter_months[next_month_index]
        return year, next_month

    def get_monday_before_second_friday(year, month):
        """Get the Monday before the second Friday of the given month and year."""
        first_day = datetime(year, month, 1)
        first_friday = first_day + timedelta(days=(4 - first_day.weekday() + 7) % 7)  # First Friday
        second_friday = first_friday + timedelta(days=7)
        return second_friday - timedelta(days=second_friday.weekday() + 7)

    def get_monday_before_third_friday(year, month):
        """Get the Monday before the third Friday of the given month and year."""
        first_day = datetime(year, month, 1)
        first_friday = first_day + timedelta(days=(4 - first_day.weekday() + 7) % 7)  # First Friday
        third_friday = first_friday + timedelta(days=14)
        return third_friday - timedelta(days=third_friday.weekday() + 7)

    def get_last_monday_before_second_last_trading_day(year, month):
        """Get the last Monday before the second last trading day of the given month and year."""
        last_day = datetime(year, month + 1, 1) - timedelta(days=1)
        second_last_trading_day = last_day
        while second_last_trading_day.weekday() in (5, 6):  # Skip weekends
            second_last_trading_day -= timedelta(days=1)
        second_last_trading_day -= timedelta(days=1)
        while second_last_trading_day.weekday() in (5, 6):  # Skip weekends
            second_last_trading_day -= timedelta(days=1)
        return second_last_trading_day - timedelta(days=second_last_trading_day.weekday() + 1)

    if ticker in ["ES", "MES", "NQ", "MNQ"]:
        year, month = next_quarter_month(now)
        monday_before_third_friday = get_monday_before_third_friday(year, month)
        if now >= monday_before_third_friday:
            year, month = next_quarter_month(datetime(year, month, 1), 1)
            if check_rollover:
                return True  # Indicate rollover is needed
        return f"{year}{month:02d}"

    elif ticker in ["TOPX", "MNTPX"]:
        year, month = next_quarter_month(now)
        monday_before_second_friday = get_monday_before_second_friday(year, month)
        if now >= monday_before_second_friday:
            year, month = next_quarter_month(datetime(year, month, 1), 1)
            if check_rollover:
                return True  # Indicate rollover is needed
        return f"{year}{month:02d}"

    elif ticker in ["HSI", "MHI"]:
        year, month = next_quarter_month(now)
        last_monday = get_last_monday_before_second_last_trading_day(year, month)
        if now >= last_monday:
            year, month = next_quarter_month(datetime(year, month, 1), 1)
            if check_rollover:
                return True  # Indicate rollover is needed
        return f"{year}{month:02d}"

    else:
        log.error(f"Invalid ticker: {ticker}. Please check the alert message.")
        return None

def contract_type_check(ticker: str, contract_month=None) -> Contract:
    contract = Contract()
    contract.symbol = ticker
    if not contract_month:
        contract_month = get_next_contract_month(ticker)
    if ticker == "SPY":
        contract.secType = "STK"
        contract.currency = "USD"
        contract.exchange = "ARCA"
    elif ticker.startswith("ETH"):
        contract.secType = "CRYPTO"
        contract.currency = "USD"
        contract.exchange = "PAXOS"
    elif ticker.startswith("EUR"):
        contract.secType = "CASH"
        contract.currency = "USD"
        contract.exchange = "IDEALPRO"
    elif ticker in ["ES", "MES", "NQ", "MNQ", "HSI", "MHI", "TOPX", "MNTPX"]:
        contract.secType = "FUT"
        contract.currency = "USD" if ticker in ["ES", "MES", "NQ", "MNQ"] else "HKD" if ticker in ["HSI", "MHI"] else "JPY"
        contract.exchange = "CME" if ticker in ["ES", "MES", "NQ", "MNQ"] else "HKFE" if ticker in ["HSI", "MHI"] else "OSE"
        contract.lastTradeDateOrContractMonth = contract_month
    else:
        log.error(f"Invalid ticker: {ticker}. Please check the alert message.")
        return None

    return contract