from __future__ import annotations

import io
import logging
from datetime import timedelta, timezone
from typing import Literal

import matplotlib

matplotlib.use("Agg")

import matplotlib.pyplot as plt
import mplfinance as mpf
import pandas as pd
from matplotlib.patches import Rectangle

logger = logging.getLogger(__name__)

ActionSide = Literal["long", "short"]
RIGHT_PAD_CANDLES = 15
UTC_PLUS_2 = timezone(timedelta(hours=2))

COLORS = {
    "bg": "#0f1115",
    "panel": "#0f1115",
    "grid": "#1e222d",
    "text": "#d1d4dc",
    "up": "#26a69a",
    "down": "#ef5350",
    "entry": "#42a5f5",
    "price": "#ffca28",
    "sl": "#ef5350",
    "tp1": "#66bb6a",
    "tp2": "#43a047",
    "tp3": "#2e7d32",
    "reward_fill": (0.15, 0.65, 0.45),
    "risk_fill": (0.85, 0.25, 0.25, 0.10),
    "label_bg": "#0f1115",
    "vol_up": "#26a69a",
    "vol_down": "#ef5350",
}

# Reward zone opacities: Entry→TP1, TP1→TP2, TP2→TP3 (decreasing)
REWARD_ALPHAS = (0.16, 0.10, 0.06)


def _parse_price(raw: str) -> float:
    return float(raw.strip().replace(" ", "").replace(",", "").replace("$", ""))


def _to_utc_plus_2(df: pd.DataFrame) -> pd.DataFrame:
    out = df.copy()
    idx = out.index
    if idx.tz is None:
        idx = idx.tz_localize("UTC")
    out.index = idx.tz_convert(UTC_PLUS_2)
    return out


def _pad_right(df: pd.DataFrame, candles: int = RIGHT_PAD_CANDLES) -> pd.DataFrame:
    """Append empty (NaN) candles so there is free space to the right of price action."""
    if len(df) < 2:
        delta = pd.Timedelta(minutes=15)
    else:
        delta = df.index[-1] - df.index[-2]
        if not isinstance(delta, pd.Timedelta) or delta <= pd.Timedelta(0):
            delta = pd.Timedelta(minutes=15)

    future_index = pd.date_range(
        start=df.index[-1] + delta,
        periods=candles,
        freq=delta,
        tz=df.index.tz,
    )
    pad = pd.DataFrame(index=future_index, columns=df.columns, dtype=float)
    return pd.concat([df, pad])


def _volume_overlay(
    df: pd.DataFrame,
    *,
    y_low: float,
    y_high: float,
    fraction: float = 0.18,
) -> tuple[pd.Series, list[str]]:
    """Scale volume into the bottom of the price pane (TradingView-style)."""
    span = y_high - y_low
    height = span * fraction
    base = y_low - span * 0.01
    vol = df["Volume"].astype(float)
    # Right-pad / empty candles: NaN so mplfinance skips the bar entirely
    # (zero height still draws a stub from y=0 → base, often as black).
    empty = df["Open"].isna() | df["Close"].isna() | vol.isna()
    vol_filled = vol.fillna(0.0)
    real = vol_filled[~empty]
    vmax = float(real.max()) if len(real) else 1.0
    if vmax <= 0:
        vmax = 1.0
    scaled = base + (vol_filled / vmax) * height
    scaled = scaled.mask(empty)

    colors: list[str] = []
    for _, row in df.iterrows():
        if pd.isna(row["Close"]) or pd.isna(row["Open"]):
            colors.append(COLORS["bg"])
        elif row["Close"] >= row["Open"]:
            colors.append(COLORS["vol_up"])
        else:
            colors.append(COLORS["vol_down"])
    return scaled, colors


