"""
Real-time market data integration for Indian and global stocks
Optimized for Hugging Face free tier
"""

import requests
import yfinance as yf
import pandas as pd
import numpy as np
from typing import Dict, List, Optional, Tuple
import time
from datetime import datetime, timedelta

class IndianStockAPI:
    """
    Free Indian Stock Market API client for NSE/BSE data
    No API key required [citation:1]
    """
    
    BASE_URL = "https://military-jobye-haiqstudios-14f59639.koyeb.app"
    
    def __init__(self):
        self.session = requests.Session()
        self.cache = {}
        self.cache_timestamp = {}
        self.cache_duration = 60  # Cache for 60 seconds to avoid rate limits
    
    def search_stocks(self, query: str) -> List[Dict]:
        """
        Search for stocks by company name or symbol
        Example: /search?q=reliance [citation:1]
        """
        try:
            url = f"{self.BASE_URL}/search"
            params = {'q': query}
            response = self.session.get(url, params=params, timeout=10)
            
            if response.status_code == 200:
                data = response.json()
                if data.get('status') == 'success':
                    return data.get('results', [])
            return []
        except Exception as e:
            print(f"Search error: {e}")
            return []
    
    def get_stock_price(self, symbol: str, exchange: str = "NSE") -> Optional[Dict]:
        """
        Get real-time stock price and details
        Examples: 
          - ITC → NSE (default) [citation:1]
          - ITC.NS → NSE explicit
          - ITC.BO → BSE
        """
        # Check cache first
        cache_key = f"{symbol}_{exchange}"
        if cache_key in self.cache:
            if time.time() - self.cache_timestamp[cache_key] < self.cache_duration:
                return self.cache[cache_key]
        
        try:
            # Format symbol with exchange suffix
            if exchange.upper() == "NSE":
                api_symbol = f"{symbol}.NS" if not symbol.endswith('.NS') else symbol
            else:  # BSE
                api_symbol = f"{symbol}.BO" if not symbol.endswith('.BO') else symbol
            
            url = f"{self.BASE_URL}/stock"
            params = {
                'symbol': api_symbol,
                'res': 'val'  # Get values with units [citation:1]
            }
            
            response = self.session.get(url, params=params, timeout=10)
            
            if response.status_code == 200:
                data = response.json()
                if data.get('status') == 'success':
                    # Cache the result
                    self.cache[cache_key] = data
                    self.cache_timestamp[cache_key] = time.time()
                    return data
            return None
        except Exception as e:
            print(f"Stock price error: {e}")
            return None
    
    def get_multiple_stocks(self, symbols: List[str], exchange: str = "NSE") -> Dict:
        """
        Get prices for multiple stocks efficiently
        """
        results = {}
        for symbol in symbols:
            data = self.get_stock_price(symbol, exchange)
            if data:
                results[symbol] = data
            time.sleep(0.5)  # Rate limiting protection
        return results
    
    def format_for_index(self, stock_data: Dict) -> Dict:
        """
        Convert stock API response to index calculator format
        """
        if not stock_data or 'data' not in stock_data:
            return {}
        
        d = stock_data['data']
        
        # Handle both numeric and unit formats [citation:1]
        def extract_value(field):
            if isinstance(field, dict):
                return field.get('value', 0)
            return field
        
        return {
            'company_name': d.get('company_name', ''),
            'last_price': extract_value(d.get('last_price', 0)),
            'change': extract_value(d.get('change', 0)),
            'percent_change': extract_value(d.get('percent_change', 0)),
            'previous_close': extract_value(d.get('previous_close', 0)),
            'open': extract_value(d.get('open', 0)),
            'day_high': extract_value(d.get('day_high', 0)),
            'day_low': extract_value(d.get('day_low', 0)),
            'year_high': extract_value(d.get('year_high', 0)),
            'year_low': extract_value(d.get('year_low', 0)),
            'volume': extract_value(d.get('volume', 0)),
            'market_cap': extract_value(d.get('market_cap', 0)),
            'pe_ratio': extract_value(d.get('pe_ratio', 0)),
            'dividend_yield': extract_value(d.get('dividend_yield', 0)),
            'sector': d.get('sector', ''),
            'currency': d.get('currency', 'INR')
        }


