import sys
import json
import pytz
import pandas as pd
from datetime import datetime, timedelta
from firstock import firstock
import pprint

class OrderBlockFetcher:
    def __init__(self, client_details):
        self.client_details = client_details
        self.user_id = client_details[0]
        self.ist = pytz.timezone('Asia/Kolkata')
        self.output_data = {
            'timestamp': None,
            'symbol': 'NIFTY',
            'timeframes': {
                '5min': {
                    'historical_data': [],
                    'order_blocks': []
                }
            }
        }

    def login(self):
        response = firstock.login(*self.client_details)
        if response.get("status") != "success":
            print("[ERROR] Login failed:", response)
            sys.exit()
        print("[INFO] Login successful")

    def fetch_time_price_series(self, exchange, trading_symbol, start_time, end_time, interval):
        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_temp = pd.DataFrame(data)
                if 'time' in df_temp.columns:
                    df_temp['datetime'] = pd.to_datetime(df_temp['time'], format='%H:%M:%S %d-%m-%Y')
                return df_temp
        return pd.DataFrame()

    def fetch_candles_last_2days(self, symbol="NIFTY", interval_minutes=5):
        now_ist = datetime.now(self.ist)
        end_time = now_ist.replace(hour=15, minute=30, second=0) if now_ist.hour >= 15 else now_ist
        start_time = end_time - timedelta(days=20)
        start_time = start_time.replace(hour=9, minute=15, second=0)

        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")

        translated_symbol = 'NSE:NIFTY'
        exchange, trading_symbol = translated_symbol.split(":")

        print(f"[INFO] Fetching {interval_minutes}min candles for last 2 days")
        df = self.fetch_time_price_series(exchange, trading_symbol, start_str, end_str, interval_str)

        if df.empty:
            print("[ERROR] No candle data fetched")
            return df

        df = df.sort_values(by='datetime', ascending=True)
        df = df[(df['datetime'].dt.hour >= 9) & (df['datetime'].dt.hour <= 15)]
        df = df[(df['datetime'].dt.hour != 9) | (df['datetime'].dt.minute >= 15)]
        return df

    def detect_order_blocks(self, candles):
        df = pd.DataFrame(candles)
        df['datetime'] = pd.to_datetime(df['epochTime'], unit='s').dt.tz_localize('UTC').dt.tz_convert(self.ist)

        order_blocks = []
        lookback = 2  # past 10 candles

        for i in range(2, len(df)):
            c1, c2, c3 = df.iloc[i-2], df.iloc[i-1], df.iloc[i]

            c1_color = "green" if c1['close'] > c1['open'] else "red"
            c2_color = "green" if c2['close'] > c2['open'] else "red"
            c3_color = "green" if c3['close'] > c3['open'] else "red"

            # Get last 10 candles before c2
            if i-1 >= lookback:
                prev10 = df.iloc[i-lookback-1:i-1]
            else:
                prev10 = df.iloc[:i-1]

            # Compute extended end index (at least 10 candles ahead)
            end_idx = min(i+10, len(df)-1)
            end_candle = df.iloc[end_idx]

            # Bullish OB: red, red, green
            if c1_color == "red" and c2_color == "red" and c3_color == "green":
                if (c2['high'] < c1['high'] and c2['low'] < c1['low'] and
                    c2['low'] < prev10['low'].min() and c2['high'] < prev10['high'].min() and
                    c3['close'] > c2['close'] and c3['high'] > c2['high']):
                    order_blocks.append({
                        "type": "Bullish",
                        "startEpoch": int(c1['epochTime']),
                        "endEpoch": int(end_candle['epochTime']),  # extend 10 candles forward
                        "high": c2['high'],
                        "low": c2['low'],
                        "datetime": str(c3['datetime'])
                    })

            # Bearish OB: green, green, red
            if c1_color == "green" and c2_color == "green" and c3_color == "red":
                if (c2['high'] > c1['high'] and c2['low'] > c1['low'] and
                    c2['high'] > prev10['high'].max() and c2['low'] > prev10['low'].max() and
                    c3['close'] < c2['close'] and c3['low'] < c2['low']):
                    order_blocks.append({
                        "type": "Bearish",
                        "startEpoch": int(c1['epochTime']),
                        "endEpoch": int(end_candle['epochTime']),  # extend 10 candles forward
                        "high": c2['high'],
                        "low": c2['low'],
                        "datetime": str(c3['datetime'])
                    })

        return order_blocks

    def fetch_all_data(self, symbol="NIFTY"):
        df = self.fetch_candles_last_2days(symbol, 5)
        if df.empty:
            print("[ERROR] Failed to fetch candles")
            return

        ohlc_records = df[['epochTime','open','high','low','close']].to_dict(orient='records')
        blocks = self.detect_order_blocks(ohlc_records)

        self.output_data['symbol'] = symbol.upper()
        self.output_data['timestamp'] = int(datetime.now(self.ist).timestamp())
        self.output_data['timeframes']['5min']['historical_data'] = ohlc_records
        self.output_data['timeframes']['5min']['order_blocks'] = blocks

        json_filename = f"{symbol.upper()}_orderblocks.json"
        with open(json_filename, "w") as f:
            json.dump(self.output_data, f, indent=4)

        print(f"✅ Data saved to {json_filename}")


if __name__ == "__main__":
    print("[INFO] Script started.")
    client_details = ['DB1485', 'ABcd#1234', '14121985', 'DB1485_API', 'd9ce8d1261a834458847929c11ea1047']
    symbol = sys.argv[1] if len(sys.argv) > 1 else 'nifty'

    fetcher = OrderBlockFetcher(client_details)
    fetcher.login()
    fetcher.fetch_all_data(symbol)
    print("[INFO] Execution completed.")