import streamlit as st
import tensorflow as tf
import numpy as np
import pandas as pd
import pickle
import json
import requests
import yfinance as yf
from tensorflow.keras.preprocessing.sequence import pad_sequences

# =====================================================
# PAGE CONFIG
# =====================================================

st.set_page_config(
    page_title="Stock Market AI Chatbot",
    page_icon="🤖",
    layout="centered"
)

# =====================================================
# LOAD MODEL
# =====================================================

from pathlib import Path
import tensorflow as tf
import pickle
import json
import streamlit as st

BASE_DIR = Path(__file__).resolve().parent.parent
MODEL_DIR = BASE_DIR / "saved_chatbot"

@st.cache_resource
def load_chatbot():

    encoder = tf.keras.models.load_model(
        MODEL_DIR / "encoder.keras",
        compile=False
    )

    decoder = tf.keras.models.load_model(
        MODEL_DIR / "decoder.keras",
        compile=False
    )

    with open(MODEL_DIR / "tokenizer.pkl", "rb") as f:
        tokenizer = pickle.load(f)

    with open(MODEL_DIR / "config.json", "r") as f:
        config = json.load(f)

    index_word = {
        v: k
        for k, v in tokenizer.word_index.items()
    }

    return {
        "encoder": encoder,
        "decoder": decoder,
        "tokenizer": tokenizer,
        "index_word": index_word,
        "MAX_ENC_LEN": config["MAX_ENC_LEN"],
        "START_IDX": config["START_IDX"],
        "END_IDX": config["END_IDX"]
    }

model_data = load_chatbot()
# =====================================================
# STOCK DATABASE
# =====================================================

STOCK_SYMBOLS = {
    "reliance": "RELIANCE.NS",
    "tcs": "TCS.NS",
    "infosys": "INFY.NS",
    "hdfc": "HDFCBANK.NS",
    "icici": "ICICIBANK.NS",
    "sbi": "SBIN.NS",
    "wipro": "WIPRO.NS",
    "adani": "ADANIENT.NS",
    "zomato": "ZOMATO.NS",
    "tata motors": "TATAMOTORS.NS"
}

# =====================================================
# STOCK DATA
# =====================================================
import matplotlib.pyplot as plt
from datetime import datetime
import plotly.graph_objects as go
import yfinance as yf

def plot_stock_chart(symbol, start_dt=None, end_dt=None):

    stock = yf.Ticker(symbol)

    if start_dt and end_dt:
        df = stock.history(start=start_dt, end=end_dt)
    else:
        df = stock.history(period="6mo")

    if df.empty:
        st.error("No stock data found.")
        return

    # Fix MultiIndex columns if present
    if isinstance(df.columns, pd.MultiIndex):
        df.columns = df.columns.get_level_values(0)

    if "Close" not in df.columns:
        st.error("Close price data not available.")
        return

    # Index IS the date with stock.history()
    date_series  = pd.to_datetime(df.index)
    close_series = df["Close"].squeeze()

    current_time = datetime.now().strftime("%d-%m-%Y %H:%M:%S")
    st.info(f"Chart generated on: {current_time}")

    start_label  = date_series[0].strftime("%d %b %Y")
    end_label    = date_series[-1].strftime("%d %b %Y")
    start_price  = round(float(close_series.iloc[0]), 2)
    end_price    = round(float(close_series.iloc[-1]), 2)

    fig = go.Figure()

    fig.add_trace(
        go.Scatter(
            x=date_series,
            y=close_series,
            mode="lines",
            line=dict(color="#00C9FF", width=2),
            name="Close Price"
        )
    )

    # Start marker
    fig.add_trace(
        go.Scatter(
            x=[date_series[0]],
            y=[start_price],
            mode="markers+text",
            marker=dict(color="#00FF99", size=10),
            text=[f"Start<br>{start_label}<br>₹{start_price}"],
            textposition="top right",
            textfont=dict(size=11, color="#00FF99"),
            showlegend=False
        )
    )

    # End marker
    fig.add_trace(
        go.Scatter(
            x=[date_series[-1]],
            y=[end_price],
            mode="markers+text",
            marker=dict(color="#FF6B6B", size=10),
            text=[f"End<br>{end_label}<br>₹{end_price}"],
            textposition="top left",
            textfont=dict(size=11, color="#FF6B6B"),
            showlegend=False
        )
    )

    fig.update_layout(
        title=f"{symbol} Price History  |  {start_label}  →  {end_label}",
        xaxis_title="Date",
        yaxis_title="Price (₹)",
        height=500,
        template="plotly_dark",
        hovermode="x unified"
    )

    st.plotly_chart(fig, use_container_width=True)

    st.subheader("Recent Closing Prices")
    temp = df[["Open", "Close"]].tail(20).copy()
    temp.index = pd.to_datetime(temp.index).strftime("%d %b %Y")
    st.dataframe(temp, use_container_width=True)
    
    
    
    
