import os
import gradio as gr
import pandas as pd
import numpy as np
import plotly.graph_objects as go
import yfinance as yf
from datetime import datetime
from typing import List, Tuple, Dict, Optional
import json

# ========== Optional memory monitoring ==========
try:
    import psutil
    print(f"🧠 Memory at startup: {psutil.virtual_memory().percent}% used, "
          f"{psutil.virtual_memory().available / 1e9:.2f} GB free")
except ImportError:
    pass

# ========== Lazy LLM import ==========
def get_llm():
    from langchain_openai import ChatOpenAI
    from langchain_core.messages import HumanMessage, AIMessage, SystemMessage
    api_key = os.getenv("OPENROUTER_API_KEY")
    if not api_key:
        raise ValueError("OPENROUTER_API_KEY environment variable not set.")
    return ChatOpenAI(
        base_url="https://openrouter.ai/api/v1",
        api_key=api_key,
        model="openai/gpt-3.5-turbo:free",
        temperature=0.3,
    )

# ========== Stock API classes using yfinance ==========
class IndianStockAPI:
    """Fetch Indian stock data via yfinance – expects symbol with .NS or .BO suffix."""
    
    def get_stock_price(self, symbol: str, exchange: str = "NSE"):
        """
        Get stock data for Indian market.
        symbol should already include .NS or .BO suffix.
        """
        try:
            ticker = yf.Ticker(symbol)
            info = ticker.info
            hist = ticker.history(period="1d")
            if hist.empty:
                return None
            
            last_price = hist['Close'].iloc[-1]
            prev_close = info.get('regularMarketPreviousClose', last_price)
            change = last_price - prev_close
            pct_change = (change / prev_close) * 100 if prev_close else 0
            
            return {
                'data': [{
                    'symbol': symbol.replace('.NS', '').replace('.BO', ''),  # strip suffix for display
                    'last_price': last_price,
                    'change': change,
                    'percent_change': pct_change,
                    'day_low': hist['Low'].iloc[-1],
                    'day_high': hist['High'].iloc[-1],
                    'year_low': info.get('fiftyTwoWeekLow', 0),
                    'year_high': info.get('fiftyTwoWeekHigh', 0),
                    'market_cap': info.get('marketCap', 0) / 1e7,  # in Crores
                    'pe_ratio': info.get('trailingPE', 0),
                    'company_name': info.get('longName', symbol)
                }]
            }
        except Exception as e:
            print(f"Error fetching Indian stock {symbol}: {e}")
            return None
    
    def format_for_index(self, data):
        """Format API response to match expected structure."""
        if data and 'data' in data and len(data['data']) > 0:
            return data['data'][0]
        return {}

class GlobalStockAPI:
    """Fetch global stock data via yfinance."""
    
    def get_stock_info(self, symbol: str):
        """Get stock data for global markets."""
        try:
            ticker = yf.Ticker(symbol)
            info = ticker.info
            hist = ticker.history(period="1d")
            if hist.empty:
                return None
            
            last_price = hist['Close'].iloc[-1]
            prev_close = info.get('regularMarketPreviousClose', last_price)
            change = last_price - prev_close
            pct_change = (change / prev_close) * 100 if prev_close else 0
            
            return {
                'symbol': symbol,
                'company_name': info.get('longName', symbol),
                'last_price': last_price,
                'change': change,
                'percent_change': pct_change,
                'day_low': hist['Low'].iloc[-1],
                'day_high': hist['High'].iloc[-1],
                'market_cap': info.get('marketCap', 0),
            }
        except Exception as e:
            print(f"Error fetching global stock {symbol}: {e}")
            return None

class RealTimeIndexCalculator:
    """Placeholder for index calculations (not heavily used)."""
    def __init__(self):
        pass

# ========== Helper to get API instances ==========
def get_realtime_calculator():
    return RealTimeIndexCalculator(), IndianStockAPI(), GlobalStockAPI()

