import os
import time
import requests
import pandas as pd
from datetime import datetime, date, timedelta
from zoneinfo import ZoneInfo
import pandas_market_calendars as mcal

IST = ZoneInfo("Asia/Kolkata")
DATA_FILE = os.path.join(os.path.dirname(__file__), "data", "nifty50_daily.parquet")

TICKERS = [
    'ADANIENT', 'ADANIPORTS', 'APOLLOHOSP', 'ASIANPAINT', 'AXISBANK', 'BAJAJ-AUTO', 'BAJAJFINSV', 'BAJFINANCE', 
    'BHARTIARTL', 'BPCL', 'BRITANNIA', 'CIPLA', 'COALINDIA', 'DIVISLAB', 'DRREDDY', 'EICHERMOT', 'GRASIM', 
    'HCLTECH', 'HDFCBANK', 'HDFCLIFE', 'HEROMOTOCO', 'HINDALCO', 'HINDUNILVR', 'ICICIBANK', 'INDUSINDBK', 
    'INFY', 'ITC', 'JSWSTEEL', 'KOTAKBANK', 'LT', 'M&M', 'MARUTI', 'NESTLEIND', 'NTPC', 'ONGC', 'POWERGRID', 
    'RELIANCE', 'SBILIFE', 'SBIN', 'SUNPHARMA', 'TATACONSUM', 'TATAMOTORS', 'TATASTEEL', 'TCS', 'TECHM', 
    'TITAN', 'ULTRACEMCO', 'UPL', 'WIPRO'
]

def is_trading_day(target_date: date) -> bool:
    try:
        nse = mcal.get_calendar('NSE')
        schedule = nse.schedule(start_date=target_date, end_date=target_date)
        return not schedule.empty
    except Exception as e:
        print(f"Calendar check failed: {e}")
        # Fallback: assume Monday-Friday is trading day
        return target_date.weekday() < 5

def fetch_groww_data(ticker: str, start_ts: int, end_ts: int):
    url = f"https://groww.in/v1/api/charting_service/v2/chart/exchange/NSE/segment/CASH/{ticker}?endTimeInMillis={end_ts}&intervalInMinutes=1&startTimeInMillis={start_ts}"
    headers = {
        "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64)",
        "Accept": "application/json"
    }
    
    try:
        response = requests.get(url, headers=headers, timeout=10)
        if response.status_code == 200:
            data = response.json()
            if data and 'candles' in data and len(data['candles']) > 0:
                # Candles format: [timestamp, open, high, low, close, volume]
                last_candle = data['candles'][-1]
                return last_candle[4] # Close price
        return None
    except Exception as e:
        print(f"Error fetching {ticker}: {e}")
        return None

def update_daily_data():
    now = datetime.now(IST)
    today = now.date()
    
    if not is_trading_day(today):
        print(f"{today} is not a trading day. Skipping update.")
        return {"status": "skipped", "reason": "not a trading day"}
        
    if not os.path.exists(DATA_FILE):
        print(f"Data file {DATA_FILE} not found!")
        return {"status": "error", "reason": "data file missing"}
        
    # Read existing data
    df = pd.read_parquet(DATA_FILE)
    
    # Check if we already updated today
    if not df.empty and pd.to_datetime(today) in df['date'].dt.date.values:
        # We might have partial data or want to overwrite, but for safety:
        # Let's delete today's entries if they exist so we can cleanly append
        df = df[df['date'].dt.date != today]
        
    print(f"Fetching data for {today}...")
    
    # Market hours: 09:15 to 15:30 IST
    start_dt = datetime.combine(today, datetime.strptime("09:15", "%H:%M").time()).replace(tzinfo=IST)
    end_dt = datetime.combine(today, datetime.strptime("15:30", "%H:%M").time()).replace(tzinfo=IST)
    
    start_ts = int(start_dt.timestamp() * 1000)
    end_ts = int(end_dt.timestamp() * 1000)
    
    new_rows = []
    
    for ticker in TICKERS:
        close_price = fetch_groww_data(ticker, start_ts, end_ts)
        if close_price is not None:
            new_rows.append({
                'date': pd.to_datetime(today),
                'close': float(close_price),
                'ticker': ticker
            })
        time.sleep(0.5) # Rate limiting
        
    if new_rows:
        new_df = pd.DataFrame(new_rows)
        updated_df = pd.concat([df, new_df], ignore_index=True)
        updated_df.sort_values(by=['ticker', 'date'], inplace=True)
        updated_df.to_parquet(DATA_FILE)
        print(f"Successfully updated {len(new_rows)} tickers for {today}")
        return {"status": "success", "updated_count": len(new_rows)}
    else:
        print("No new data fetched.")
        return {"status": "error", "reason": "fetch failed for all tickers"}

if __name__ == "__main__":
    update_daily_data()