import pandas as pd

from datetime import datetime
from exchange import SMART
from ib_insync import IB, Index, Option
from ib_insync.util import isNan
from market_data_type import LIVE
from typing import Callable

class Client:

    def __init__(self, host: str, port: int, client_id = 1) -> None:
        self.ib = IB()
        self.ib.connect(host, port, clientId = client_id)
        self.ib.reqMarketDataType(LIVE)
    
    def get_ticker(self, symbol: str, exchange: str):
        underlying = Index(symbol, exchange)
        self.ib.qualifyContracts(underlying)
        return self.ib.reqTickers(underlying)[0]
    
    def get_option_chain(self, symbol: str, expiration: datetime, sub_symbol: str = None, 
            contract_filter: Callable = None) -> pd.DataFrame:
        expiration_date = expiration.strftime('%Y%m%d')
        contract_details = self.ib.reqContractDetails(Option(symbol, expiration_date, exchange = SMART))
        contracts = [_.contract for _ in contract_details if not sub_symbol or _.contract.tradingClass == sub_symbol]
        if contract_filter:
            contracts = list(filter(contract_filter, contracts))
        
        option_data = []
        for contract in contracts:
            option = Option(symbol, expiration_date, contract.strike, contract.right, exchange = SMART, currency = 'USD')
            if sub_symbol:
                option.tradingClass = sub_symbol
            option_data_snapshot = self.ib.reqMktData(option, '', True, False)
            option_data.append(option_data_snapshot)
        
        option_chain = []
        for option in option_data:
            print('Processing Option: ', option)
            while isNan(option.bid) or isNan(option.ask) or option.modelGreeks is None:
                # TODO: Add a timeout?
                self.ib.sleep()
            option_chain.append({
                'Strike': option.contract.strike,
                'Type': option.contract.right, # 'C' for Call, 'P' for Put.
                'Bid': option.bid,
                'Ask': option.ask,
                'Delta': option.modelGreeks.delta,
            })
        
        return pd.DataFrame(option_chain)