# ========== System prompt ==========
SYSTEM_PROMPT = """You are an expert financial analyst specializing in index numbers and real-time market data. 
You help users understand, calculate, and analyze various types of indices including:
- Laspeyres, Paasche, and Fisher price indices
- Real-time market indices (Nifty 50, Sensex, S&P 500)
- Custom indices built from live stock data
- Sector-specific indices
You have access to:
1. Real-time Indian stock data (NSE/BSE) via free API [citation:1]
2. Global stock data via Yahoo Finance [citation:2]
3. A powerful index calculator for custom analysis
When users ask about stocks or indices, fetch live data. Provide clear explanations with numbers."""

# Mapping for common company names to tickers (for global stocks)
GLOBAL_TICKER_MAP = {
    'APPLE': 'AAPL',
    'GOOGLE': 'GOOGL',
    'MICROSOFT': 'MSFT',
    'AMAZON': 'AMZN',
    'TESLA': 'TSLA',
    'META': 'META',
    'FACEBOOK': 'META',
    'NETFLIX': 'NFLX',
}

# ========== Core conversational function ==========
def process_query(query: str, history: List[Dict], exchange: str) -> Tuple[str, List[Dict], pd.DataFrame, str]:
    """
    Process user query with real-time data awareness.
    history is a list of dicts: [{"role": "user", "content": "..."}, {"role": "assistant", "content": "..."}, ...]
    Returns:
        - response_text: str (the assistant's reply)
        - updated_history: list of dicts (new chat history)
        - data_preview: pd.DataFrame
        - status: str (unused)
    """
    rt_calc, indian_api, global_api = get_realtime_calculator()
    response_text = ""
    data_preview = pd.DataFrame()
    
    # Common Indian stocks (base names, without suffix)
    indian_stocks = ['RELIANCE', 'TCS', 'INFY', 'HDFCBANK', 'ICICIBANK', 'SBIN', 
                     'ITC', 'BHARTIARTL', 'KOTAKBANK', 'LT', 'WIPRO', 'AXISBANK']
    
    query_upper = query.upper()
    query_lower = query.lower()
    
    # Real-time data handling
    if any(stock in query_upper for stock in indian_stocks) or 'NIFTY' in query_upper or 'SENSEX' in query_upper:
        for stock in indian_stocks:
            if stock in query_upper:
                try:
                    # Build symbol with correct suffix for Indian exchange
                    if exchange == "NSE":
                        yf_symbol = stock + ".NS"
                    elif exchange == "BSE":
                        yf_symbol = stock + ".BO"
                    else:
                        yf_symbol = stock  # fallback (should not happen)
                    
                    data = indian_api.get_stock_price(yf_symbol, exchange)
                    if data and 'data' in data:
                        formatted = indian_api.format_for_index(data)
                        response_text = f"**{formatted['company_name']}**\n\n"
                        response_text += f"💰 Current Price: ₹{formatted['last_price']:,.2f}\n"
                        response_text += f"📊 Change: {formatted['change']:+.2f} ({formatted['percent_change']:+.2f}%)\n"
                        response_text += f"📈 Day Range: ₹{formatted['day_low']} - ₹{formatted['day_high']}\n"
                        response_text += f"📉 52-Week Range: ₹{formatted['year_low']} - ₹{formatted['year_high']}\n"
                        response_text += f"💹 Market Cap: ₹{formatted['market_cap']:,.0f} Cr\n"
                        response_text += f"📊 P/E Ratio: {formatted['pe_ratio']:.2f}\n"
                        data_preview = pd.DataFrame([formatted])
                except Exception as e:
                    response_text = f"Error fetching data: {str(e)}"
                break
        
        if 'NIFTY' in query_upper and not response_text:
            try:
                nifty_data = global_api.get_stock_info('^NSEI')
                if nifty_data:
                    response_text = f"**Nifty 50 Index**\nCurrent: {nifty_data['last_price']:,.2f}\n"
                    response_text += f"Change: {nifty_data['change']:+.2f} ({nifty_data['percent_change']:+.2f}%)"
            except Exception as e:
                response_text = f"Error fetching Nifty: {str(e)}"
    
    # Check for global stock tickers or company names
    else:
        # Try to find a ticker in the query
        ticker = None
        # First check if any known ticker appears directly
        known_tickers = ['AAPL', 'GOOGL', 'MSFT', 'AMZN', 'TSLA', 'META', 'NFLX']
        for t in known_tickers:
            if t in query_upper:
                ticker = t
                break
        
        # If not found, check for company names
        if not ticker:
            for company, t in GLOBAL_TICKER_MAP.items():
                if company in query_upper:
                    ticker = t
                    break
        
        if ticker:
            try:
                data = global_api.get_stock_info(ticker)
                if data:
                    response_text = f"**{data['company_name']} ({ticker})**\n\n"
                    response_text += f"💰 Current Price: ${data['last_price']:,.2f}\n"
                    response_text += f"📊 Change: ${data['change']:+.2f} ({data['percent_change']:+.2f}%)\n"
                    response_text += f"📈 Day Range: ${data['day_low']} - ${data['day_high']}\n"
                    response_text += f"💹 Market Cap: ${data['market_cap']/1e9:,.2f}B\n"
                    data_preview = pd.DataFrame([data])
            except Exception as e:
                response_text = f"Error fetching data: {str(e)}"
    
    # If no real-time match, use LLM
    if not response_text:
        try:
            llm = get_llm()
            from langchain_core.messages import HumanMessage, AIMessage, SystemMessage
            
            # Convert Gradio history (list of dicts) to flat list for LangChain
            flat_history = []
            for turn in history:
                flat_history.append(turn["content"])
            
            messages = [SystemMessage(content=SYSTEM_PROMPT)]
            for i, msg in enumerate(flat_history):
                if i % 2 == 0:
                    messages.append(HumanMessage(content=msg))
                else:
                    messages.append(AIMessage(content=msg))
            messages.append(HumanMessage(content=query))
            
            response = llm.invoke(messages)
            response_text = response.content
        except Exception as e:
            response_text = f"LLM error: {str(e)}"
    
    # Update history in Gradio's dictionary format (list of dicts)
    updated_history = history + [
        {"role": "user", "content": query},
        {"role": "assistant", "content": response_text}
    ]
    # Keep only last 10 exchanges (20 messages)
    if len(updated_history) > 20:
        updated_history = updated_history[-20:]
    
    return response_text, updated_history, data_preview, ""

