import pandas as pd
from datetime import datetime, timedelta
from typing import Dict, List
from .client import ZerodhaClient

class DataFetcher:
    TIMEFRAME_MAP = {
        '5min': '5minute',
        '15min': '15minute',
        '30min': '30minute',
        '1h': '60minute',
        'day': 'day',
        'week': 'week',
        'month': 'month'
    }
    
    def __init__(self, zerodha_client: ZerodhaClient):
        self.zerodha = zerodha_client
        self.instrument_cache = {}
    
    def resolve_symbol(self, symbol: str, exchange: str = "NSE") -> int:
        """Get instrument token for a symbol"""
        cache_key = f"{exchange}:{symbol}"
        if cache_key in self.instrument_cache:
            return self.instrument_cache[cache_key]
        
        instruments = self.zerodha.kite.instruments(exchange)
        for inst in instruments:
            if inst['tradingsymbol'] == symbol:
                self.instrument_cache[cache_key] = inst['instrument_token']
                return inst['instrument_token']
        raise ValueError(f"Instrument not found: {symbol} on {exchange}")
    
    def fetch_ohlc(
        self,
        symbol: str,
        interval: str,
        days_back: int = 60,
        end_date: datetime = None
    ) -> pd.DataFrame:
        """Fetch OHLC data for a symbol"""
        if not end_date:
            end_date = datetime.now()
        
        start_date = end_date - timedelta(days=days_back)
        instrument_token = self.resolve_symbol(symbol)
        
        # Convert to Kite interval format
        kite_interval = self.TIMEFRAME_MAP.get(interval)
        if not kite_interval:
            raise ValueError(f"Unsupported interval: {interval}")
        
        return self.zerodha.get_historical_data(
            instrument_token=instrument_token,
            interval=kite_interval,
            from_date=start_date.strftime("%Y-%m-%d"),
            to_date=end_date.strftime("%Y-%m-%d")
        )
    
    def fetch_multi_timeframe(
        self,
        symbol: str,
        timeframes: List[str],
        candles_count: int = 100
    ) -> Dict[str, pd.DataFrame]:
        """Fetch data for multiple timeframes"""
        data = {}
        for tf in timeframes:
            df = self.fetch_ohlc(symbol, tf)
            # Ensure we have enough candles
            if len(df) > candles_count:
                df = df.iloc[-candles_count:]
            data[tf] = df
        return data