import streamlit as st
import pandas as pd
import altair as alt
import os
import sys
from pathlib import Path
from dotenv import set_key

# --- Path Setup & Imports ---
# Ensure we can import backend.py regardless of how the script is run
current_dir = Path(__file__).parent
if str(current_dir) not in sys.path:
    sys.path.append(str(current_dir))

from backend import AuthManager, DataManager, AnalyticsEngine  # noqa: E402

# --- Page Config ---
st.set_page_config(page_title="Kotak Neo Portfolio", layout="wide", page_icon="📈")

# --- Constants ---
ENV_PATH = Path(".env")


def save_credentials(key, mobile, ucc, secret, mpin):
    try:
        if not ENV_PATH.exists():
            ENV_PATH.touch()
        set_key(ENV_PATH, "KOTAK_CONSUMER_KEY", key)
        set_key(ENV_PATH, "KOTAK_MOBILE", mobile)
        set_key(ENV_PATH, "KOTAK_UCC", ucc)
        set_key(ENV_PATH, "KOTAK_TOTP_SECRET", secret)
        set_key(ENV_PATH, "KOTAK_MPIN", mpin)
        return True
    except Exception as e:
        st.error(f"Failed to save config: {e}")
        return False


def render_login_sidebar():
    st.sidebar.header("🔐 Setup / Login")
    has_env = ENV_PATH.exists() and os.getenv("KOTAK_CONSUMER_KEY")

    if has_env:
        st.sidebar.success("Configuration found.")
        if st.sidebar.button("Login with Saved Config"):
            return "ENV_LOGIN"
        with st.sidebar.expander("Update Credentials"):
            return render_login_form(save_mode=True)
    else:
        st.sidebar.warning("No configuration found.")
        return render_login_form(save_mode=True)


def render_login_form(save_mode=False):
    with st.sidebar.form("creds_form"):
        key = st.text_input("Consumer Key", type="password")
        mob = st.text_input("Mobile", max_chars=10)
        ucc = st.text_input("UCC")
        secret = st.text_input("TOTP Secret", type="password")
        mpin = st.text_input("MPIN", type="password")

        btn_text = "Save & Login" if save_mode else "Login Temporary"
        if st.form_submit_button(btn_text):
            if save_mode:
                save_credentials(key, mob, ucc, secret, mpin)
                st.success("Credentials Saved! Please click Login.")
                # Force reload env in current session
                os.environ["KOTAK_CONSUMER_KEY"] = key
                os.environ["KOTAK_MOBILE"] = mob
                os.environ["KOTAK_UCC"] = ucc
                os.environ["KOTAK_TOTP_SECRET"] = secret
                os.environ["KOTAK_MPIN"] = mpin
                return "ENV_LOGIN"
            return {"key": key, "mob": mob, "ucc": ucc, "secret": secret, "mpin": mpin}
    return None


def main():
    st.title("📈 Kotak Neo Portfolio Dashboard")

    if "auth_manager" not in st.session_state:
        st.session_state.auth_manager = AuthManager()
    if "client" not in st.session_state:
        st.session_state.client = None

    # --- Authentication Flow ---
    if not st.session_state.client:
        creds = render_login_sidebar()
        if creds:
            with st.spinner("Authenticating..."):
                try:
                    if creds == "ENV_LOGIN":
                        st.session_state.client = st.session_state.auth_manager.login()
                    elif isinstance(creds, dict):
                        st.session_state.client = st.session_state.auth_manager.login(
                            creds["key"],
                            creds["mob"],
                            creds["ucc"],
                            creds["secret"],
                            creds["mpin"],
                        )
                    st.success("Logged In Successfully!")
                    st.rerun()
                except Exception as e:
                    st.error(f"Login Error: {e}")
            return

    # --- Dashboard Flow ---
    if st.session_state.client:
        if st.sidebar.button("Logout"):
            st.session_state.client = None
            st.rerun()

        data_mgr = DataManager(st.session_state.client)

        with st.spinner("Fetching Live Data..."):
            try:
                holdings_json, positions_json = data_mgr.fetch_portfolio()

                # FIXED: Handle return values correctly
                hydrate_result = data_mgr.hydrate_live_prices(holdings_json)

                # Unpack safely regardless of tuple size
                if isinstance(hydrate_result, tuple):
                    holdings_json = hydrate_result[0]
                    updates = hydrate_result[1]
                else:
                    holdings_json = hydrate_result
                    updates = 0

                if updates > 0:
                    st.toast(f"Updated {updates} scrips with live data.", icon="⚡")
                else:
                    st.warning(
                        "Using Cached Data (No Live Updates - Check Market Hours)"
                    )

                df = AnalyticsEngine.process_holdings(holdings_json)
                render_dashboard(df, positions_json)

            except Exception as e:
                st.error(f"Data Error: {e}")
                if st.button("Retry"):
                    st.rerun()