def get_stock_info(symbol):

    try:

        stock = yf.Ticker(symbol)

        df = stock.history(period="1mo")

        if df.empty:
            return None

        return {
            "price": round(df["Close"].iloc[-1], 2),
            "high": round(df["High"].iloc[-1], 2),
            "low": round(df["Low"].iloc[-1], 2),
            "volume": int(df["Volume"].iloc[-1])
        }

    except:
        return None

# =====================================================
# CHATBOT RESPONSE
# =====================================================

def generate_response(user_text):

    tokenizer = model_data["tokenizer"]
    encoder = model_data["encoder"]
    decoder = model_data["decoder"]

    seq = tokenizer.texts_to_sequences(
        [user_text.lower()]
    )

    seq = pad_sequences(
        seq,
        maxlen=model_data["MAX_ENC_LEN"],
        padding="post"
    )

    lstm_h, lstm_c, gru_h = encoder.predict(
        seq,
        verbose=0
    )

    target_seq = np.array(
        [[model_data["START_IDX"]]]
    )

    reply = []

    for _ in range(25):

        output, lstm_h, lstm_c, gru_h = decoder.predict(
            [
                target_seq,
                lstm_h,
                lstm_c,
                gru_h
            ],
            verbose=0
        )

        idx = np.argmax(
            output[0, 0, :]
        )

        if idx == model_data["END_IDX"]:
            break

        word = model_data["index_word"].get(
            idx,
            ""
        )

        if word:
            reply.append(word)

        target_seq = np.array([[idx]])

    result = " ".join(reply)

    if len(result.strip()) == 0:
        return "Sorry, I don't know that yet."

    return result

# =====================================================
# WEATHER
# =====================================================

