"""Test script for candlestick chart generation."""
import logging
from polygon_api import PolygonAPI
from chart_generator import ChartGenerator
import os
from dotenv import load_dotenv

load_dotenv()

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


def test_candlestick_charts():
    """Test candlestick chart generation with Polygon data."""

    api_key = os.getenv('POLYGON_API_KEY')
    if not api_key:
        logger.error("POLYGON_API_KEY not set in .env file")
        return False

    logger.info("Initializing Polygon API")
    api = PolygonAPI(api_key)

    # Test tickers
    tickers = ["AAPL", "TSLA"]

    for ticker in tickers:
        logger.info(f"\nTesting candlestick chart for {ticker}...")

        # Fetch candle data
        logger.info(f"Fetching 30-day candle data for {ticker}")
        candle_data = api.get_candles(ticker, days=30, resolution='D')

        if not candle_data:
            logger.error(f"Failed to fetch candle data for {ticker}")
            continue

        logger.info(f"Successfully fetched {len(candle_data['timestamps'])} candles")
        logger.info(f"Price range: ${min(candle_data['low']):.2f} - ${max(candle_data['high']):.2f}")

        # Generate candlestick chart
        logger.info(f"Generating candlestick chart for {ticker}")
        chart = ChartGenerator.create_candlestick_chart(ticker, candle_data, f"{ticker}")

        if not chart:
            logger.error(f"Failed to generate chart for {ticker}")
            continue

        # Save chart
        filename = f"/tmp/{ticker}_candlestick_30d.png"
        with open(filename, "wb") as f:
            f.write(chart.getvalue())
        logger.info(f"Chart saved: {filename} ({len(chart.getvalue())} bytes)")

    logger.info("\nCandlestick chart test completed successfully!")
    return True


if __name__ == "__main__":
    try:
        success = test_candlestick_charts()
        exit(0 if success else 1)
    except Exception as e:
        logger.error(f"Test failed: {e}", exc_info=True)
        exit(1)