# ========== Real-time UI functions ==========
def fetch_live_stock(symbol: str, exchange: str):
    """Fetch live stock data and return JSON, price, change, pct, chart."""
    _, indian_api, global_api = get_realtime_calculator()
    if exchange in ["NSE", "BSE"]:
        # Construct symbol with suffix
        yf_symbol = symbol + ".NS" if exchange == "NSE" else symbol + ".BO"
        data = indian_api.get_stock_price(yf_symbol, exchange)
        if data and 'data' in data:
            formatted = indian_api.format_for_index(data)
            fig = go.Figure()
            fig.add_trace(go.Indicator(
                mode="number+delta",
                value=formatted['last_price'],
                number={'prefix': '₹'},
                delta={'reference': formatted['last_price'] - formatted['change']}
            ))
            fig.update_layout(height=300)
            return formatted, formatted['last_price'], formatted['change'], formatted['percent_change'], fig
    else:
        # For global, try to map common names
        ticker = symbol.upper()
        if ticker in GLOBAL_TICKER_MAP.values() or ticker in ['AAPL', 'GOOGL', 'MSFT', 'AMZN', 'TSLA', 'META', 'NFLX']:
            data = global_api.get_stock_info(ticker)
            if data:
                fig = go.Figure()
                fig.add_trace(go.Indicator(
                    mode="number+delta",
                    value=data['last_price'],
                    number={'prefix': '$'},
                    delta={'reference': data['last_price'] - data['change']}
                ))
                fig.update_layout(height=300)
                return data, data['last_price'], data['change'], data['percent_change'], fig
    return {}, 0, 0, 0, go.Figure()