def _position_start_x(df: pd.DataFrame, signal_time: int | None) -> int:
    """Integer x of the candle where the seq==1 position starts (fallback: last)."""
    last_x = len(df) - 1
    if signal_time is None or last_x < 0:
        return max(last_x, 0)
    ts = pd.Timestamp(int(signal_time), unit="s", tz="UTC")
    if df.index.tz is not None:
        ts = ts.tz_convert(df.index.tz)
    # Last candle whose open time is <= signal time
    pos = int(df.index.searchsorted(ts, side="right") - 1)
    if pos < 0:
        return 0
    return min(pos, last_x)


def _draw_previous_trade(
    ax,
    df: pd.DataFrame,
    *,
    prev_entry: float,
    prev_side: str,
    prev_signal_time: int,
    exit_price: float,
    exit_x: int,
) -> None:
    """Entry → exit of the trade that this reversal closes."""
    prev_x = _position_start_x(df, prev_signal_time)
    if prev_x >= exit_x:
        return
    profitable = (
        exit_price >= prev_entry if prev_side == "long" else exit_price <= prev_entry
    )
    color = COLORS["up"] if profitable else COLORS["down"]
    y_low = min(prev_entry, exit_price)
    y_high = max(prev_entry, exit_price)
    height = y_high - y_low
    if height <= 0:
        height = abs(prev_entry) * 1e-6 or 1e-8
    ax.add_patch(
        Rectangle(
            (prev_x, y_low),
            exit_x - prev_x,
            height,
            facecolor=color,
            edgecolor="none",
            alpha=0.04,
            zorder=1,
        )
    )
    closes = df["Close"].iloc[prev_x : exit_x + 1].astype(float)
    xs = list(range(prev_x, prev_x + len(closes)))
    ax.plot(
        xs,
        closes.to_numpy(),
        color=color,
        linewidth=0.9,
        linestyle=(0, (3, 4)),
        alpha=0.28,
        zorder=5,
        solid_capstyle="round",
    )
    ax.plot(
        [prev_x, exit_x],
        [prev_entry, exit_price],
        color=color,
        linewidth=1.05,
        linestyle=(0, (4, 5)),
        alpha=0.45,
        zorder=6,
    )
    ax.scatter(
        [prev_x],
        [prev_entry],
        s=22,
        c=color,
        marker="o",
        zorder=7,
        edgecolors="#ffffff",
        linewidths=0.7,
    )
    ax.annotate(
        "Prev",
        xy=(prev_x, prev_entry),
        xytext=(6, 8),
        textcoords="offset points",
        va="bottom",
        ha="left",
        fontsize=7,
        color=color,
        zorder=8,
        bbox={
            "boxstyle": "round,pad=0.2",
            "facecolor": COLORS["label_bg"],
            "edgecolor": color,
            "linewidth": 0.6,
            "alpha": 0.88,
        },
    )


