services/chart_service.py

"""
Stock price + volume chart image generator (matplotlib, Agg backend).

Generates a compact dark-themed PNG chart suitable for LINE messages:
- Close price line with MA5 / MA20
- Volume bars (green = up day, red = down day)
- Stats overlay: period high/low, % change, avg volume
- Output: 800x500 px PNG bytes
"""

import io
import logging
from datetime import datetime, date
from typing import Optional

import matplotlib
matplotlib.use("Agg")  # headless — no display required

import matplotlib.pyplot as plt
import matplotlib.dates as mdates
import matplotlib.ticker as mticker
import matplotlib.font_manager as fm
import numpy as np
import pandas as pd

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# CJK font setup — try system fonts, fall back to sans-serif
# ---------------------------------------------------------------------------

def _setup_cjk_font() -> None:
    """Configure matplotlib to use a CJK-capable font if available."""
    # Priority list of CJK fonts (Windows → Linux → Docker)
    candidates = [
        "Microsoft JhengHei",   # Windows 繁中
        "Microsoft YaHei",      # Windows 简中
        "Noto Sans CJK TC",     # Linux / Docker (apt: fonts-noto-cjk)
        "Noto Sans CJK SC",
        "WenQuanYi Micro Hei",  # Linux fallback
        "SimHei",               # older Windows
        "Arial Unicode MS",     # macOS
    ]
    available = {f.name for f in fm.fontManager.ttflist}
    for font_name in candidates:
        if font_name in available:
            plt.rcParams["font.sans-serif"] = [font_name, "DejaVu Sans"]
            plt.rcParams["axes.unicode_minus"] = False
            logger.info("Chart CJK font: %s", font_name)
            return
    # No CJK font found — Chinese characters will show as boxes
    logger.warning("No CJK font found; chart titles with Chinese may not render correctly")

_setup_cjk_font()

# ---------------------------------------------------------------------------
# Dark theme colours (matches the web UI palette)
# ---------------------------------------------------------------------------

_BG        = "#1a1a2e"
_PANEL_BG  = "#16213e"
_TEXT       = "#e0e0e0"
_TEXT_DIM   = "#888888"
_GRID      = "#2a2a4a"
_PRICE     = "#00d2ff"
_MA5       = "#ffd700"
_MA20      = "#ff6b9d"
_VOL_UP    = "#27ACB2"
_VOL_DOWN  = "#FF6B35"


