import logging
from datetime import datetime, timedelta
import pytz
import pandas as pd
from firstock import firstock

class StockDataFetcher:
    def __init__(self, client_details):
        self.user_id = client_details[0]
        self.client_details = client_details
        self.ist = pytz.timezone('Asia/Kolkata')
        self.setup_logger()

    def setup_logger(self):
        logging.basicConfig(
            level=logging.INFO,
            format='%(asctime)s - %(levelname)s - %(message)s',
            datefmt='%Y-%m-%d %H:%M:%S'
        )
        self.logger = logging.getLogger(__name__)

    def login(self):
        try:
            response = firstock.login(*self.client_details)
            if response.get("status") == "success":
                self.logger.info("Login successful")
                return True
            else:
                self.logger.error(f"Login failed: {response}")
                return False
        except Exception as e:
            self.logger.error(f"Login error: {e}")
            return False

    def fetch_time_price_series(self, exchange, trading_symbol, start_time, end_time, interval):
        """Fetch time price series data"""
        try:
            response = firstock.timePriceSeries(
                userId=self.user_id,
                exchange=exchange,
                tradingSymbol=trading_symbol,
                startTime=start_time,
                endTime=end_time,
                interval=interval
            )
            
            if response.get("status") == "success":
                data = response.get("data", [])
                if data:
                    df = pd.DataFrame(data)
                    if 'time' in df.columns:
                        df['datetime'] = pd.to_datetime(df['time'], format='%H:%M:%S %d-%m-%Y')
                    self.logger.info(f"Fetched {len(df)} records")
                    return df
            return pd.DataFrame()
            
        except Exception as e:
            self.logger.error(f"Error fetching data: {e}")
            return pd.DataFrame()

    def fetch_data(self, symbol, interval_minutes=5, count=100):
        """Fetch candle data for a symbol"""
        now_ist = datetime.now(self.ist)
        
        # Calculate time range
        end_time = now_ist
        if now_ist.hour >= 15 and now_ist.minute >= 31:
            end_time = now_ist.replace(hour=15, minute=30, second=0)
        
        start_time = end_time - timedelta(days=(count // 75) + 5)
        start_time = start_time.replace(hour=9, minute=15, second=0)
        
        # Format times
        interval_str = f"{interval_minutes}mi"
        start_str = start_time.strftime("%H:%M:%S %d-%m-%Y")
        end_str = end_time.strftime("%H:%M:%S %d-%m-%Y")
        
        # Get exchange and symbol
        translated_symbol = 'NSE:NIFTY'
        exchange, trading_symbol = translated_symbol.split(":")
        
        self.logger.info(f"Fetching {interval_minutes}min data for {symbol}")
        
        df = self.fetch_time_price_series(exchange, trading_symbol, start_str, end_str, interval_str)
        
        if not df.empty:
            df = df.sort_values(by='datetime', ascending=True)
            # Filter trading hours
            df = df[(df['datetime'].dt.hour >= 9) & (df['datetime'].dt.hour <= 15)]
            df = df[(df['datetime'].dt.hour != 9) | (df['datetime'].dt.minute >= 15)]
            df = df.tail(count)
            
        return df


if __name__ == "__main__":
    # Configuration
    client_details = ['DB1485', 'ABcd$1234', '14121985', 'DB1485_API', 'c18634e5002598698626e3590b52c520']
    
    # Initialize and login
    fetcher = StockDataFetcher(client_details)
    if not fetcher.login():
        exit(1)
    
    # Fetch data
    symbol = 'nifty'
    df = fetcher.fetch_data(symbol, interval_minutes=5, count=100)
    
    if not df.empty:
        print(f"\nFetched {len(df)} candles")
        print(df[['datetime', 'open', 'high', 'low', 'close']].head())
        print(f"\nDate range: {df['datetime'].min()} to {df['datetime'].max()}")
    else:
        print("No data fetched")