class GlobalStockAPI:
    """
    Yahoo Finance client for global stocks and indices
    Uses yfinance library [citation:2]
    """
    
    def __init__(self):
        self.cache = {}
        self.cache_timestamp = {}
        self.cache_duration = 60  # Cache for 60 seconds
    
    def get_stock_info(self, symbol: str) -> Optional[Dict]:
        """
        Get real-time stock data from Yahoo Finance
        Examples: 'RELIANCE.NS', 'AAPL', '^NSEI' (Nifty 50)
        """
        # Check cache
        if symbol in self.cache:
            if time.time() - self.cache_timestamp[symbol] < self.cache_duration:
                return self.cache[symbol]
        
        try:
            ticker = yf.Ticker(symbol)
            info = ticker.info
            
            # Get current price (fastest method)
            hist = ticker.history(period="1d", interval="1m")
            
            if not hist.empty:
                current_price = hist['Close'].iloc[-1]
            else:
                current_price = info.get('regularMarketPrice', info.get('currentPrice', 0))
            
            result = {
                'symbol': symbol,
                'company_name': info.get('longName', info.get('shortName', symbol)),
                'last_price': current_price,
                'previous_close': info.get('regularMarketPreviousClose', 0),
                'open': info.get('regularMarketOpen', 0),
                'day_high': info.get('dayHigh', info.get('regularMarketDayHigh', 0)),
                'day_low': info.get('dayLow', info.get('regularMarketDayLow', 0)),
                'volume': info.get('volume', info.get('regularMarketVolume', 0)),
                'market_cap': info.get('marketCap', 0),
                'pe_ratio': info.get('trailingPE', 0),
                'dividend_yield': info.get('dividendYield', 0) * 100 if info.get('dividendYield') else 0,
                'sector': info.get('sector', ''),
                'industry': info.get('industry', ''),
                'currency': info.get('currency', 'USD'),
                'timestamp': datetime.now().isoformat()
            }
            
            # Cache the result
            self.cache[symbol] = result
            self.cache_timestamp[symbol] = time.time()
            
            return result
            
        except Exception as e:
            print(f"YFinance error for {symbol}: {e}")
            return None
    
    def get_nifty50_components(self) -> List[str]:
        """
        Get Nifty 50 constituent symbols
        """
        # Common Nifty 50 stocks with .NS suffix for Yahoo Finance
        return [
            'RELIANCE.NS', 'TCS.NS', 'HDFCBANK.NS', 'INFY.NS', 'ICICIBANK.NS',
            'HINDUNILVR.NS', 'ITC.NS', 'SBIN.NS', 'BHARTIARTL.NS', 'KOTAKBANK.NS',
            'BAJFINANCE.NS', 'LT.NS', 'WIPRO.NS', 'AXISBANK.NS', 'TITAN.NS',
            'ASIANPAINT.NS', 'MARUTI.NS', 'SUNPHARMA.NS', 'HCLTECH.NS', 'ULTRACEMCO.NS'
        ]
    
    def get_multiple_stocks(self, symbols: List[str]) -> Dict:
        """
        Get data for multiple symbols efficiently
        """
        results = {}
        for symbol in symbols:
            data = self.get_stock_info(symbol)
            if data:
                results[symbol] = data
        return results


class IndexBuilder:
    """
    Build custom indices from real-time stock data
    """
    
    def __init__(self):
        self.indian_api = IndianStockAPI()
        self.global_api = GlobalStockAPI()
    
    def create_price_weighted_index(self, symbols: List[str], exchange: str = "NSE") -> Dict:
        """
        Create a price-weighted index (like Dow Jones)
        Index = (Sum of stock prices) / Divisor
        """
        stocks_data = []
        total_price = 0
        
        for symbol in symbols:
            if exchange in ["NSE", "BSE"]:
                data = self.indian_api.get_stock_price(symbol, exchange)
                if data and 'data' in data:
                    price_data = self.indian_api.format_for_index(data)
                    stocks_data.append(price_data)
                    total_price += price_data['last_price']
            else:
                data = self.global_api.get_stock_info(symbol)
                if data:
                    stocks_data.append(data)
                    total_price += data['last_price']
        
        # Simple divisor (number of stocks for demo)
        divisor = len(stocks_data)
        index_value = total_price / divisor if divisor > 0 else 0
        
        return {
            'index_name': f"Price-Weighted Index ({exchange})",
            'index_value': round(index_value, 2),
            'components': len(stocks_data),
            'total_price': round(total_price, 2),
            'divisor': divisor,
            'stocks': stocks_data,
            'timestamp': datetime.now().isoformat()
        }
    
    def create_market_cap_weighted_index(self, symbols: List[str], exchange: str = "NSE") -> Dict:
        """
        Create a market capitalization-weighted index (like S&P 500, Nifty 50)
        Index = (Sum of market caps) / Base market cap * Base index value
        """
        stocks_data = []
        total_market_cap = 0
        
        for symbol in symbols:
            if exchange in ["NSE", "BSE"]:
                data = self.indian_api.get_stock_price(symbol, exchange)
                if data and 'data' in data:
                    price_data = self.indian_api.format_for_index(data)
                    stocks_data.append(price_data)
                    total_market_cap += price_data['market_cap']
            else:
                data = self.global_api.get_stock_info(symbol)
                if data:
                    stocks_data.append(data)
                    total_market_cap += data['market_cap']
        
        # Use first stock's market cap as base for demo
        base_market_cap = stocks_data[0]['market_cap'] if stocks_data else 1
        index_value = (total_market_cap / base_market_cap) * 100 if base_market_cap > 0 else 0
        
        return {
            'index_name': f"Market Cap-Weighted Index ({exchange})",
            'index_value': round(index_value, 2),
            'components': len(stocks_data),
            'total_market_cap': total_market_cap,
            'base_market_cap': base_market_cap,
            'stocks': stocks_data,
            'timestamp': datetime.now().isoformat()
        }
    
    def calculate_sector_index(self, sector: str, exchange: str = "NSE") -> Dict:
        """
        Create an index for a specific sector
        """
        # This would normally filter stocks by sector
        # For demo, we'll use a few known stocks
        sector_stocks = {
            'technology': ['TCS.NS', 'INFY.NS', 'WIPRO.NS', 'HCLTECH.NS'],
            'banking': ['HDFCBANK.NS', 'ICICIBANK.NS', 'SBIN.NS', 'KOTAKBANK.NS'],
            'energy': ['RELIANCE.NS', 'ONGC.NS', 'POWERGRID.NS']
        }
        
        symbols = sector_stocks.get(sector.lower(), [])
        if not symbols:
            return {'error': f'Sector {sector} not found'}
        
        return self.create_market_cap_weighted_index(symbols, exchange)