def generate_chart_image(
    df: pd.DataFrame,
    stock_no: str,
    name: str = "",
    months: int = 2,
) -> bytes:
    """
    Generate a price + volume chart image as PNG bytes.

    Parameters
    ----------
    df : DataFrame with columns [date, open, high, low, close, volume].
         Should contain at least ``months`` months of data.
    stock_no : stock code (displayed in title).
    name : stock Chinese/English name.
    months : how many trailing months to show (default 2).

    Returns
    -------
    PNG image bytes.
    """
    # Trim to requested months
    df = df.copy()
    df["date"] = pd.to_datetime(df["date"])
    cutoff = df["date"].max() - pd.DateOffset(months=months)
    df = df[df["date"] >= cutoff].reset_index(drop=True)

    if df.empty or len(df) < 3:
        raise ValueError(f"Not enough data to render chart for {stock_no}")

    # Compute moving averages
    df["ma5"]  = df["close"].rolling(5, min_periods=1).mean()
    df["ma20"] = df["close"].rolling(20, min_periods=1).mean()

    dates  = df["date"]
    close  = df["close"]
    volume = df["volume"]
    ma5    = df["ma5"]
    ma20   = df["ma20"]

    # Up/down for volume bar colouring
    prev_close = close.shift(1).fillna(close.iloc[0])
    vol_colors = [_VOL_UP if c >= p else _VOL_DOWN for c, p in zip(close, prev_close)]

    # Stats
    period_high  = float(df["high"].max()) if "high" in df.columns else float(close.max())
    period_low   = float(df["low"].min())  if "low"  in df.columns else float(close.min())
    total_change  = (float(close.iloc[-1]) - float(close.iloc[0])) / float(close.iloc[0]) * 100
    avg_vol      = float(volume.mean())

    # -----------------------------------------------------------------------
    # Create figure
    # -----------------------------------------------------------------------
    fig, (ax_price, ax_vol) = plt.subplots(
        2, 1,
        figsize=(8, 5),
        dpi=100,
        gridspec_kw={"height_ratios": [3, 1], "hspace": 0.08},
        facecolor=_BG,
    )

    for ax in (ax_price, ax_vol):
        ax.set_facecolor(_PANEL_BG)
        ax.tick_params(colors=_TEXT_DIM, labelsize=8)
        ax.grid(True, color=_GRID, linewidth=0.4, alpha=0.6)
        for spine in ax.spines.values():
            spine.set_color(_GRID)

    # -- Price subplot --
    ax_price.plot(dates, close, color=_PRICE, linewidth=1.5, label="Close")
    ax_price.plot(dates, ma5,   color=_MA5,   linewidth=1.0, label="MA5", alpha=0.8)
    ax_price.plot(dates, ma20,  color=_MA20,  linewidth=1.0, label="MA20", alpha=0.8)

    # Fill between close and MA20 for visual emphasis
    ax_price.fill_between(dates, close, ma20, alpha=0.08, color=_PRICE)

    # Legend
    ax_price.legend(
        loc="upper left", fontsize=7,
        facecolor=_PANEL_BG, edgecolor=_GRID,
        labelcolor=_TEXT_DIM,
    )

    # Title
    title = f"{name}  {stock_no}" if name else stock_no
    ax_price.set_title(title, color=_TEXT, fontsize=12, fontweight="bold", pad=10)

    # Stats text box
    change_color = _VOL_UP if total_change >= 0 else _VOL_DOWN
    stats_text = (
        f"H: {period_high:,.2f}  L: {period_low:,.2f}  "
        f"Chg: {total_change:+.2f}%  "
        f"AvgVol: {_format_volume(avg_vol)}"
    )
    ax_price.text(
        0.99, 0.97, stats_text,
        transform=ax_price.transAxes,
        fontsize=7, color=_TEXT_DIM,
        ha="right", va="top",
        bbox=dict(boxstyle="round,pad=0.3", facecolor=_BG, edgecolor=_GRID, alpha=0.8),
    )

    # Remove x-tick labels from price subplot (shared with volume)
    ax_price.set_xticklabels([])

    # Y-axis formatting
    ax_price.yaxis.set_major_formatter(mticker.FormatStrFormatter("%.1f"))

    # -- Volume subplot --
    ax_vol.bar(dates, volume, color=vol_colors, width=0.8, alpha=0.85)
    ax_vol.set_ylabel("Vol", color=_TEXT_DIM, fontsize=8)

    # Volume y-axis: abbreviated labels
    ax_vol.yaxis.set_major_formatter(
        mticker.FuncFormatter(lambda x, _: _format_volume(x))
    )

    # Date formatting on volume x-axis
    ax_vol.xaxis.set_major_formatter(mdates.DateFormatter("%m/%d"))
    ax_vol.xaxis.set_major_locator(mdates.WeekdayLocator(byweekday=mdates.MO, interval=2))
    fig.autofmt_xdate(rotation=30, ha="right")

    # Watermark
    fig.text(
        0.99, 0.01, "stock-predictor",
        fontsize=6, color=_TEXT_DIM, alpha=0.4,
        ha="right", va="bottom",
    )

    fig.subplots_adjust(left=0.10, right=0.95, top=0.92, bottom=0.12, hspace=0.08)

    # -----------------------------------------------------------------------
    # Export to PNG bytes
    # -----------------------------------------------------------------------
    buf = io.BytesIO()
    fig.savefig(buf, format="png", facecolor=fig.get_facecolor(), bbox_inches="tight")
    plt.close(fig)
    buf.seek(0)
    return buf.read()


def _format_volume(v: float) -> str:
    """Abbreviate volume numbers: 1.2M, 345K, etc."""
    if v >= 1_000_000_000:
        return f"{v / 1_000_000_000:.1f}B"
    if v >= 1_000_000:
        return f"{v / 1_000_000:.1f}M"
    if v >= 1_000:
        return f"{v / 1_000:.0f}K"
    return f"{v:.0f}"