import streamlit as st
import pandas as pd
import numpy as np
import yfinance as yf
from datetime import datetime, timedelta, time as dtime
import pytz
import time

# =========================
# CONFIG
# =========================
SP500 = "SPY"
GOLD = "GLD"

SMA_MONTHS = 36
BASE_EQUITY = 0.60
TILT_MULT = 1.5
MAX_TILT = 0.25
MIN_EQUITY = 0.20
MAX_EQUITY = 0.90

PRICE_REFRESH_SEC = 5

# =========================
# DATA
# =========================
@st.cache_data(ttl=3600)
def load_history(ticker):
    df = yf.download(
        ticker,
        period="10y",
        auto_adjust=True,
        progress=False,
    )

    if df.empty:
        raise RuntimeError(f"No historical data for {ticker}")

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

    if "Close" not in df.columns:
        raise RuntimeError(f"'Close' not found for {ticker}")

    series = df["Close"].copy()

    if not isinstance(series.index, pd.DatetimeIndex):
        raise RuntimeError(f"Index is not DatetimeIndex for {ticker}")

    series = series.resample("M").last()

    return series


@st.cache_data(ttl=15)
def live_price(ticker):
    df = yf.download(
        ticker,
        period="1d",
        interval="1m",
        progress=False,
    )

    if df.empty:
        raise RuntimeError(f"No live data for {ticker}")

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

    price = df["Close"].iloc[-1]

    return float(price)


def us_market_status():
    tz = pytz.timezone("America/New_York")
    now = datetime.now(tz)

    market_open = dtime(9, 30)
    market_close = dtime(16, 0)

    is_weekday = now.weekday() < 5
    is_open = is_weekday and market_open <= now.time() <= market_close

    return {
        "is_open": is_open,
        "now_et": now.strftime("%Y-%m-%d %H:%M:%S"),
        "session": "OPEN" if is_open else "CLOSED",
    }


sp_hist = load_history(SP500)
gd_hist = load_history(GOLD)

hist_prices = pd.concat(
    [sp_hist, gd_hist],
    axis=1,
    keys=["SP500", "GOLD"],
).dropna()

sma_sp = hist_prices["SP500"].rolling(SMA_MONTHS).mean()
sma_gd = hist_prices["GOLD"].rolling(SMA_MONTHS).mean()

# =========================
# STATE INIT
# =========================
if "running" not in st.session_state:
    st.session_state.running = False
    st.session_state.last_sip = None
    st.session_state.total_invested = 0.0
    st.session_state.sp_units = 0.0
    st.session_state.gd_units = 0.0
    st.session_state.last_sp_price = None
    st.session_state.last_gd_price = None
if "pnl_ledger" not in st.session_state:
    st.session_state.pnl_ledger = pd.DataFrame(
        columns=["timestamp", "pnl"]
    )

# =========================
# UI
# =========================
st.title("SIPXAR - Live SIP Portfolio (US)")
market = us_market_status()

if market["is_open"]:
    st.success(f"US Market OPEN (ET {market['now_et']})")
else:
    st.warning(f"US Market CLOSED (ET {market['now_et']})")

sip_amount = st.number_input("SIP Amount ($)", min_value=10, value=1000, step=100)
sip_minutes = st.number_input(
    "SIP Frequency (minutes) - TEST MODE",
    min_value=1,
    value=2,
    step=1,
)
sip_interval = timedelta(minutes=sip_minutes)

col1, col2 = st.columns(2)
if col1.button("START"):
    st.session_state.running = True
    if st.session_state.last_sip is None:
        st.session_state.last_sip = datetime.utcnow() - sip_interval

if col2.button("STOP"):
    st.session_state.running = False

# =========================
# ENGINE LOOP
# =========================
if st.session_state.running:
    now = datetime.utcnow()

    if market["is_open"]:
        sp_price = live_price(SP500)
        gd_price = live_price(GOLD)
        st.session_state.last_sp_price = sp_price
        st.session_state.last_gd_price = gd_price
    else:
        sp_price = st.session_state.get("last_sp_price")
        gd_price = st.session_state.get("last_gd_price")

        if sp_price is None or pd.isna(sp_price):
            sp_price = hist_prices["SP500"].iloc[-1]
        if gd_price is None or pd.isna(gd_price):
            gd_price = hist_prices["GOLD"].iloc[-1]

    if st.session_state.last_sip is None:
        st.session_state.last_sip = now - sip_interval

    # SIP trigger only when market is open
    if (
        market["is_open"]
        and (now - st.session_state.last_sip) >= sip_interval
    ):
        rd = (
            (hist_prices["SP500"].iloc[-1] / sma_sp.iloc[-1])
            - (hist_prices["GOLD"].iloc[-1] / sma_gd.iloc[-1])
        )

        tilt = np.clip(-rd * TILT_MULT, -MAX_TILT, MAX_TILT)
        eq_w = np.clip(BASE_EQUITY * (1 + tilt), MIN_EQUITY, MAX_EQUITY)
        gd_w = 1 - eq_w

        sp_buy = sip_amount * eq_w
        gd_buy = sip_amount * gd_w

        st.session_state.sp_units += sp_buy / sp_price
        st.session_state.gd_units += gd_buy / gd_price
        st.session_state.total_invested += sip_amount
        st.session_state.last_sip = now

    # MTM
    sp_val = st.session_state.sp_units * sp_price
    gd_val = st.session_state.gd_units * gd_price
    port_val = sp_val + gd_val
    pnl = port_val - st.session_state.total_invested
    st.session_state.pnl_ledger = pd.concat(
        [
            st.session_state.pnl_ledger,
            pd.DataFrame(
                {
                    "timestamp": [datetime.utcnow()],
                    "pnl": [pnl],
                }
            ),
        ],
        ignore_index=True,
    )

    # =========================
    # DISPLAY
    # =========================
    st.subheader("Portfolio Snapshot")
    st.caption(
        "Prices updating live" if market["is_open"] else "Prices frozen - market closed"
    )

    c1, c2, c3 = st.columns(3)
    c1.metric("Total Invested", f"${st.session_state.total_invested:,.2f}")
    c2.metric("Portfolio Value", f"${port_val:,.2f}")
    c3.metric("Unrealized PnL", f"${pnl:,.2f}", delta=f"{pnl:,.2f}")

    next_sip_in = sip_interval - (now - st.session_state.last_sip)
    next_sip_sec = max(0, int(next_sip_in.total_seconds()))
    st.caption(f"Next SIP in ~ {next_sip_sec} seconds (TEST MODE)")

    st.subheader("Equity Curve (PnL)")
    if len(st.session_state.pnl_ledger) > 1:
        pnl_df = st.session_state.pnl_ledger.copy()
        pnl_df["timestamp"] = pd.to_datetime(pnl_df["timestamp"])
        pnl_df.set_index("timestamp", inplace=True)

        st.line_chart(
            pnl_df["pnl"],
            height=350,
        )
    else:
        st.info("Equity curve will appear after portfolio updates.")

    st.subheader("Holdings")
    st.dataframe(
        pd.DataFrame(
            {
                "Asset": ["SP500 (SPY)", "Gold (GLD)"],
                "Units": [st.session_state.sp_units, st.session_state.gd_units],
                "Price": [sp_price, gd_price],
                "Value": [sp_val, gd_val],
            }
        )
    )

    time.sleep(PRICE_REFRESH_SEC)
    st.rerun()
