app/chart_generator.py

"""
Generates and visualizes market data charts correlated with news events.
"""
import io
import base64
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.dates as mdates

def generate_price_chart(price_data: list, event_timestamp: pd.Timestamp, entity: str) -> str:
    """
    Generates a base64-encoded price chart image with an event annotation.

    Args:
        price_data: A list of [timestamp, price] pairs from CoinGecko.
        event_timestamp: The timestamp of the news event.
        entity: The cryptocurrency entity (e.g., 'Bitcoin').

    Returns:
        A base64 encoded string of the PNG chart image.
    """
    if not price_data:
        return ""

    # Use a dark theme for the chart
    plt.style.use('dark_background')
    
    # Create a pandas DataFrame
    df = pd.DataFrame(price_data, columns=['timestamp', 'price'])
    df['timestamp'] = pd.to_datetime(df['timestamp'], unit='ms')
    df = df.set_index('timestamp')

    fig, ax = plt.subplots(figsize=(10, 4))
    
    # Plot the price data
    ax.plot(df.index, df['price'], color='cyan', linewidth=2)

    # Annotate the event
    try:
        event_price = df.asof(event_timestamp)['price']
        ax.axvline(event_timestamp, color='red', linestyle='--', linewidth=1.5, label=f'Event: {entity}')
        ax.plot(event_timestamp, event_price, 'ro', markersize=8) # Red dot on the event
        ax.annotate(f'Event',
                    xy=(event_timestamp, event_price),
                    xytext=(event_timestamp, event_price * 1.01),
                    ha='center',
                    arrowprops=dict(facecolor='white', shrink=0.05, width=1, headwidth=4),
                    bbox=dict(boxstyle='round,pad=0.3', fc='yellow', ec='k', lw=1, alpha=0.8),
                    color='black'
                   )
    except KeyError:
        # Event timestamp might be out of range
        ax.axvline(event_timestamp, color='red', linestyle='--', linewidth=1.5)

    # Formatting the chart
    ax.set_title(f'{entity.upper()} Price Action around Event', fontsize=14)
    ax.set_ylabel('Price (USD)')
    ax.grid(True, linestyle='--', alpha=0.3)
    fig.autofmt_xdate()
    ax.xaxis.set_major_formatter(mdates.DateFormatter('%H:%M'))
    ax.tick_params(axis='x', rotation=45)
    plt.tight_layout()

    # Save to an in-memory buffer
    buf = io.BytesIO()
    fig.savefig(buf, format='png', transparent=True)
    buf.seek(0)
    img_base64 = base64.b64encode(buf.read()).decode('utf-8')
    plt.close(fig)

    return f"data:image/png;base64,{img_base64}"