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",
    "long_fill": (0.15, 0.65, 0.45),
    "short_fill": (0.85, 0.25, 0.25),
    "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 render_setup_chart(
    df: pd.DataFrame,
    *,
    ticker: str,
    action: ActionSide,
    entry: str,
    stop_loss: str,
    tp1: str,
    tp2: str,
    tp3: str,
    timeframe: str,
    current_price: str | None = None,
    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)
    tp3_p = _parse_price(tp3)
    current_p = _parse_price(current_price) if current_price is not None else None

    is_long = action == "long"
    reward_rgb = COLORS["long_fill"] if is_long else COLORS["short_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)

    level_prices = [entry_p, sl_p, tp1_p, tp2_p, tp3_p]
    if current_p is not None:
        level_prices.append(current_p)
    y_min = min(float(df["Low"].min()), *level_prices)
    y_max = max(float(df["High"].max()), *level_prices)
    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)

    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 = [
        (entry_p, COLORS["entry"], "-", 1.4),
        (sl_p, COLORS["sl"], "--", 1.2),
        (tp1_p, COLORS["tp1"], ":", 1.0),
        (tp2_p, COLORS["tp2"], ":", 1.0),
        (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,
        )
    )

    # Three reward bands with decreasing opacity toward farther TPs
    reward_bands = (
        (entry_p, tp1_p, REWARD_ALPHAS[0]),
        (tp1_p, tp2_p, REWARD_ALPHAS[1]),
        (tp2_p, tp3_p, REWARD_ALPHAS[2]),
    )
    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 current_p is not None:
        ax.hlines(
            current_p,
            xmin=entry_x,
            xmax=x_right,
            colors=COLORS["price"],
            linestyles="-.",
            linewidths=1.3,
            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"]),
        (tp2_p, f"TP2 {tp2}", COLORS["tp2"]),
        (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()