def render_setup_chart(
    df: pd.DataFrame,
    *,
    ticker: str,
    action: ActionSide,
    entry: str,
    stop_loss: str,
    tp1: str,
    timeframe: str,
    tp2: str | None = None,
    tp3: str | None = None,
    current_price: str | None = None,
    signal_time: int | None = None,
    right_pad: int | None = None,
    prev_entry: str | None = None,
    prev_side: str | None = None,
    prev_signal_time: int | None = None,
) -> bytes:
    entry_p = _parse_price(entry)
    sl_p = _parse_price(stop_loss)
    tp1_p = _parse_price(tp1)
    tp2_p = _parse_price(tp2) if tp2 else None
    tp3_p = _parse_price(tp3) if tp3 else None
    current_p = _parse_price(current_price) if current_price is not None else None
    prev_entry_p = _parse_price(prev_entry) if prev_entry else None

    is_long = action == "long"
    reward_rgb = COLORS["reward_fill"]

    df = _to_utc_plus_2(df)

    # Position tool starts at seq==1 candle; "now" is the last real candle
    now_x = len(df) - 1
    entry_x = _position_start_x(df, signal_time)
    plot_df = _pad_right(df, RIGHT_PAD_CANDLES if right_pad is None else right_pad)

    level_prices = [entry_p, sl_p, tp1_p]
    if tp2_p is not None:
        level_prices.append(tp2_p)
    if tp3_p is not None:
        level_prices.append(tp3_p)
    if current_p is not None:
        level_prices.append(current_p)
    if prev_entry_p is not None:
        level_prices.append(prev_entry_p)
    y_min = min(float(df["Low"].min()), *level_prices)
    y_max = max(float(df["High"].max()), *level_prices)
    hold_remainder = tp2_p is None and tp3_p is None
    price_pad = (y_max - y_min) * 0.06 or y_max * 0.002

    vol_scaled, vol_colors = _volume_overlay(plot_df, y_low=y_min, y_high=y_max)

    addplots = [
        mpf.make_addplot(
            vol_scaled,
            type="bar",
            panel=0,
            color=vol_colors,
            width=0.8,
            alpha=0.15,  # ~85% transparent
            secondary_y=False,
        ),
    ]

    mc = mpf.make_marketcolors(
        up=COLORS["up"],
        down=COLORS["down"],
        edge="inherit",
        wick="inherit",
        volume="in",
    )
    style = mpf.make_mpf_style(
        base_mpf_style="nightclouds",
        marketcolors=mc,
        facecolor=COLORS["bg"],
        figcolor=COLORS["bg"],
        gridcolor=COLORS["grid"],
        gridstyle="--",
        y_on_right=True,
        rc={
            "axes.labelcolor": COLORS["text"],
            "xtick.color": COLORS["text"],
            "ytick.color": COLORS["text"],
            "axes.edgecolor": COLORS["grid"],
            "figure.facecolor": COLORS["bg"],
            "axes.facecolor": COLORS["panel"],
            "font.size": 9,
        },
    )

    fig, axes = mpf.plot(
        plot_df,
        type="candle",
        style=style,
        volume=False,
        addplot=addplots,
        returnfig=True,
        figsize=(12, 7),
        tight_layout=True,
        datetime_format="%m-%d\n%H:%M",
        warn_too_much_data=10_000,
        xrotation=0,
        ylabel="",
    )
    ax = axes[0]
    ax.set_ylabel("")
    for label in ax.get_xticklabels():
        label.set_horizontalalignment("center")
        label.set_fontsize(8)
        label.set_linespacing(1.35)

    # Room for volume bars under candles
    vol_floor = float(vol_scaled.dropna().min()) if vol_scaled.notna().any() else y_min
    ax.set_ylim(min(y_min - price_pad, vol_floor) - price_pad * 0.3, y_max + price_pad)
    ylim_lo, ylim_hi = ax.get_ylim()
    remainder_end = (ylim_hi if is_long else ylim_lo) if hold_remainder else None

    x_right = ax.get_xlim()[1]
    zone_width = x_right - entry_x

    # Levels + zones start at the entry (last real) candle, not full chart width
    level_specs: list[tuple[float, str, str, float]] = [
        (entry_p, COLORS["entry"], "--", 1.2),
        (sl_p, COLORS["sl"], "-", 1.2),
        (tp1_p, COLORS["tp1"], ":", 1.0),
    ]
    if tp2_p is not None:
        level_specs.append((tp2_p, COLORS["tp2"], ":", 1.0))
    if tp3_p is not None:
        level_specs.append((tp3_p, COLORS["tp3"], ":", 1.0))
    for price, color, ls, lw in level_specs:
        ax.hlines(
            price,
            xmin=entry_x,
            xmax=x_right,
            colors=color,
            linestyles=ls,
            linewidths=lw,
            alpha=0.95,
            zorder=4,
        )

    risk_low = min(entry_p, sl_p)
    risk_high = max(entry_p, sl_p)
    ax.add_patch(
        Rectangle(
            (entry_x, risk_low),
            zone_width,
            risk_high - risk_low,
            facecolor=COLORS["risk_fill"],
            edgecolor="none",
            zorder=0,
        )
    )

    reward_bands: list[tuple[float, float, float]] = [
        (entry_p, tp1_p, REWARD_ALPHAS[0]),
    ]
    if tp2_p is not None:
        reward_bands.append((tp1_p, tp2_p, REWARD_ALPHAS[1]))
    if tp2_p is not None and tp3_p is not None:
        reward_bands.append((tp2_p, tp3_p, REWARD_ALPHAS[2]))
    elif tp3_p is not None:
        reward_bands.append((tp1_p, tp3_p, REWARD_ALPHAS[2]))
    elif remainder_end is not None:
        reward_bands.append((tp1_p, remainder_end, REWARD_ALPHAS[1]))
    for price_a, price_b, alpha in reward_bands:
        band_low = min(price_a, price_b)
        band_high = max(price_a, price_b)
        ax.add_patch(
            Rectangle(
                (entry_x, band_low),
                zone_width,
                band_high - band_low,
                facecolor=(*reward_rgb, alpha),
                edgecolor="none",
                zorder=0,
            )
        )

    ax.scatter(
        [entry_x],
        [entry_p],
        s=22,
        c=COLORS["entry"],
        marker="o",
        zorder=7,
        edgecolors="#ffffff",
        linewidths=0.7,
    )

    if (
        prev_entry_p is not None
        and prev_side is not None
        and prev_signal_time is not None
        and prev_signal_time > 0
    ):
        _draw_previous_trade(
            ax,
            df,
            prev_entry=prev_entry_p,
            prev_side=prev_side,
            prev_signal_time=prev_signal_time,
            exit_price=entry_p,
            exit_x=entry_x,
        )

    if current_p is not None:
        ax.hlines(
            current_p,
            xmin=now_x,
            xmax=x_right,
            colors=COLORS["price"],
            linestyles="-.",
            linewidths=1.0,
            alpha=0.95,
            zorder=5,
        )
        ax.scatter(
            [now_x],
            [current_p],
            s=28,
            c=COLORS["price"],
            marker="D",
            zorder=7,
            edgecolors="#ffffff",
            linewidths=0.7,
        )

    labels = [
        (entry_p, f"Entry {entry}", COLORS["entry"]),
        (sl_p, f"SL {stop_loss}", COLORS["sl"]),
        (tp1_p, f"TP1 {tp1}", COLORS["tp1"]),
    ]
    if tp2_p is not None and tp2 is not None:
        labels.append((tp2_p, f"TP2 {tp2}", COLORS["tp2"]))
    if tp3_p is not None and tp3 is not None:
        labels.append((tp3_p, f"TP3 {tp3}", COLORS["tp3"]))
    if current_p is not None and current_price is not None:
        labels.append((current_p, f"Price {current_price}", COLORS["price"]))
    for price, text, color in labels:
        ax.annotate(
            text,
            xy=(x_right, price),
            xytext=(6, 0),
            textcoords="offset points",
            va="center",
            ha="left",
            fontsize=8,
            color=color,
            clip_on=False,
            zorder=8,
            bbox={
                "boxstyle": "round,pad=0.28",
                "facecolor": COLORS["label_bg"],
                "edgecolor": color,
                "linewidth": 0.8,
                "alpha": 0.92,
            },
        )

    side = "LONG" if is_long else "SHORT"
    ax.set_title(
        f"{ticker}  ·  {timeframe}  ·  {side}",
        color=COLORS["text"],
        fontsize=12,
        pad=12,
    )

    fig.subplots_adjust(right=0.82)

    buf = io.BytesIO()
    fig.savefig(buf, format="png", dpi=140, facecolor=COLORS["bg"], bbox_inches="tight")
    plt.close(fig)
    buf.seek(0)
    return buf.read()