# export_image.py

import argparse
import os
import time
from playwright.sync_api import sync_playwright
from pathlib import Path

# Define the output directory as a constant
OUTPUT_DIR = "images"


def export_image(url_path, filename_prefix):
    """
    Launches a headless browser, navigates to a URL path, and screenshots
    the main content element, saving it to the OUTPUT_DIR.
    """
    print(f"-> Preparing to export image for: {filename_prefix}")

    os.makedirs(OUTPUT_DIR, exist_ok=True)

    base_url = "http://localhost"
    # Ensure the URL path starts correctly
    url_path = url_path.lstrip("/")
    url = f"{base_url}/{url_path}"

    output_file = os.path.join(OUTPUT_DIR, f"{filename_prefix}_{int(time.time())}.png")

    with sync_playwright() as p:
        try:
            browser = p.chromium.launch()
            page = browser.new_page()

            page.set_viewport_size({"width": 1920, "height": 1080})

            print(f"  Navigating to {url}...")
            # Use 'domcontentloaded' for faster navigation when possible
            page.goto(url, wait_until="domcontentloaded")

            # Give fonts and styles a moment to render after the DOM is ready
            page.wait_for_timeout(500)

            # --- THIS IS THE CORRECTED LINE ---
            # Target the new, correct class for our main content card
            element = page.locator(".max-w-3xl")

            print(f"  Saving screenshot to {output_file}...")
            element.screenshot(path=output_file)

            browser.close()
            print(f"-> Export complete! Image saved to {output_file}")

        except Exception as e:
            print(f"\nAn error occurred during export: {e}")
            print(
                "Please ensure the 'rstat-dashboard' server is running in another terminal."
            )


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Export subreddit sentiment images.")
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument("-s", "--subreddit", help="The name of the subreddit to export.")
    group.add_argument(
        "-o", "--overall", action="store_true", help="Export the overall summary image."
    )

    parser.add_argument(
        "-w",
        "--weekly",
        action="store_true",
        help="Export the weekly view instead of the daily view.",
    )
    args = parser.parse_args()

    view_type = "weekly" if args.weekly else "daily"

    if args.subreddit:
        url_path_to_render = f"subreddit/{args.subreddit}?view={view_type}&image=true"
        filename_prefix_to_save = f"{args.subreddit}_{view_type}"
        export_image(url_path_to_render, filename_prefix_to_save)

    elif args.overall:
        url_path_to_render = f"/?view={view_type}&image=true"
        filename_prefix_to_save = f"overall_summary_{view_type}"
        export_image(url_path_to_render, filename_prefix_to_save)