mod url_allowlist;

use anyhow::{anyhow, Result};
use aws_sdk_s3::primitives::ByteStream;
use chromiumoxide::browser::{Browser, BrowserConfig};
use chromiumoxide::cdp::browser_protocol::page::CaptureScreenshotFormat;
use chromiumoxide::page::ScreenshotParams;
use futures::StreamExt;
use std::sync::Arc;
use tokio::sync::Semaphore;
use uuid::Uuid;

pub use url_allowlist::assert_url_allowed;

/// Env-backed settings for [`ScreenshotService`] (bucket, prefix, concurrency).
#[derive(Debug, Clone)]
pub struct ScreenshotConfig {
    pub bucket: String,
    pub prefix: String,
    /// Maximum concurrent screenshot captures (Chromium launches). `0` means unlimited.
    pub max_concurrent: usize,
}

impl ScreenshotConfig {
    pub fn from_env() -> Option<Self> {
        let bucket: String = std::env::var("S3_SCREENSHOT_BUCKET").ok()?;
        let prefix: String = std::env::var("S3_SCREENSHOT_PREFIX").unwrap_or_else(|_| "screenshots/".to_string());
        let prefix: String = if prefix.ends_with('/') {
            prefix
        } else {
            format!("{}/", prefix)
        };
        let max_concurrent: usize = std::env::var("SCREENSHOT_MAX_CONCURRENT")
            .ok()
            .and_then(|s| s.parse().ok())
            .unwrap_or(2);
        Some(Self {
            bucket,
            prefix,
            max_concurrent,
        })
    }
}

/// S3 upload + Chromium capture. Build with [`ScreenshotService::from_env`].
pub struct ScreenshotService {
    s3: aws_sdk_s3::Client,
    config: ScreenshotConfig,
    concurrent: Option<Arc<Semaphore>>,
}

impl ScreenshotService {
    /// Requires `S3_SCREENSHOT_BUCKET`. Uses [`aws_config::load_from_env`] for credentials/region.
    pub async fn from_env() -> Option<Self> {
        let config: ScreenshotConfig = ScreenshotConfig::from_env()?;
        let cfg: aws_config::SdkConfig = aws_config::load_from_env().await;
        let s3 = aws_sdk_s3::Client::new(&cfg);
        let concurrent = if config.max_concurrent > 0 {
            Some(Arc::new(Semaphore::new(config.max_concurrent)))
        } else {
            None
        };
        Some(Self {
            s3,
            config,
            concurrent,
        })
    }

    /// Validate URL, render PNG with headless Chromium, upload; returns S3 object key.
    pub async fn capture_and_upload_png(&self, target_url: &str) -> Result<String> {
        assert_url_allowed(target_url)?;
        let _permit = match &self.concurrent {
            Some(sem) => Some(
                sem.clone()
                    .acquire_owned()
                    .await
                    .map_err(|e| anyhow!("screenshot concurrency: {}", e))?,
            ),
            None => None,
        };
        let png = capture_viewport_png(target_url).await?;
        let key = format!("{}{}.png", self.config.prefix, Uuid::new_v4());
        self.s3
            .put_object()
            .bucket(&self.config.bucket)
            .key(&key)
            .content_type("image/png")
            .body(ByteStream::from(png))
            .send()
            .await
            .map_err(|e| anyhow!("S3 put_object: {}", e))?;
        Ok(key)
    }
}

async fn capture_viewport_png(url: &str) -> Result<Vec<u8>> {
    let mut builder = BrowserConfig::builder();
    if let Ok(exe) = std::env::var("CHROME_EXECUTABLE") {
        builder = builder.chrome_executable(exe);
    }
    let cfg = builder.build().map_err(|e| anyhow!("BrowserConfig: {}", e))?;
    let (mut browser, mut handler) = Browser::launch(cfg)
        .await
        .map_err(|e| anyhow!("Chromium launch: {}", e))?;

    let h = tokio::spawn(async move {
        while handler.next().await.is_some() {}
    });

    let page = browser
        .new_page(url)
        .await
        .map_err(|e| anyhow!("new_page: {}", e))?;

    let png = page
        .screenshot(
            ScreenshotParams::builder()
                .format(CaptureScreenshotFormat::Png)
                .build(),
        )
        .await
        .map_err(|e| anyhow!("screenshot: {}", e))?;

    h.abort();
    let _ = browser.close().await;

    Ok(png)
}