def get_weather(city):
    """Fetch weather using Open-Meteo (free, no API key needed)."""
    try:
        city_clean = city.strip()

        # Step 1: Geocode city name → lat/lon
        geo_url = (
            f"https://geocoding-api.open-meteo.com/v1/search"
            f"?name={requests.utils.quote(city_clean)}&count=1&language=en&format=json"
        )
        geo_res = requests.get(geo_url, timeout=10)

        if geo_res.status_code != 200 or not geo_res.json().get("results"):
            return f"⚠️ Could not find city **{city_clean}**. Try a different spelling."

        geo      = geo_res.json()["results"][0]
        lat      = geo["latitude"]
        lon      = geo["longitude"]
        area     = geo.get("name", city_clean)
        country  = geo.get("country", "")

        # Step 2: Fetch weather from Open-Meteo
        weather_url = (
            f"https://api.open-meteo.com/v1/forecast"
            f"?latitude={lat}&longitude={lon}"
            f"&current=temperature_2m,relative_humidity_2m,apparent_temperature,"
            f"weather_code,wind_speed_10m,wind_direction_10m,surface_pressure,visibility"
            f"&daily=weather_code,temperature_2m_max,temperature_2m_min"
            f"&timezone=auto&forecast_days=3"
        )
        w_res = requests.get(weather_url, timeout=10)

        if w_res.status_code != 200:
            return f"⚠️ Weather data unavailable for **{area}**."

        w = w_res.json()
        curr    = w["current"]
        daily   = w["daily"]

        # WMO weather code → description
        WMO = {
            0:"Clear sky", 1:"Mainly clear", 2:"Partly cloudy", 3:"Overcast",
            45:"Foggy", 48:"Icy fog", 51:"Light drizzle", 53:"Drizzle",
            55:"Heavy drizzle", 61:"Slight rain", 63:"Rain", 65:"Heavy rain",
            71:"Slight snow", 73:"Snow", 75:"Heavy snow", 80:"Rain showers",
            81:"Heavy showers", 82:"Violent showers", 95:"Thunderstorm",
            96:"Thunderstorm w/ hail", 99:"Heavy thunderstorm"
        }
        code        = curr.get("weather_code", 0)
        description = WMO.get(code, f"Code {code}")
        wind_deg    = curr.get("wind_direction_10m", 0)
        directions  = ["N","NE","E","SE","S","SW","W","NW"]
        wind_dir    = directions[round(wind_deg / 45) % 8]

        # 3-day forecast
        forecast_lines = []
        for i in range(len(daily["time"])):
            d_code = daily["weather_code"][i]
            d_desc = WMO.get(d_code, f"Code {d_code}")
            forecast_lines.append(
                f"  📅 {daily['time'][i]} — {d_desc}, "
                f"Max: {daily['temperature_2m_max'][i]}°C / Min: {daily['temperature_2m_min'][i]}°C"
            )
        forecast_str = "\n".join(forecast_lines)

        return f"""
🌤️ **Weather Report — {area}, {country}**

🌡️ Temperature  : {curr['temperature_2m']}°C  (Feels like {curr['apparent_temperature']}°C)
☁️  Condition    : {description}
💧 Humidity     : {curr['relative_humidity_2m']}%
💨 Wind         : {curr['wind_speed_10m']} km/h {wind_dir}
🔵 Pressure     : {curr['surface_pressure']} hPa

📆 **3-Day Forecast:**
{forecast_str}
"""

    except requests.exceptions.ConnectionError:
        return "⚠️ No internet connection. Could not fetch weather data."
    except requests.exceptions.Timeout:
        return "⚠️ Weather request timed out. Please try again."
    except (KeyError, ValueError, IndexError) as e:
        return f"⚠️ Could not parse weather data for **{city}**. Try e.g. `weather in Mumbai`."
    except Exception as e:
        return f"⚠️ Unexpected error: {str(e)}"


# =====================================================
# ROUTER
# =====================================================

from datetime import datetime
from datetime import datetime

def route_message(msg):

    text = msg.lower()

    current_datetime = datetime.now().strftime(
        "%d-%m-%Y %I:%M:%S %p"
    )

    if text == "time":
        return f"Current date & time: {current_datetime}"

    # Show all available stocks
    if any(kw in text for kw in ["stocks", "stock list", "show stocks", "all stocks", "available stocks"]):
        lines = []
        for i, (company, symbol) in enumerate(STOCK_SYMBOLS.items(), 1):
            lines.append(f"{i}. **{company.title()}** — `{symbol}`")
        return "📊 **Available Stocks:**\n\n" + "\n".join(lines)

    # Weather trigger: "weather in mumbai" / "weather delhi" / "mumbai weather"
    if "weather" in text:
        import re
        # Remove the word "weather", "in", "of", "the" carefully
        city = re.sub(r'\bweather\b', '', text)
        city = re.sub(r'^\s*(in|of|the)\s+', '', city.strip())
        city = city.strip()
        if city:
            return get_weather(city)
        else:
            return "🌍 Please tell me a city name, e.g. **weather in Mumbai**"

    for company, symbol in STOCK_SYMBOLS.items():

        if company.lower() in text:

            # Show graph — ask for date range
            if "graph" in text or "chart" in text:

                st.session_state.pending_chart = {
                    "symbol": symbol,
                    "company": company
                }

                return f"📅 Please select the **start and end date** for the **{company.upper()}** chart using the date pickers below."

            stock = get_stock_info(symbol)

            if stock:

                return f"""
📈 {company.upper()}

📅 Today : {current_datetime}

Current Price : ₹{stock['price']}

Day High : ₹{stock['high']}

Day Low : ₹{stock['low']}

Volume : {stock['volume']:,}
"""

            else:
                return "Unable to fetch stock data."

    return generate_response(text)


