# app.py - Gradio app for HuggingFace Spaces
# Place this file in the same folder with your helper modules and saved model file.

import importlib
import os
import joblib
from datetime import date
import traceback

import gradio as gr
import matplotlib.pyplot as plt

# --- Load / reload helper modules (these should be in the same folder) ---
# If a module is missing, we will handle it later to avoid crashing the Space.
try:
    import data_utils
    import ticker_utils
    import ml_utils_v2
    import eda_utils
    import news_utils
    importlib.reload(data_utils)
    importlib.reload(ticker_utils)
    importlib.reload(ml_utils_v2)
    importlib.reload(eda_utils)
    importlib.reload(news_utils)
except Exception as e:
    # We'll still run Gradio but features depending on these modules will show friendly messages.
    missing_modules_err = str(e)
    data_utils = ticker_utils = ml_utils_v2 = eda_utils = news_utils = None
    print("Warning: one or more helper modules failed to import:", missing_modules_err)

# --- Local imports from helpers (safe guarded) ---
try:
    from data_utils import fetch_historical_data, fetch_live_quote, create_pdf_report
except Exception:
    fetch_historical_data = fetch_live_quote = create_pdf_report = None

try:
    from ml_utils_v2 import add_features
except Exception:
    add_features = None

try:
    from eda_utils import compute_daily_returns, compute_drawdown, compute_volatility, add_sma
except Exception:
    compute_daily_returns = compute_drawdown = compute_volatility = add_sma = None

try:
    from news_utils import fetch_news as fetch_news_live
except Exception:
    fetch_news_live = None

# --- Saved model file (change name if you used a different filename) ---
MODEL_FILE = "best_model_h7_t3.joblib"

saved_model = None
saved_features = None
saved_label_map = None
saved_accuracy = "N/A"
if os.path.exists(MODEL_FILE):
    try:
        saved = joblib.load(MODEL_FILE)
        # saved is expected to be a dict with keys 'model','features','label_map' and optionally 'accuracy'
        saved_model = saved.get("model")
        saved_features = saved.get("features")
        saved_label_map = saved.get("label_map")
        saved_accuracy = saved.get("accuracy", "N/A")
        print(f"Loaded model from {MODEL_FILE} (accuracy={saved_accuracy})")
    except Exception as e:
        print("Could not load saved model:", e)
        saved_model = saved_features = saved_label_map = None
else:
    print(f"No saved model found at {MODEL_FILE}. Predictions will be unavailable until you upload a model file.")

# --- Helpers ---


def render_signal_badge(signal: str) -> str:
    """Return a small HTML badge for signal highlighting."""
    color = {"BUY": "#1aab2b", "HOLD": "#ffb703", "SELL": "#e63946"}.get(signal, "#999999")
    return f'<div style="display:inline-block;padding:10px 16px;border-radius:10px;background:{color};color:#fff;font-weight:700;font-size:18px">{signal}</div>'


def predict_with_saved_model(model, features_list, label_map, df_with_features):
    """
    Predict latest signal using a pre-saved model.
    - model: trained model (trained on mapped labels)
    - features_list: list of features that model expects
    - label_map: mapping original_label -> mapped_label (e.g., {-1:0,0:1,1:2})
    - df_with_features: DataFrame after running add_features()
    Returns string "BUY"/"HOLD"/"SELL" or error code.
    """
    if model is None or features_list is None or label_map is None:
        return "MODEL_NOT_LOADED"

    # keep only features present in the DataFrame
    feat = [f for f in features_list if f in df_with_features.columns]
    if not feat:
        return "NO_FEATURES"

    last = df_with_features.iloc[-1:][feat].replace([float('inf'), float('-inf')], float('nan'))
    # try small fills (safe) to avoid NaNs preventing prediction
    last = last.fillna(method="ffill").fillna(method="bfill")
    if last.isna().all(axis=None):
        return "INSUFFICIENT_DATA"

    # model expects mapped labels, decode prediction back to original
    inv_map = {v: k for k, v in label_map.items()}
    try:
        mapped_val = model.predict(last)[0]
        decoded = inv_map.get(int(mapped_val), None)
        if decoded is None:
            return "DECODE_FAIL"
        return {1: "BUY", 0: "HOLD", -1: "SELL"}[decoded]
    except Exception as e:
        print("Prediction error:", e)
        return "PRED_ERROR"