def build_custom_index(symbols_df: pd.DataFrame, index_type: str, exchange: str):
    """Build custom index from symbols and weights."""
    _, indian_api, global_api = get_realtime_calculator()
    weights = []
    prices = []
    total_weight = 0
    for _, row in symbols_df.iterrows():
        symbol = row['Symbol']
        weight = float(row['Weight (%)']) / 100.0
        if exchange in ["NSE", "BSE"]:
            # Add suffix for Yahoo Finance
            yf_symbol = symbol + ".NS" if exchange == "NSE" else symbol + ".BO"
            data = indian_api.get_stock_price(yf_symbol, exchange)
            if data and 'data' in data:
                price = data['data'][0]['last_price']
            else:
                price = 0
        else:
            # For global, try direct ticker
            data = global_api.get_stock_info(symbol)
            price = data.get('last_price', 0) if data else 0
        weights.append(weight)
        prices.append(price)
        total_weight += weight
    
    # Normalize weights if total != 1
    if total_weight != 1 and total_weight > 0:
        weights = [w/total_weight for w in weights]
    
    if prices:
        if index_type == "Price-Weighted":
            index_value = sum(prices) / len(prices)
        elif index_type == "Equal-Weighted":
            index_value = sum(prices) / len(prices)
        elif index_type == "Market Cap-Weighted":
            index_value = sum(w * p for w, p in zip(weights, prices))
        else:
            index_value = 0
    else:
        index_value = 0
    
    return {"index_value": round(index_value, 2), "components": prices, "weights": weights}

# ========== Index Calculator functions ==========
def calculate_indices(df: pd.DataFrame):
    """
    Compute Laspeyres, Paasche, and Fisher indices from input dataframe.
    Expected columns: Item, p0, q0, p1, q1
    """
    if df.empty:
        return "Please enter data.", 0, 0, 0, go.Figure()
    
    # Convert to numeric, handle errors
    try:
        df['p0'] = pd.to_numeric(df['p0'], errors='coerce')
        df['q0'] = pd.to_numeric(df['q0'], errors='coerce')
        df['p1'] = pd.to_numeric(df['p1'], errors='coerce')
        df['q1'] = pd.to_numeric(df['q1'], errors='coerce')
    except Exception as e:
        return f"Error converting data: {e}", 0, 0, 0, go.Figure()
    
    # Drop rows with NaN
    df = df.dropna()
    if df.empty:
        return "No valid numeric data.", 0, 0, 0, go.Figure()
    
    # Calculate sums
    sum_p0q0 = (df['p0'] * df['q0']).sum()
    sum_p1q0 = (df['p1'] * df['q0']).sum()
    sum_p0q1 = (df['p0'] * df['q1']).sum()
    sum_p1q1 = (df['p1'] * df['q1']).sum()
    
    # Avoid division by zero
    if sum_p0q0 == 0 or sum_p0q1 == 0:
        return "Base period expenditure is zero, cannot compute indices.", 0, 0, 0, go.Figure()
    
    laspeyres = (sum_p1q0 / sum_p0q0) * 100
    paasche = (sum_p1q1 / sum_p0q1) * 100
    fisher = np.sqrt(laspeyres * paasche)
    
    # Create comparison chart
    fig = go.Figure(data=[
        go.Bar(name='Index Value', x=['Laspeyres', 'Paasche', 'Fisher'], 
               y=[laspeyres, paasche, fisher])
    ])
    fig.update_layout(title="Price Index Comparison", yaxis_title="Index Value")
    
    result_text = f"**Laspeyres Index:** {laspeyres:.2f}\n**Paasche Index:** {paasche:.2f}\n**Fisher Index:** {fisher:.2f}"
    return result_text, laspeyres, paasche, fisher, fig