# =====================================================
# UI
# =====================================================

st.title("🤖 Stock Market AI Chatbot")

st.caption("chatbot BiLSTM + BiGRU ")

# ── Available Stocks Panel ─────────────────────────────────────────────
with st.expander("📊 Available Stocks — click to see all", expanded=False):

    cols = st.columns(2)

    for i, (company, symbol) in enumerate(STOCK_SYMBOLS.items()):
        with cols[i % 2]:
            st.markdown(
                f"🔹 **{company.title()}** &nbsp; `{symbol}`"
            )

    st.caption("💡 Try: *hdfc graph*, *reliance price*, *tcs chart*")

if "messages" not in st.session_state:

    st.session_state.messages = [
        {
            "role": "assistant",
            "content": "Hello! Ask me about stocks or general questions."
        }
    ]

# Show messages
# Show messages - render charts inline inside history

for msg in st.session_state.messages:

    with st.chat_message(msg["role"]):

        st.markdown(msg["content"])

        # Re-render chart inline if this message had one
        if msg.get("chart"):
            c = msg["chart"]
            plot_stock_chart(c["symbol"], c.get("start"), c.get("end"))

# Chat input

prompt = st.chat_input("Type your question...")

if prompt:

    st.session_state.messages.append({"role": "user", "content": prompt})

    with st.chat_message("user"):
        st.markdown(prompt)

    answer = route_message(prompt)

    with st.chat_message("assistant"):
        st.markdown(answer)

    st.session_state.messages.append({"role": "assistant", "content": answer})


# =====================================================
# DATE PICKER + CHART RENDER
# =====================================================

if "pending_chart" in st.session_state:

    info    = st.session_state.pending_chart
    company = info["company"]
    symbol  = info["symbol"]

    st.markdown("---")
    st.markdown(f"### 📅 Select Date Range for **{company.upper()}** Chart")

    col1, col2 = st.columns(2)

    with col1:
        start_date = st.date_input(
            "Start Date",
            value=pd.Timestamp.today() - pd.DateOffset(months=6),
            max_value=pd.Timestamp.today() - pd.DateOffset(days=1),
            key="chart_start_date"
        )

    with col2:
        end_date = st.date_input(
            "End Date",
            value=pd.Timestamp.today(),
            max_value=pd.Timestamp.today(),
            key="chart_end_date"
        )

    if st.button("📈 Show Chart", key="show_chart_btn"):

        if start_date >= end_date:
            st.error("⚠️ Start date must be before end date.")
        else:
            start_str = start_date.strftime("%Y-%m-%d")
            end_str   = end_date.strftime("%Y-%m-%d")

            # Attach chart to last assistant message so it renders inline in history
            for msg in reversed(st.session_state.messages):
                if msg["role"] == "assistant":
                    msg["chart"] = {
                        "symbol": symbol,
                        "start":  start_str,
                        "end":    end_str
                    }
                    break

            del st.session_state.pending_chart
            st.rerun()


# Clear chat

if st.button("Clear Chat"):

    st.session_state.messages = [
        {"role": "assistant", "content": "Chat cleared."}
    ]

    for key in ["pending_chart", "chart_symbol", "chart_start", "chart_end"]:
        if key in st.session_state:
            del st.session_state[key]

    st.rerun()