# --- Core backend function used by Gradio ---
def analyze_and_report(company_name: str, exchange: str, start_date: str, end_date: str):
    """
    Main function called by Gradio.
    Returns: info_html, price_fig, hist_fig, dd_fig, vol_fig, news_md, pdf_path
    """
    # basic validations and module availability checks
    if not company_name or not company_name.strip():
        return "❌ Enter a company name", None, None, None, None, "No news", None

    if ticker_utils is None:
        return "❌ ticker_utils not available in Space. Upload it.", None, None, None, None, "No news", None
    if fetch_historical_data is None:
        return "❌ data_utils.fetch_historical_data not available. Upload data_utils.py", None, None, None, None, "No news", None

    # resolve ticker (safe)
    try:
        resolved = ticker_utils.find_ticker(company_name.lower(), exchange_preference=exchange)
    except Exception as e:
        return f"❌ Could not resolve ticker ({e})", None, None, None, None, "No news", None

    # fetch historical data
    try:
        df = fetch_historical_data(resolved, str(start_date), str(end_date))
    except Exception as e:
        return f"❌ Error fetching historical data: {e}", None, None, None, None, "No news", None

    if df is None or df.empty:
        return f"❌ No historical data found for {resolved}", None, None, None, None, "No news", None

    # attempt to append latest live day (data_utils may provide a helper)
    try:
        if hasattr(data_utils, "ensure_latest"):
            df, _ = data_utils.ensure_latest(df, resolved)
    except Exception:
        pass

    # get live snapshot (best-effort)
    try:
        live = fetch_live_quote(resolved) if fetch_live_quote else {}
    except Exception:
        live = {}

    # ----- EDA figures -----
    price_fig = hist_fig = dd_fig = vol_fig = None
    try:
        if add_sma is not None:
            df_plot = add_sma(df.copy(), windows=[20, 50])
        else:
            df_plot = df.copy()
        price_fig, ax = plt.subplots(figsize=(8, 3))
        ax.plot(df_plot.index, df_plot["Close"], label="Close")
        if "SMA_20" in df_plot.columns:
            ax.plot(df_plot.index, df_plot["SMA_20"], label="SMA20")
        if "SMA_50" in df_plot.columns:
            ax.plot(df_plot.index, df_plot["SMA_50"], label="SMA50")
        ax.set_title(f"{resolved} Close & SMAs")
        ax.legend()
        plt.tight_layout()
    except Exception as e:
        print("Price plot error:", e)
        price_fig = None

    try:
        if compute_daily_returns:
            fig_hist, ax2 = plt.subplots(figsize=(6, 3))
            compute_daily_returns(df).hist(bins=40, ax=ax2)
            ax2.set_title("Histogram of daily returns")
            plt.tight_layout()
            hist_fig = fig_hist
    except Exception as e:
        print("Histogram error:", e)
        hist_fig = None

    try:
        if compute_drawdown:
            fig_dd, ax3 = plt.subplots(figsize=(6, 3))
            compute_drawdown(df).plot(ax=ax3)
            ax3.set_title("Drawdown")
            plt.tight_layout()
            dd_fig = fig_dd
    except Exception as e:
        print("Drawdown error:", e)
        dd_fig = None

    try:
        if compute_volatility:
            fig_vol, ax4 = plt.subplots(figsize=(6, 3))
            compute_volatility(df, 30).plot(ax=ax4)
            ax4.set_title("30-day rolling volatility")
            plt.tight_layout()
            vol_fig = fig_vol
    except Exception as e:
        print("Volatility plot error:", e)
        vol_fig = None

    # ----- Prediction using saved model -----
    pred_html = "<i>Prediction unavailable</i>"
    pred_text = "N/A"
    if saved_model is not None and add_features is not None:
        try:
            df_feat = add_features(df)
            signal = predict_with_saved_model(saved_model, saved_features, saved_label_map, df_feat)
            if signal in ("MODEL_NOT_LOADED", "NO_FEATURES", "INSUFFICIENT_DATA", "DECODE_FAIL", "PRED_ERROR"):
                pred_html = f"<i>Prediction unavailable: {signal}</i>"
                pred_text = signal
            else:
                pred_html = render_signal_badge(signal)
                pred_text = signal
        except Exception as e:
            print("Prediction pipeline error:", e)
            pred_html = "<i>Prediction error</i>"
            pred_text = "ERR"
    else:
        pred_html = "<i>Model not loaded or features missing</i>"

    # ----- News -----
    try:
        if fetch_news_live is not None:
            news_items = fetch_news_live(resolved.split(".")[0].lower(), ticker=resolved)
        else:
            news_items = []
    except Exception as e:
        print("News fetch error:", e)
        news_items = []

    news_md = ""
    if news_items:
        for n in news_items:
            title = n.get("title", "No title")
            src = n.get("source", "")
            summary = n.get("summary", "")
            url = n.get("url", "")
            if url:
                news_md += f"### [{title}]({url})\n**{src}**\n{summary}\n\n---\n"
            else:
                news_md += f"### {title}\n**{src}**\n{summary}\n\n---\n"
    else:
        news_md = "No news available."

    # ----- PDF report (try to create; return path if successful) -----
    pdf_path = None
    try:
        if create_pdf_report is not None:
            pdf_path = f"report_{resolved.replace('.','_')}.pdf"
            create_pdf_report(pdf_path, resolved, df, live_info=live)
        else:
            pdf_path = None
    except Exception as e:
        print("PDF generation error:", e)
        pdf_path = None

    # ----- Info HTML block (includes model accuracy if available) -----
    model_acc_text = saved_accuracy if saved_accuracy != "N/A" else "N/A"
    info_html = f"""
    ### Resolved ticker: `{resolved}`
    **Live LTP:** {live.get('LTP','N/A')} &nbsp;&nbsp; **High:** {live.get('DayHigh','N/A')} &nbsp;&nbsp; **Low:** {live.get('DayLow','N/A')}
    <br><br>
    **Model accuracy (saved):** {model_acc_text}
    <br><br>
    {pred_html}
    """

    return info_html, price_fig, hist_fig, dd_fig, vol_fig, news_md, pdf_path