class RealTimeIndexCalculator:
    """
    Main class that combines real-time data with index calculations
    """
    
    def __init__(self):
        self.index_builder = IndexBuilder()
        self.indian_api = IndianStockAPI()
        self.global_api = GlobalStockAPI()
        from index_calculator import IndexCalculator
        self.calculator = IndexCalculator()
    
    def calculate_live_laspeyres(self, symbols: List[str], base_prices: List[float], 
                                 base_quantities: List[float], exchange: str = "NSE") -> Dict:
        """
        Calculate Laspeyres index using live prices
        """
        # Get live prices
        live_prices = []
        stock_details = []
        
        for i, symbol in enumerate(symbols):
            if exchange in ["NSE", "BSE"]:
                data = self.indian_api.get_stock_price(symbol, exchange)
                if data and 'data' in data:
                    price_data = self.indian_api.format_for_index(data)
                    live_prices.append(price_data['last_price'])
                    stock_details.append(price_data)
            else:
                data = self.global_api.get_stock_info(symbol)
                if data:
                    live_prices.append(data['last_price'])
                    stock_details.append(data)
        
        # Create DataFrame for index calculator
        df = pd.DataFrame({
            'item': symbols,
            'p0': base_prices[:len(live_prices)],
            'q0': base_quantities[:len(live_prices)],
            'p1': live_prices
        })
        
        self.calculator.load_data(df)
        laspeyres = self.calculator.calculate_laspeyres()
        
        return {
            'index_value': laspeyres['value'],
            'interpretation': laspeyres['interpretation'],
            'live_prices': live_prices,
            'stock_details': stock_details,
            'timestamp': datetime.now().isoformat()
        }
    
    def compare_indices(self, symbols: List[str], exchange: str = "NSE") -> pd.DataFrame:
        """
        Compare multiple index calculations on the same basket
        """
        # Get live data
        stocks_data = []
        for symbol in symbols:
            if exchange in ["NSE", "BSE"]:
                data = self.indian_api.get_stock_price(symbol, exchange)
                if data and 'data' in data:
                    stocks_data.append(self.indian_api.format_for_index(data))
            else:
                data = self.global_api.get_stock_info(symbol)
                if data:
                    stocks_data.append(data)
        
        if not stocks_data:
            return pd.DataFrame()
        
        # Prepare data for index calculator
        df = pd.DataFrame({
            'item': [s.get('company_name', f'Stock{i}') for i, s in enumerate(stocks_data)],
            'p0': [100.0] * len(stocks_data),  # Assume base price 100
            'q0': [1.0] * len(stocks_data),    # Equal quantities
            'p1': [s['last_price'] for s in stocks_data]
        })
        
        self.calculator.load_data(df)
        results = self.calculator.calculate_all()
        
        # Add live data
        results['live_data'] = stocks_data
        results['symbols'] = symbols
        results['exchange'] = exchange
        
        return results