from datetime import datetime, timedelta
import pandas as pd
import plotly.graph_objs as go
import streamlit as st
import pytz

from datasource import get_stock_bars


def get_ny_today():
    ny_tz = pytz.timezone("America/New_York")
    return datetime.now(ny_tz).date()


st.set_page_config(layout="wide")
st.title("Candlestick Chart")
st.sidebar.title("Filters")

symbol = st.sidebar.text_input("Ticker symbol", value="TSLA").upper()
date_start = st.sidebar.date_input(
    "Start date",
    get_ny_today() - timedelta(days=1),
    max_value=get_ny_today() - timedelta(days=1),
)
timeframe = st.sidebar.selectbox(
    "Timeframe", options=["1m", "5m", "15m", "30m", "1h", "1d"], index=2
)

try:
    if timeframe == "1d":
        actual_start_date = date_start - timedelta(days=29)  # 29 + 1 = 30 days total
        actual_end_date = date_start + timedelta(days=1)
    else:
        actual_start_date = date_start
        actual_end_date = date_start + timedelta(days=1)

    bars = get_stock_bars(
        symbol, actual_start_date, actual_end_date, interval=timeframe
    )
    if bars.empty:
        st.warning("No data. Check symbol and dates.")
    else:
        bars = bars.reset_index()
        bars["timestamp"] = bars["timestamp"].dt.tz_convert("America/New_York")
        bars = bars.set_index("timestamp")

        bars["volume"] = pd.to_numeric(bars["volume"], errors="coerce").fillna(0)
        typical_price = (bars["high"] + bars["low"] + bars["close"]) / 3.0
        bars["vwap"] = (typical_price * bars["volume"]).cumsum() / bars[
            "volume"
        ].cumsum()

        if timeframe != "1d":
            premarket_mask = bars.index.time < pd.to_datetime("09:30:00").time()
            premarket_high = (
                bars.loc[premarket_mask, "high"].max() if premarket_mask.any() else None
            )
        else:
            premarket_high = None

        timestamps = [ts.strftime("%Y-%m-%d %H:%M:%S") for ts in bars.index]
        open_vals = bars["open"].tolist()
        high_vals = bars["high"].tolist()
        low_vals = bars["low"].tolist()
        close_vals = bars["close"].tolist()
        vwap_vals = bars["vwap"].tolist()
        volume_vals = bars["volume"].tolist()

        fig = go.Figure()
        fig.add_trace(
            go.Candlestick(
                x=timestamps,
                open=open_vals,
                high=high_vals,
                low=low_vals,
                close=close_vals,
                name="Candlestick",
            )
        )
        fig.add_trace(
            go.Scatter(
                x=timestamps,
                y=vwap_vals,
                mode="lines",
                line=dict(color="yellow", width=1),
                name="VWAP",
            )
        )
        if premarket_high is not None and pd.notna(premarket_high):
            fig.add_trace(
                go.Scatter(
                    x=[timestamps, timestamps[-1]],
                    y=[premarket_high, premarket_high],
                    mode="lines",
                    line=dict(color="red", width=1, dash="dash"),
                    name="Premarket High",
                )
            )
        fig.add_trace(
            go.Bar(
                x=timestamps,
                y=volume_vals,
                yaxis="y2",
                marker=dict(color="rgba(200,200,200,0.5)"),
                name="Volume",
                opacity=0.5,
            )
        )

        if timeframe != "1d":
            bars_dates = pd.to_datetime(bars.index.date).unique()
            for day in bars_dates:
                dm = pd.Timestamp(day).strftime("%Y-%m-%d")
                fig.add_vrect(
                    x0=f"{dm} 04:00:00",
                    x1=f"{dm} 09:30:00",
                    fillcolor="rgba(0, 200, 255, 0.10)",
                    layer="below",
                    line_width=0,
                    annotation_text="Pre-market",
                    annotation_position="top left",
                )
                fig.add_vrect(
                    x0=f"{dm} 16:00:00",
                    x1=f"{dm} 20:00:00",
                    fillcolor="rgba(255, 200, 0, 0.08)",
                    layer="below",
                    line_width=0,
                    annotation_text="After-hours",
                    annotation_position="top left",
                )

        if timeframe == "1d":
            title = f"{symbol} - {timeframe} ({actual_start_date.strftime('%Y-%m-%d')} to {date_start.strftime('%Y-%m-%d')})"
        else:
            title = f"{symbol} - {timeframe}"

        fig.update_layout(
            title=title,
            xaxis_title="Date/Time",
            yaxis_title="Price",
            xaxis_rangeslider_visible=False,
            yaxis=dict(domain=[0.3, 1]),
            yaxis2=dict(domain=[0, 0.25], title="Volume"),
            legend=dict(orientation="h"),
            margin=dict(t=40, b=20),
            hovermode="x unified",
            height=720,
        )

        st.plotly_chart(fig, use_container_width=True)
        preview_cols = ["open", "high", "low", "close", "volume", "vwap"]
        if premarket_high is not None and pd.notna(premarket_high):
            preview_cols.append("premarket_high")
            bars["premarket_high"] = premarket_high
        st.write("Data:")
        st.dataframe(bars, use_container_width=True)

except Exception as e:
    st.error(f"Error fetching or plotting data: {e}")
    import traceback

    st.write("Full traceback:")
    st.code(traceback.format_exc())