def render_dashboard(df, positions_json):
    if df.empty:
        st.info("No holdings data available.")
        return

    # KPI Calculation
    invested = df["investedAmount"].sum()
    curr_val = df["currentValue"].sum()
    pnl = curr_val - invested
    pnl_pct = (pnl / invested * 100) if invested else 0
    day_pnl = df["dayPnL"].sum()

    # Top Row Metrics
    c1, c2, c3, c4 = st.columns(4)
    c1.metric("Total Invested", f"₹{invested:,.2f}")
    c2.metric("Current Value", f"₹{curr_val:,.2f}")
    c3.metric("Total P&L", f"₹{pnl:,.2f}", f"{pnl_pct:.2f}%")
    c4.metric("Day's P&L", f"₹{day_pnl:,.2f}", "Today")

    st.divider()

    # Tabs
    t1, t2, t3, t4 = st.tabs(
        ["💼 Holdings", "⚡ Positions", "📈 Analysis", "🔍 Raw Data"]
    )

    with t1:
        # --- NEW: Allocation Bar Chart ---
        if not df.empty:
            st.subheader("Allocation by Stock (Current Value)")
            if "tradingSymbol" in df.columns:
                # Matches the logic from the old script
                chart_data = df.set_index("tradingSymbol")[["currentValue"]]
                st.bar_chart(chart_data)
                st.divider()

        # Prepare Display DataFrame with Renamed Columns
        display_cols = [
            "tradingSymbol",
            "quantity",
            "averagePrice",
            "currentPrice",
            "investedAmount",
            "currentValue",
            "dayChangePct",
            "profitAndLoss",
            "pnl_pct",
        ]
        final_cols = [c for c in display_cols if c in df.columns]

        display_df = df[final_cols].rename(
            columns={
                "tradingSymbol": "Scrip",
                "averagePrice": "Avg Price",
                "currentPrice": "CMP",
                "investedAmount": "Invested",
                "currentValue": "Value",
                "dayChangePct": "Day %",
                "profitAndLoss": "P&L",
                "pnl_pct": "P&L %",
            }
        )

        # Styling Function (Green/Red text)
        def style_pnl(v):
            if isinstance(v, (int, float)):
                if v > 0:
                    return "color: #28a745; font-weight: bold"
                elif v < 0:
                    return "color: #dc3545; font-weight: bold"
            return ""

        # Render Styled Table using Pandas Styler (as requested)
        st.dataframe(
            display_df.style.format(
                {
                    "Avg Price": "₹{:.2f}",
                    "CMP": "₹{:.2f}",
                    "Invested": "₹{:.2f}",
                    "Value": "₹{:.2f}",
                    "Day %": "{:.2f}%",
                    "P&L": "₹{:.2f}",
                    "P&L %": "{:.2f}%",
                }
            ).applymap(style_pnl, subset=["P&L", "P&L %", "Day %"]),
            use_container_width=True,
            hide_index=True,
        )

    with t2:
        if positions_json and "data" in positions_json:
            st.dataframe(pd.DataFrame(positions_json["data"]), use_container_width=True)
        else:
            st.info("No active positions.")

    with t3:
        st.subheader("Sector Allocation")
        sector_df = AnalyticsEngine.get_sector_data(df)
        if not sector_df.empty:
            chart = (
                alt.Chart(sector_df)
                .mark_arc(innerRadius=50)
                .encode(
                    theta=alt.Theta(field="currentValue", type="quantitative"),
                    color=alt.Color(field="sector", type="nominal"),
                    tooltip=["sector", alt.Tooltip("currentValue", format=",.2f")],
                )
            )
            st.altair_chart(chart, use_container_width=True)
        else:
            st.info("Sector data not available.")

    with t4:
        st.dataframe(df)


if __name__ == "__main__":
    main()