# ========== Comparison UI functions ==========
def create_comparison_ui():
    """Return UI components for index comparison tab."""
    with gr.Blocks() as comp_ui:
        with gr.Row():
            symbols_input = gr.Textbox(label="Symbols (comma-separated)", placeholder="RELIANCE, TCS, INFY")
            exchange = gr.Radio(["NSE", "BSE", "Global"], value="NSE", label="Exchange")
        compare_btn = gr.Button("Compare Indices")
        results_output = gr.JSON(label="Comparison Results")
        chart_output = gr.Plot(label="Comparison Chart")
    return {"symbols": symbols_input, "exchange": exchange, "btn": compare_btn,
            "results": results_output, "chart": chart_output}

def compare_live_indices(symbols_str: str, exchange: str):
    """Compare multiple indices."""
    symbols = [s.strip() for s in symbols_str.split(",") if s.strip()]
    _, indian_api, global_api = get_realtime_calculator()
    data = {}
    for sym in symbols:
        if exchange in ["NSE", "BSE"]:
            yf_symbol = sym + ".NS" if exchange == "NSE" else sym + ".BO"
            d = indian_api.get_stock_price(yf_symbol, exchange)
            if d and 'data' in d and d['data']:
                data[sym] = d['data'][0]['last_price']
        else:
            d = global_api.get_stock_info(sym)
            if d:
                data[sym] = d['last_price']
    
    fig = go.Figure(data=[go.Bar(x=list(data.keys()), y=list(data.values()))])
    fig.update_layout(title="Index Comparison", yaxis_title="Price")
    return data, fig