# --- Build the Gradio UI ---
with gr.Blocks(title="Indian Stock Analyzer") as demo:
    gr.Markdown("# 📈 Indian Stock Analyzer")

    with gr.Row():
        with gr.Column(scale=3):
            company = gr.Textbox(label="Company name (e.g., reliance, hudco, rvnl)", placeholder="Type company...")
            exchange = gr.Dropdown(["NSE", "BSE"], value="NSE", label="Exchange")
            start_dt = gr.Textbox(label="Start Date (YYYY-MM-DD)", value="2015-01-01")
            end_dt = gr.Textbox(label="End Date (YYYY-MM-DD)", value=str(date.today()))
            submit = gr.Button("Analyze")

        with gr.Column(scale=2):
            info = gr.HTML("<i>Results will appear here...</i>")
            download_pdf = gr.File(label="Download PDF (generated after Analyze)")

    with gr.Tabs():
        with gr.TabItem("Price & EDA"):
            price_out = gr.Plot()
            hist_out = gr.Plot()
            dd_out = gr.Plot()
            vol_out = gr.Plot()
        with gr.TabItem("News"):
            news_out = gr.Markdown()

    # Wire up the button
    submit.click(
        analyze_and_report,
        inputs=[company, exchange, start_dt, end_dt],
        outputs=[info, price_out, hist_out, dd_out, vol_out, news_out, download_pdf],
    )

# Launch the app (HuggingFace will run this automatically)
if __name__ == "__main__":
    # Use ssr_mode=False and a common port to avoid long SSR startup on Spaces
    demo.launch(server_name="0.0.0.0", server_port=7860, ssr_mode=False)