# ========== Main Gradio interface ==========
with gr.Blocks(title="AI-Powered Index Analyst with Real-Time Data") as demo:
    gr.Markdown("""
    # 📈 AI-Powered Index Analyst with Real-Time Market Data
    
    Welcome to your intelligent financial index assistant! This enhanced version includes:
    - **Real-time Indian stock data** from NSE/BSE via Yahoo Finance
    - **Global market data** via Yahoo Finance
    - **Custom index builder** with live prices
    - **Conversational AI** for market insights
    - **Multiple index formulas** (Laspeyres, Paasche, Fisher, etc.)
    
    Built with ❤️ for CFA candidates and finance professionals
    """)
    
    with gr.Tab("💬 Conversational Assistant"):
        chat_history = gr.Chatbot(label="Conversation", height=400)
        msg = gr.Textbox(
            label="Ask about stocks or indices",
            placeholder="E.g., 'What's the price of RELIANCE?' or 'Compare Nifty 50 index using Laspeyres and Fisher'"
        )
        
        with gr.Row():
            exchange_selector = gr.Radio(
                choices=["NSE", "BSE", "Global"],
                value="NSE",
                label="Exchange",
                interactive=True
            )
            clear = gr.Button("Clear Conversation")
            submit = gr.Button("Submit", variant="primary")
        
        data_output = gr.DataFrame(label="Data Preview")
        calc_output = gr.Textbox(label="Calculation Results", lines=3)
        
        submit.click(
            process_query,
            inputs=[msg, chat_history, exchange_selector],
            outputs=[calc_output, chat_history, data_output, calc_output]
        )
        
        msg.submit(
            process_query,
            inputs=[msg, chat_history, exchange_selector],
            outputs=[calc_output, chat_history, data_output, calc_output]
        )
        
        clear.click(
            lambda: ([], pd.DataFrame(), ""),
            outputs=[chat_history, data_output, calc_output]
        )
    
    with gr.Tab("📡 Real-Time Market Data"):
        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown("""
                ### 📡 Real-Time Market Data
                
                Fetch live stock prices from:
                - **NSE** (National Stock Exchange of India) via Yahoo Finance
                - **BSE** (Bombay Stock Exchange) via Yahoo Finance
                - **Global Markets** via Yahoo Finance
                """)
                
                exchange = gr.Radio(
                    choices=["NSE", "BSE", "Global"],
                    value="NSE",
                    label="Select Exchange"
                )
                
                symbol_input = gr.Textbox(
                    label="Stock Symbol",
                    placeholder="RELIANCE, TCS, INFY, or AAPL, GOOGL for global"
                )
                
                fetch_btn = gr.Button("Fetch Live Data", variant="primary")
                
                gr.Markdown("""
                ### 📊 Build Custom Index
                Enter up to 5 symbols to create your own index
                """)
                
                symbols_list = gr.Dataframe(
                    headers=["Symbol", "Weight (%)"],
                    datatype=["str", "number"],
                    row_count=5,
                    col_count=2,
                    value=[["RELIANCE", 20], ["TCS", 20], ["INFY", 20], ["HDFCBANK", 20], ["ITC", 20]]
                )
                
                index_type = gr.Dropdown(
                    choices=["Price-Weighted", "Market Cap-Weighted", "Equal-Weighted"],
                    value="Market Cap-Weighted",
                    label="Index Type"
                )
                
                build_index_btn = gr.Button("Build Custom Index", variant="secondary")
            
            with gr.Column(scale=1):
                stock_output = gr.JSON(label="Stock Details")
                
                with gr.Row():
                    price = gr.Number(label="Current Price")
                    change = gr.Number(label="Change")
                    pct_change = gr.Number(label="% Change")
                
                chart_output = gr.Plot(label="Price Chart")
                
                index_result = gr.JSON(label="Custom Index Result")
        
        fetch_btn.click(
            fetch_live_stock,
            inputs=[symbol_input, exchange],
            outputs=[stock_output, price, change, pct_change, chart_output]
        )
        
        build_index_btn.click(
            build_custom_index,
            inputs=[symbols_list, index_type, exchange],
            outputs=[index_result]
        )
    
    with gr.Tab("📊 Index Calculator"):
        gr.Markdown("""
        ### Price Index Calculator
        Enter items with base period prices (p0) and quantities (q0), and current period prices (p1) and quantities (q1).
        The calculator will compute Laspeyres, Paasche, and Fisher indices.
        """)
        
        index_df = gr.Dataframe(
            headers=["Item", "p0", "q0", "p1", "q1"],
            datatype=["str", "number", "number", "number", "number"],
            row_count=5,
            col_count=5,
            value=[["Item1", 10, 100, 12, 110],
                   ["Item2", 20, 50, 22, 55],
                   ["Item3", 30, 30, 33, 32],
                   ["Item4", 40, 20, 44, 18],
                   ["Item5", 50, 10, 55, 9]]
        )
        
        calc_index_btn = gr.Button("Calculate Indices", variant="primary")
        
        with gr.Row():
            laspeyres_out = gr.Number(label="Laspeyres Index")
            paasche_out = gr.Number(label="Paasche Index")
            fisher_out = gr.Number(label="Fisher Index")
        
        index_results = gr.Textbox(label="Detailed Results", lines=4)
        index_plot = gr.Plot(label="Index Comparison Chart")
        
        calc_index_btn.click(
            calculate_indices,
            inputs=[index_df],
            outputs=[index_results, laspeyres_out, paasche_out, fisher_out, index_plot]
        )
    
    with gr.Tab("🔄 Index Comparison"):
        comp_ui = create_comparison_ui()
        comp_ui["btn"].click(
            compare_live_indices,
            inputs=[comp_ui["symbols"], comp_ui["exchange"]],
            outputs=[comp_ui["results"], comp_ui["chart"]]
        )
    
    gr.Markdown("""
    ---
    ### 📚 Resources & Credits
    - Data provided by [Yahoo Finance](https://finance.yahoo.com/) via [yfinance](https://pypi.org/project/yfinance/)
    - [CFA Institute](https://www.cfainstitute.org)
    
    *Note: Free APIs may have rate limits. Data is cached for performance.*
    """)

if __name__ == "__main__":
    demo.launch()